"""
Sandbox - 基于 AgentRun multiport 转发方案的 Desktop SDK 核心类

架构说明
--------
1. 创建 Sandbox：使用 agentrun SDK（同 test_agentrun_benchmark.py 方式）
   - agentrun.sandbox.Sandbox.create(template_type=TemplateType.AIO, template_name=...)

2. 访问 envd 兼容层（port 49983）：
   - 从 sandbox.get_cdp_url() 解析出域名和 sandbox_id
   - base_url = http://{domain}/sandboxes/{sid}
   - 所有 envd 请求加 Header: X-Sandbox-Port: 49983（nginx multiport 转发）

3. noVNC URL（port 6080）：
   - 流 URL = {base_url}/vnc.html，加 X-Sandbox-Port: 6080

接口对照（e2b_desktop → desktop_sdk）
--------------------------------------
e2b_desktop.Sandbox.create()          → desktop_sdk.Sandbox.create()
desktop.commands.run(cmd)             → envd Process Start/Connect
desktop.files.read/write(path)        → envd /files REST
desktop.stream.start/stop/get_url()   → x11vnc + noVNC via commands.run
desktop.screenshot()                  → scrot + files.read
desktop.left_click/right_click/...    → xdotool via commands.run
desktop.write(text)                   → xdotool type
desktop.press(key)                    → xdotool key
desktop.launch(app)                   → gtk-launch
desktop.open(path)                    → xdg-open
desktop.wait(ms)                      → sleep
"""

import re
import time
from shlex import quote as shell_quote
from typing import Iterator, Literal, Optional, Tuple, Union, overload
from uuid import uuid4

from .commands import CommandHandle, CommandResult, CommandsManager
from .files import FilesManager
from .stream import StreamManager

# 键名映射表（与 e2b_desktop 保持一致）
KEYS = {
    "alt": "Alt_L",
    "alt_left": "Alt_L",
    "alt_right": "Alt_R",
    "backspace": "BackSpace",
    "break": "Pause",
    "caps_lock": "Caps_Lock",
    "cmd": "Super_L",
    "command": "Super_L",
    "control": "Control_L",
    "control_left": "Control_L",
    "control_right": "Control_R",
    "ctrl": "Control_L",
    "del": "Delete",
    "delete": "Delete",
    "down": "Down",
    "end": "End",
    "enter": "Return",
    "esc": "Escape",
    "escape": "Escape",
    "f1": "F1",
    "f2": "F2",
    "f3": "F3",
    "f4": "F4",
    "f5": "F5",
    "f6": "F6",
    "f7": "F7",
    "f8": "F8",
    "f9": "F9",
    "f10": "F10",
    "f11": "F11",
    "f12": "F12",
    "home": "Home",
    "insert": "Insert",
    "left": "Left",
    "menu": "Menu",
    "meta": "Meta_L",
    "num_lock": "Num_Lock",
    "page_down": "Page_Down",
    "page_up": "Page_Up",
    "pause": "Pause",
    "print": "Print",
    "right": "Right",
    "scroll_lock": "Scroll_Lock",
    "shift": "Shift_L",
    "shift_left": "Shift_L",
    "shift_right": "Shift_R",
    "space": "space",
    "super": "Super_L",
    "super_left": "Super_L",
    "super_right": "Super_R",
    "tab": "Tab",
    "up": "Up",
    "win": "Super_L",
    "windows": "Super_L",
}

MOUSE_BUTTONS = {"left": 1, "right": 3, "middle": 2}

# envd 端口（由 skills-agent-desktop 监听）
ENVD_PORT = 49983
# noVNC 端口
NOVNC_PORT = 6080
# 桌面 DISPLAY 默认值
DEFAULT_DISPLAY = ":1"


def _map_key(key: str) -> str:
    lower_key = key.lower()
    return KEYS.get(lower_key, lower_key)


