"""ComputerUseSession - 虚拟桌面会话管理

管理虚拟桌面会话，维护操作上下文。
提供 InMemoryComputerUseSession 用于测试，LocalComputerUseSession 用于本地桌面，
DockerComputerUseSession 用于 Docker 沙箱（stub）。
"""

from __future__ import annotations

import asyncio
import base64
import io
import logging
import time
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import TypeAlias

logger = logging.getLogger(__name__)

# Scalar session state values (text/number/flag/none). Screen-state dicts use
# dict[str, object] because they also hold tuple cursor positions.
SessionState: TypeAlias = dict[str, str | int | bool | None]


@dataclass
class ScreenInfo:
    """屏幕信息"""

    width: int = 1280
    height: int = 720


@dataclass
class ActionResult:
    """操作执行结果"""

    success: bool
    action: str
    output: str = ""
    screenshot_base64: str = ""
    error: str = ""
    metadata: dict[str, object] = field(default_factory=dict)


class ComputerUseSession(ABC):
    """虚拟桌面会话抽象基类

    管理虚拟桌面会话，维护操作上下文。
    子类实现具体的会话管理（Docker 沙箱、内存模拟等）。
    """

    def __init__(
        self,
        session_id: str | None = None,
        screen_width: int = 1280,
        screen_height: int = 720,
    ):
        self.session_id = session_id or str(uuid.uuid4())
        self.screen = ScreenInfo(width=screen_width, height=screen_height)
        self._started = False
        self._action_history: list[dict[str, object]] = []

    @property
    def is_started(self) -> bool:
        return self._started

    @abstractmethod
    async def start(self) -> None:
        """启动虚拟桌面会话"""
        ...

    @abstractmethod
    async def stop(self) -> None:
        """停止虚拟桌面会话"""
        ...

    @abstractmethod
    async def screenshot(self) -> ActionResult:
        """截取当前屏幕

        Returns:
            ActionResult 包含 screenshot_base64 和屏幕内容描述
        """
        ...

    @abstractmethod
    async def execute_action(self, action: str, **params: object) -> ActionResult:
        """执行 UI 操作

        Args:
            action: 操作类型 (click/type/scroll/drag/key/wait)
            **params: 操作参数

        Returns:
            ActionResult 包含操作结果
        """
        ...

    def record_action(self, action: str, params: dict[str, object], result: ActionResult) -> None:
        """记录操作历史"""
        self._action_history.append(
            {
                "timestamp": time.time(),
                "action": action,
                "params": params,
                "success": result.success,
                "output": result.output[:200] if result.output else "",
            }
        )

    @property
    def action_history(self) -> list[dict[str, object]]:
        """获取操作历史（副本）"""
        return list(self._action_history)

    def __repr__(self) -> str:
        status = "started" if self._started else "stopped"
        return (
            f"<{type(self).__name__} id={self.session_id[:8]} "
            f"screen={self.screen.width}x{self.screen.height} {status}>"
        )


