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

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

from __future__ import annotations

import logging
import time
import uuid
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


@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, Any] = 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, Any]] = []

    @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: Any) -> ActionResult:
        """执行 UI 操作

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

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

    def record_action(self, action: str, params: dict[str, Any], 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, Any]]:
        """获取操作历史（副本）"""
        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, Any] = {
            "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: Any) -> 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: Any) -> 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 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: Any) -> 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: Any,
    ) -> 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()