class Sandbox:
    """
    基于 AgentRun multiport 转发方案的桌面 Sandbox SDK。

    与 e2b_desktop.Sandbox 接口完全兼容，底层通过 agentrun SDK 管理 Sandbox 生命周期。

    快速上手：
        from desktop_sdk import Sandbox
        from agentrun.sandbox import TemplateType, Template, TemplateInput
        from agentrun.sandbox.model import TemplateContainerConfiguration

        # 步骤 1：创建 Template
        Template.create(input=TemplateInput(
            template_type=TemplateType.AIO,
            template_name="my-template",
            container_configuration=TemplateContainerConfiguration(
                image="tf-test-registry.cn-shanghai.cr.aliyuncs.com/agentrun-ci-test/sandbox-computeruse:d318c",
            ),
            execution_role_arn="acs:ram::xxx:role/aliyunfcdefaultrole",
        ))

        # 步骤 2：创建 Sandbox
        desktop = Sandbox.create(template_name="my-template")
        desktop.launch("google-chrome")
        desktop.wait(3000)
        img = desktop.screenshot()
        with open("shot.png", "wb") as f:
            f.write(img)
        desktop.kill()
    """

    # agentrun Sandbox 对象
    _agentrun_sandbox = None
    _sandbox_id: str = ""
    _base_url: str = ""
    _display: str = DEFAULT_DISPLAY
    _commands: Optional[CommandsManager] = None
    _files: Optional[FilesManager] = None
    _stream_mgr: Optional[StreamManager] = None

    # ------------------------------------------------------------------
    # 工厂方法
    # ------------------------------------------------------------------

    @classmethod
    def create(
        cls,
        template_name: str,
        display: str = DEFAULT_DISPLAY,
        wait_timeout: int = 120,
        idle_timeout_seconds: int = 600,
        **kwargs,
    ) -> "Sandbox":
        """
        使用已存在的 AgentRun Template 创建 Sandbox。

        与 test_agentrun_benchmark.py 中的方式完全一致：
        先在外部通过 Template.create() 创建好 template，再调用此方法。

        :param template_name: AgentRun template 名称（需已通过 Template.create() 创建）
        :param display: X11 DISPLAY 变量（默认 :1）
        :param wait_timeout: 等待 Sandbox 就绪的超时秒数
        :param idle_timeout_seconds: Sandbox 空闲超时秒数
        :return: Sandbox 实例

        典型用法（对应 test_agentrun_benchmark.py 流程）：
            from agentrun.sandbox import TemplateType, Template, TemplateInput
            from agentrun.sandbox.model import TemplateContainerConfiguration
            from desktop_sdk import Sandbox

            # 步骤 1：创建 Template
            Template.create(input=TemplateInput(
                template_type=TemplateType.AIO,
                template_name=template_name,
                container_configuration=TemplateContainerConfiguration(
                    image="tf-test-registry.cn-shanghai.cr.aliyuncs.com/agentrun-ci-test/sandbox-computeruse:d318c",
                ),
                execution_role_arn=EXECUTION_ROLE_ARN,
            ))

            # 步骤 2：创建 Sandbox
            desktop = Sandbox.create(template_name=template_name)
        """
        from agentrun.sandbox import Sandbox as ARSandbox, TemplateType

        instance = cls.__new__(cls)
        instance._display = display
        instance._agentrun_sandbox = None
        instance._sandbox_id = ""
        instance._base_url = ""
        instance._commands = None
        instance._files = None
        instance._stream_mgr = None

        # 创建 Sandbox 实例（与 test_agentrun_benchmark.py 完全一致）
        ar_sandbox = ARSandbox.create(
            template_type=TemplateType.AIO,
            template_name=template_name,
            sandbox_idle_timeout_seconds=idle_timeout_seconds,
        )
        instance._agentrun_sandbox = ar_sandbox
        instance._sandbox_id = ar_sandbox.sandbox_id

        # 初始化各 manager（_wait_health 需要它们）
        instance._commands = CommandsManager(instance)
        instance._files = FilesManager(instance)
        instance._stream_mgr = StreamManager(instance)

        # 解析 base_url（_wait_health 需要它）
        instance._base_url = instance._resolve_base_url()

        # 等待 agentrun status 就绪，再等待 /health 200
        instance._wait_ready(max_wait=wait_timeout)

        return instance

    def _wait_ready(self, max_wait: int = 180, interval: int = 3):
        """
        等待 Sandbox 完全就绪：
          1. agentrun status 变为 RUNNING/READY
          2. envd /health 返回 HTTP 200/204
          3. 桌面环境（Xvfb + XFCE）就绪：xdotool getmouselocation 成功
        """
        import logging
        import httpx as _httpx

        logger = logging.getLogger(__name__)
        start = time.time()

        # 阶段 1：等待 agentrun status
        while time.time() - start < max_wait:
            status = getattr(self._agentrun_sandbox, "status", "UNKNOWN")
            if status in ("RUNNING", "READY"):
                break
            logger.debug(f"Sandbox status={status}, waiting...")
            time.sleep(interval)
        else:
            raise TimeoutError(
                f"Sandbox did not enter RUNNING state within {max_wait}s"
            )

        # 阶段 2：轮询 /health 直到 200/204
        health_url = f"{self._base_url}/health"
        headers = self._get_envd_headers()

        logger.info(f"Waiting for /health 200 at {health_url} ...")
        while time.time() - start < max_wait:
            try:
                resp = _httpx.get(health_url, headers=headers, timeout=5)
                if resp.status_code in (200, 204):
                    logger.info(
                        f"envd ready (health OK, elapsed {time.time()-start:.1f}s)"
                    )
                    break
                logger.debug(
                    f"/health returned {resp.status_code}, retrying...")
            except Exception as exc:
                logger.debug(f"/health error: {exc}, retrying...")
            time.sleep(interval)
        else:
            raise TimeoutError(
                f"Sandbox /health did not return 200 within {max_wait}s"
            )

        # 阶段 3：等待 XFCE 桌面完全渲染
        # xdotool 成功只代表 Xvfb 可用，还需等 xfdesktop 进程启动才算桌面就绪
        logger.info("Waiting for XFCE desktop (xfdesktop) to be ready...")
        desktop_ready = False
        while time.time() - start < max_wait:
            try:
                # 检测 xfdesktop 进程，它是 XFCE 桌面渲染的最后一个关键进程
                result = self._commands.run(
                    "pgrep -x xfdesktop",
                    timeout=5,
                )
                if result.exit_code == 0 and result.stdout.strip():
                    # xfdesktop 进程存在，额外等 2s 让它完成初始绘制
                    time.sleep(2)
                    desktop_ready = True
                    break
            except Exception:
                pass
            logger.debug("xfdesktop not ready yet, retrying...")
            time.sleep(2)

        if not desktop_ready:
            logger.warning(
                "xfdesktop did not start within timeout, proceeding anyway"
            )
        else:
            logger.info(
                f"Sandbox fully ready (xfdesktop OK, total wait {time.time()-start:.1f}s)"
            )

    def _resolve_base_url(self) -> str:
        """从 agentrun Sandbox 解析出 HTTP base URL"""
        # 优先通过 CDP URL 解析
        try:
            cdp_url = self._agentrun_sandbox.get_cdp_url()
            if cdp_url:
                # 格式：ws(s)://{domain}/sandboxes/{sid}/ws/automation...
                #        http(s)://{domain}/sandboxes/{sid}/...
                m = re.search(
                    r"(?:wss?|https?)://(.+?)/sandboxes/(.+?)/", cdp_url)
                if m:
                    domain = m.group(1)
                    sid = m.group(2)
                    return f"http://{domain}/sandboxes/{sid}"
        except Exception:
            pass

        # 备用：通过 sandbox_id 和 account_id/region 构造
        try:
            ar = self._agentrun_sandbox
            account_id = getattr(ar, "account_id", None) or getattr(
                ar, "_account_id", None)
            region = getattr(ar, "region", None) or getattr(
                ar, "_region", None)
            if account_id and region:
                sid = self._sandbox_id
                return f"http://{account_id}.agentrun-data.{region}.aliyuncs.com/sandboxes/{sid}"
        except Exception:
            pass

        raise RuntimeError(
            "Unable to resolve sandbox base URL from CDP URL or sandbox metadata")

    # ------------------------------------------------------------------
    # 内部工具
    # ------------------------------------------------------------------

    def _get_envd_base_url(self) -> str:
        """返回 envd 兼容层的 base URL（直接使用 sandbox base URL）"""
        return self._base_url

    def _get_envd_headers(self) -> dict:
        """返回访问 envd 兼容层所需的 HTTP headers（multiport 转发）"""
        return {"X-Sandbox-Port": str(ENVD_PORT)}

    def _get_novnc_url(self) -> str:
        """返回 noVNC web 界面 URL（通过 multiport 转发 6080）"""
        return f"{self._base_url}/vnc.html"

    # ------------------------------------------------------------------
    # 公开属性
    # ------------------------------------------------------------------

    @property
    def sandbox_id(self) -> str:
        return self._sandbox_id

    @property
    def commands(self) -> CommandsManager:
        return self._commands

    @property
    def files(self) -> FilesManager:
        return self._files

    @property
    def stream(self) -> StreamManager:
        return self._stream_mgr

    # ------------------------------------------------------------------
    # 截图
    # ------------------------------------------------------------------

    @overload
    def screenshot(self, format: Literal["stream"]) -> Iterator[bytes]: ...

    @overload
    def screenshot(self, format: Literal["bytes"] = ...) -> bytearray: ...

    def screenshot(
        self,
        format: Literal["bytes", "stream"] = "bytes",
    ):
        """
        截取当前桌面截图。

        :param format: 返回格式：'bytes'（bytearray）| 'stream'（字节流生成器）
        :return: 截图数据
        """
        screenshot_path = f"/tmp/screenshot-{uuid4()}.png"
        self.commands.run(
            f"DISPLAY={self._display} scrot --pointer {screenshot_path}",
            timeout=15,
        )
        data = self.files.read(screenshot_path, format=format)
        try:
            self.files.remove(screenshot_path, force=True)
        except Exception:
            pass
        return data

    # ------------------------------------------------------------------
    # 鼠标控制
    # ------------------------------------------------------------------

    def left_click(self, x: Optional[int] = None, y: Optional[int] = None):
        """左键单击"""
        if x is not None and y is not None:
            self.move_mouse(x, y)
        self.commands.run(f"DISPLAY={self._display} xdotool click 1")

    def double_click(self, x: Optional[int] = None, y: Optional[int] = None):
        """左键双击"""
        if x is not None and y is not None:
            self.move_mouse(x, y)
        self.commands.run(
            f"DISPLAY={self._display} xdotool click --repeat 2 1")

    def right_click(self, x: Optional[int] = None, y: Optional[int] = None):
        """右键单击"""
        if (x is None) != (y is None):
            raise ValueError("Both x and y must be provided together")
        if x is not None and y is not None:
            self.move_mouse(x, y)
        self.commands.run(f"DISPLAY={self._display} xdotool click 3")

    def middle_click(self, x: Optional[int] = None, y: Optional[int] = None):
        """中键单击"""
        if x is not None and y is not None:
            self.move_mouse(x, y)
        self.commands.run(f"DISPLAY={self._display} xdotool click 2")

    def scroll(self, direction: Literal["up", "down"] = "down", amount: int = 1):
        """
        滚动鼠标滚轮。

        :param direction: 'up' 向上 | 'down' 向下
        :param amount: 滚动次数
        """
        btn = "4" if direction == "up" else "5"
        self.commands.run(
            f"DISPLAY={self._display} xdotool click --repeat {amount} {btn}"
        )

    def move_mouse(self, x: int, y: int):
        """移动鼠标到指定坐标"""
        self.commands.run(
            f"DISPLAY={self._display} xdotool mousemove --sync {x} {y}")

    def mouse_press(self, button: Literal["left", "right", "middle"] = "left"):
        """按下鼠标键（不释放）"""
        self.commands.run(
            f"DISPLAY={self._display} xdotool mousedown {MOUSE_BUTTONS[button]}"
        )

    def mouse_release(self, button: Literal["left", "right", "middle"] = "left"):
        """释放鼠标键"""
        self.commands.run(
            f"DISPLAY={self._display} xdotool mouseup {MOUSE_BUTTONS[button]}"
        )

    def drag(self, fr: Tuple[int, int], to: Tuple[int, int]):
        """
        拖拽鼠标。

        :param fr: 起始坐标 (x, y)
        :param to: 目标坐标 (x, y)
        """
        self.move_mouse(fr[0], fr[1])
        self.mouse_press()
        self.move_mouse(to[0], to[1])
        self.mouse_release()

    def get_cursor_position(self) -> Tuple[int, int]:
        """获取当前鼠标坐标"""
        result = self.commands.run(
            f"DISPLAY={self._display} xdotool getmouselocation"
        )
        groups = re.search(r"x:(\d+)\s+y:(\d+)", result.stdout)
        if not groups:
            raise RuntimeError(
                f"Failed to parse cursor position from output: {result.stdout}"
            )
        return int(groups.group(1)), int(groups.group(2))

    # ------------------------------------------------------------------
    # 键盘控制
    # ------------------------------------------------------------------

    def write(self, text: str, *, chunk_size: int = 25, delay_in_ms: int = 75) -> None:
        """
        在当前光标位置输入文本。

        :param text: 要输入的文本
        :param chunk_size: 每次输入的字符数
        :param delay_in_ms: 每次输入之间的延迟（毫秒）
        """
        def chunks(s, n):
            for i in range(0, len(s), n):
                yield s[i:i + n]

        for chunk in chunks(text, chunk_size):
            self.commands.run(
                f"DISPLAY={self._display} xdotool type --delay {delay_in_ms} -- {shell_quote(chunk)}"
            )

    def press(self, key: Union[str, list]):
        """
        按键（单个或组合键）。

        :param key: 键名字符串（如 'enter'、'ctrl'）或键名列表（组合键）
        """
        if isinstance(key, list):
            key_str = "+".join(_map_key(k) for k in key)
        else:
            key_str = _map_key(key)
        self.commands.run(f"DISPLAY={self._display} xdotool key {key_str}")

    # ------------------------------------------------------------------
    # 窗口控制
    # ------------------------------------------------------------------

    def get_screen_size(self) -> Tuple[int, int]:
        """获取屏幕分辨率"""
        result = self.commands.run(
            f"DISPLAY={self._display} xrandr"
        )
        m = re.search(r"(\d+x\d+)", result.stdout)
        if not m:
            raise RuntimeError(
                f"Failed to parse screen size from xrandr output: {result.stdout}"
            )
        w, h = m.group(1).split("x")
        return int(w), int(h)

    def get_current_window_id(self) -> str:
        """获取当前活动窗口的 ID"""
        result = self.commands.run(
            f"DISPLAY={self._display} xdotool getwindowfocus"
        )
        return result.stdout.strip()

    def get_application_windows(self, application: str) -> list:
        """
        获取指定应用程序所有可见窗口的 ID 列表。

        :param application: 应用程序 class 名称（如 'Firefox'、'Google-chrome'）
        """
        result = self.commands.run(
            f"DISPLAY={self._display} xdotool search --onlyvisible --class {application}"
        )
        return [wid for wid in result.stdout.strip().split("\n") if wid]

    def get_window_title(self, window_id: str) -> str:
        """获取指定窗口 ID 的标题"""
        result = self.commands.run(
            f"DISPLAY={self._display} xdotool getwindowname {window_id}"
        )
        return result.stdout.strip()

    # ------------------------------------------------------------------
    # 应用控制
    # ------------------------------------------------------------------

    def launch(self, application: str, uri: Optional[str] = None):
        """
        启动应用程序。

        envd 以 root 运行，gtk-launch/xdg-open 在 root 下找不到 user 的 .desktop 关联，
        因此改用 runuser 切换到 user 身份，通过 gtk-launch 查找 /usr/share/applications/
        和 ~/.local/share/applications/ 下的 .desktop 文件。

        :param application: .desktop 文件名（不带 .desktop 后缀），如 'google-chrome'、'code'
        :param uri: 可选的 URI 参数
        """
        uri_part = f" {uri}" if uri else ""
        # runuser 切换到 user，保持 DISPLAY 和 XDG 环境变量
        cmd = (
            f"runuser -u user -- env "
            f"DISPLAY={self._display} "
            f"XDG_RUNTIME_DIR=/run/user/1000 "
            f"HOME=/home/user "
            f"gtk-launch {application}{uri_part}"
        )
        self.commands.run(cmd, background=True, timeout=0)

    def open(self, file_or_url: str):
        """
        用默认应用打开文件或 URL（xdg-open）。

        envd 以 root 运行，须切换到 user 身份才能找到 MIME 关联。

        :param file_or_url: 文件路径或 URL
        """
        cmd = (
            f"runuser -u user -- env "
            f"DISPLAY={self._display} "
            f"XDG_RUNTIME_DIR=/run/user/1000 "
            f"HOME=/home/user "
            f"xdg-open {file_or_url}"
        )
        self.commands.run(cmd, background=True, timeout=0)

    def wait(self, ms: int):
        """
        等待指定毫秒数。

        :param ms: 等待时长（毫秒）
        """
        self.commands.run(f"sleep {ms / 1000}",
                          timeout=max(int(ms / 1000) + 10, 10))

    # ------------------------------------------------------------------
    # 健康检查
    # ------------------------------------------------------------------

    def check_health(self) -> dict:
        """检查 sandbox 健康状态（通过 envd /health）"""
        import httpx
        try:
            base = self._get_envd_base_url()
            headers = self._get_envd_headers()
            resp = httpx.get(f"{base}/health", headers=headers, timeout=10)
            if resp.status_code in (200, 204):
                return {"status": "ok"}
            return {"status": "error", "code": resp.status_code}
        except Exception as e:
            return {"status": "error", "error": str(e)}

    # ------------------------------------------------------------------
    # 生命周期
    # ------------------------------------------------------------------

    def kill(self):
        """销毁 Sandbox"""
        try:
            if self._agentrun_sandbox is not None:
                if hasattr(self._agentrun_sandbox, "delete"):
                    self._agentrun_sandbox.delete()
                elif hasattr(self._agentrun_sandbox, "stop"):
                    self._agentrun_sandbox.stop()
        except Exception:
            pass

    def __repr__(self) -> str:
        return f"<Sandbox sandbox_id={self._sandbox_id!r}>"