class InMemoryComputerUseSession(ComputerUseSession):
    """内存模拟会话，用于测试

    不连接真实虚拟桌面，模拟截屏和操作结果。
    """

    def __init__(
        self,
        session_id: str | None = None,
        screen_width: int = 1280,
        screen_height: int = 720,
    ):
        super().__init__(
            session_id=session_id,
            screen_width=screen_width,
            screen_height=screen_height,
        )
        self._screen_state: dict[str, object] = {
            "focused_element": None,
            "cursor_position": (0, 0),
            "typed_text": "",
        }

    async def start(self) -> None:
        """启动内存会话"""
        self._started = True
        logger.info("InMemory session started: %s", self.session_id[:8])

    async def stop(self) -> None:
        """停止内存会话"""
        self._started = False
        logger.info("InMemory session stopped: %s", self.session_id[:8])

    async def screenshot(self) -> ActionResult:
        """模拟截屏，返回空截图和当前状态描述"""
        if not self._started:
            return ActionResult(
                success=False,
                action="screenshot",
                error="Session not started",
            )

        state_desc = (
            f"Screen {self.screen.width}x{self.screen.height}, "
            f"cursor at {self._screen_state['cursor_position']}, "
            f"focused: {self._screen_state['focused_element']}, "
            f"typed: '{self._screen_state['typed_text']}'"
        )
        return ActionResult(
            success=True,
            action="screenshot",
            output=state_desc,
            screenshot_base64="",
            metadata={"screen_state": dict(self._screen_state)},
        )

    async def execute_action(self, action: str, **params: object) -> ActionResult:
        """模拟执行 UI 操作"""
        if not self._started:
            return ActionResult(
                success=False,
                action=action,
                error="Session not started",
            )

        result = self._simulate_action(action, **params)
        self.record_action(action, params, result)
        return result

    def _simulate_action(self, action: str, **params: object) -> ActionResult:
        """模拟具体操作"""
        if action == "click":
            x = params.get("x", 0)
            y = params.get("y", 0)
            self._screen_state["cursor_position"] = (x, y)
            self._screen_state["focused_element"] = f"element@({x},{y})"
            return ActionResult(
                success=True,
                action="click",
                output=f"Clicked at ({x}, {y})",
            )

        if action == "type":
            text = params.get("text", "")
            self._screen_state["typed_text"] += text
            return ActionResult(
                success=True,
                action="type",
                output=f"Typed: {text}",
            )

        if action == "scroll":
            direction = params.get("direction", "down")
            amount = params.get("amount", 3)
            return ActionResult(
                success=True,
                action="scroll",
                output=f"Scrolled {direction} by {amount}",
            )

        if action == "drag":
            start_x = params.get("start_x", 0)
            start_y = params.get("start_y", 0)
            end_x = params.get("end_x", 0)
            end_y = params.get("end_y", 0)
            self._screen_state["cursor_position"] = (end_x, end_y)
            return ActionResult(
                success=True,
                action="drag",
                output=f"Dragged from ({start_x},{start_y}) to ({end_x},{end_y})",
            )

        if action == "key":
            key_name = params.get("key_name", "")
            return ActionResult(
                success=True,
                action="key",
                output=f"Pressed key: {key_name}",
            )

        if action == "wait":
            duration = params.get("duration", 1.0)
            return ActionResult(
                success=True,
                action="wait",
                output=f"Waited {duration}s",
            )

        return ActionResult(
            success=False,
            action=action,
            error=f"Unknown action: {action}",
        )


class LocalComputerUseSession(ComputerUseSession):
    """本地桌面会话，使用 pyautogui + screencapture。

    在本地桌面执行真实的 UI 操作。需要辅助功能权限（macOS）。
    适用于开发/演示环境，生产环境应使用 DockerComputerUseSession。
    """

    def __init__(
        self,
        session_id: str | None = None,
        screen_width: int = 1280,
        screen_height: int = 720,
    ):
        super().__init__(
            session_id=session_id,
            screen_width=screen_width,
            screen_height=screen_height,
        )
        self._pyautogui: object | None = None

    # Allowed keys for the `key` action — prevents OS-level shortcut abuse
    _ALLOWED_KEYS: set[str] = {
        "enter",
        "return",
        "tab",
        "backspace",
        "delete",
        "home",
        "end",
        "up",
        "down",
        "left",
        "right",
        "pageup",
        "pagedown",
        "space",
        "escape",
        "insert",
        "f1",
        "f2",
        "f3",
        "f4",
        "f5",
        "f6",
        "f7",
        "f8",
        "f9",
        "f10",
        "f11",
        "f12",
        "shift",
        "ctrl",
        "alt",
        "command",
        "a",
        "b",
        "c",
        "d",
        "e",
        "f",
        "g",
        "h",
        "i",
        "j",
        "k",
        "l",
        "m",
        "n",
        "o",
        "p",
        "q",
        "r",
        "s",
        "t",
        "u",
        "v",
        "w",
        "x",
        "y",
        "z",
        "0",
        "1",
        "2",
        "3",
        "4",
        "5",
        "6",
        "7",
        "8",
        "9",
    }
    _ALLOWED_BUTTONS: set[str] = {"left", "right", "middle"}
    _MAX_TEXT_LENGTH: int = 1000
    _MAX_WAIT_SECONDS: float = 10.0

    async def start(self) -> None:
        """启动本地桌面会话"""
        try:
            import pyautogui

            self._pyautogui = pyautogui
            pyautogui.FAILSAFE = True
            pyautogui.PAUSE = 0.1
        except ImportError:
            raise RuntimeError(
                "pyautogui is required for LocalComputerUseSession. "
                "Install with: pip install pyautogui"
            )
        self._started = True
        logger.info("Local desktop session started: %s", self.session_id[:8])

    async def stop(self) -> None:
        """停止本地桌面会话"""
        self._started = False
        self._pyautogui = None
        logger.info("Local desktop session stopped: %s", self.session_id[:8])

    async def screenshot(self) -> ActionResult:
        """截取本地桌面屏幕"""
        if not self._started:
            return ActionResult(success=False, action="screenshot", error="Session not started")

        try:
            img = self._pyautogui.screenshot()
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            screenshot_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
            return ActionResult(
                success=True,
                action="screenshot",
                output=f"Screenshot {img.size[0]}x{img.size[1]}",
                screenshot_base64=screenshot_b64,
            )
        except Exception as e:
            return ActionResult(success=False, action="screenshot", error=str(e))

    async def execute_action(self, action: str, **params: object) -> ActionResult:
        """在本地桌面执行 UI 操作"""
        if not self._started:
            return ActionResult(success=False, action=action, error="Session not started")

        try:
            result = self._execute_local_action(action, **params)
            self.record_action(action, params, result)
            return result
        except Exception as e:
            return ActionResult(success=False, action=action, error=str(e))

    def _validate_coordinates(self, x: int, y: int) -> bool:
        """Check if coordinates are within screen bounds."""
        return 0 <= x <= self.screen_width and 0 <= y <= self.screen_height

    async def _execute_local_action(self, action: str, **params: object) -> ActionResult:
        """Execute a local UI action with input validation."""
        pg = self._pyautogui

        if action == "click":
            x, y = params.get("x", 0), params.get("y", 0)
            button = params.get("button", "left")
            if button not in self._ALLOWED_BUTTONS:
                return ActionResult(
                    success=False, action="click", error=f"Invalid button: {button}"
                )
            if not self._validate_coordinates(x, y):
                return ActionResult(
                    success=False, action="click", error=f"Coordinates out of bounds: ({x}, {y})"
                )
            pg.click(x, y, button=button)
            return ActionResult(success=True, action="click", output=f"Clicked at ({x}, {y})")

        if action == "type":
            text = params.get("text", "")
            if len(text) > self._MAX_TEXT_LENGTH:
                return ActionResult(
                    success=False,
                    action="type",
                    error=f"Text too long: {len(text)} > {self._MAX_TEXT_LENGTH}",
                )
            pg.write(text)
            return ActionResult(success=True, action="type", output=f"Typed: {text[:50]}")

        if action == "scroll":
            direction = params.get("direction", "down")
            amount = params.get("amount", 3)
            clicks = amount if direction == "down" else -amount
            pg.scroll(clicks)
            return ActionResult(
                success=True, action="scroll", output=f"Scrolled {direction} by {amount}"
            )

        if action == "drag":
            sx, sy = params.get("start_x", 0), params.get("start_y", 0)
            ex, ey = params.get("end_x", 0), params.get("end_y", 0)
            if not (self._validate_coordinates(sx, sy) and self._validate_coordinates(ex, ey)):
                return ActionResult(
                    success=False, action="drag", error="Drag coordinates out of bounds"
                )
            pg.moveTo(sx, sy)
            pg.dragTo(ex, ey, duration=0.5)
            return ActionResult(
                success=True, action="drag", output=f"Dragged from ({sx},{sy}) to ({ex},{ey})"
            )

        if action == "key":
            key_name = params.get("key_name", "")
            keys = key_name.split("+")
            for k in keys:
                if k.lower() not in self._ALLOWED_KEYS:
                    return ActionResult(success=False, action="key", error=f"Key not allowed: {k}")
            if len(keys) > 1:
                pg.hotkey(*keys)
            else:
                pg.press(key_name)
            return ActionResult(success=True, action="key", output=f"Pressed key: {key_name}")

        if action == "wait":
            duration = min(params.get("duration", 1.0), self._MAX_WAIT_SECONDS)
            await asyncio.sleep(duration)
            return ActionResult(success=True, action="wait", output=f"Waited {duration}s")

        return ActionResult(success=False, action=action, error=f"Unknown action: {action}")


class DockerComputerUseSession(ComputerUseSession):
    """Docker 沙箱虚拟桌面会话

    通过 Docker 容器运行虚拟桌面，连接 Anthropic Computer Use API。
    当前为 stub 实现，可后续接入真实 Docker 环境。
    """

    def __init__(
        self,
        session_id: str | None = None,
        screen_width: int = 1280,
        screen_height: int = 720,
        container_image: str = "anthropic/computer-use-demo:latest",
        docker_url: str = "unix:///var/run/docker.sock",
    ):
        super().__init__(
            session_id=session_id,
            screen_width=screen_width,
            screen_height=screen_height,
        )
        self._container_image = container_image
        self._docker_url = docker_url
        self._container_id: str | None = None

    @property
    def container_id(self) -> str | None:
        return self._container_id

    async def start(self) -> None:
        """启动 Docker 容器

        Stub: 实际实现需要通过 Docker API 创建容器。
        """
        logger.info(
            "Docker session start requested: %s (image=%s)",
            self.session_id[:8],
            self._container_image,
        )
        # TODO: 实际 Docker 容器创建逻辑
        # async with docker.DockerClient(base_url=self._docker_url) as client:
        #     container = await client.containers.run(
        #         self._container_image,
        #         detach=True,
        #         environment={"SCREEN_WIDTH": self.screen.width, ...},
        #     )
        #     self._container_id = container.id
        self._started = True
        self._container_id = f"stub-{self.session_id[:8]}"
        logger.info("Docker session started (stub): %s", self.session_id[:8])

    async def stop(self) -> None:
        """停止 Docker 容器"""
        if self._container_id:
            logger.info(
                "Docker session stop requested: %s (container=%s)",
                self.session_id[:8],
                self._container_id[:12],
            )
            # TODO: 实际 Docker 容器停止逻辑
            # async with docker.DockerClient(base_url=self._docker_url) as client:
            #     container = await client.containers.get(self._container_id)
            #     await container.stop()
            #     await container.remove()
        self._started = False
        self._container_id = None
        logger.info("Docker session stopped: %s", self.session_id[:8])

    async def screenshot(self) -> ActionResult:
        """截取 Docker 虚拟桌面屏幕

        Stub: 实际实现需要通过 VNC/HTTP 获取截图。
        """
        if not self._started:
            return ActionResult(
                success=False,
                action="screenshot",
                error="Session not started",
            )

        # TODO: 实际截屏逻辑
        return ActionResult(
            success=True,
            action="screenshot",
            output=f"Screenshot from container {self._container_id}",
            screenshot_base64="",
        )

    async def execute_action(self, action: str, **params: object) -> ActionResult:
        """在 Docker 虚拟桌面执行操作

        Stub: 实际实现需要通过 Anthropic Computer Use API。
        """
        if not self._started:
            return ActionResult(
                success=False,
                action=action,
                error="Session not started",
            )

        # TODO: 实际操作执行逻辑
        result = ActionResult(
            success=True,
            action=action,
            output=f"Executed {action} in container {self._container_id}",
        )
        self.record_action(action, params, result)
        return result


class ComputerUseSessionManager:
    """Computer Use 会话管理器

    管理多个 ComputerUseSession 实例的生命周期。
    """

    def __init__(
        self,
        max_sessions: int = 10,
        session_factory: type[ComputerUseSession] | None = None,
    ):
        self._sessions: dict[str, ComputerUseSession] = {}
        self._max_sessions = max_sessions
        self._session_factory = session_factory or InMemoryComputerUseSession

    def get_or_create(
        self,
        session_id: str | None = None,
        **kwargs: object,
    ) -> ComputerUseSession:
        """获取或创建会话"""
        if session_id and session_id in self._sessions:
            return self._sessions[session_id]

        session = self._session_factory(session_id=session_id, **kwargs)
        self._sessions[session.session_id] = session

        # 超过上限时移除最旧会话
        if len(self._sessions) > self._max_sessions:
            oldest_id = next(iter(self._sessions))
            self.remove(oldest_id)

        return session

    def get(self, session_id: str) -> ComputerUseSession | None:
        """获取指定会话"""
        return self._sessions.get(session_id)

    def remove(self, session_id: str) -> None:
        """移除会话"""
        session = self._sessions.pop(session_id, None)
        if session and session.is_started:
            logger.warning("Removing started session: %s", session_id[:8])

    def list_sessions(self) -> list[str]:
        """列出所有会话 ID"""
        return list(self._sessions.keys())

    def has_session(self, session_id: str) -> bool:
        """检查会话是否存在"""
        return session_id in self._sessions

    async def close_all(self) -> None:
        """关闭所有会话"""
        for session in self._sessions.values():
            if session.is_started:
                await session.stop()
        self._sessions.clear()
