"""ComputerUseTool - Anthropic Computer Use API 集成

封装 Anthropic Computer Use API 调用，支持截屏识别、UI 操作、降级策略。
操作类型：screenshot / click / type / scroll / drag / key / wait
降级链：ComputerUseTool 失败 → 检查是否有 API/CLI 替代 → ShellTool → AskHumanTool
"""

from __future__ import annotations

import asyncio
import base64
import logging
from typing import Any, Callable, Awaitable

import httpx

from agentkit.tools.base import Tool
from agentkit.tools.computer_use_session import (
    ComputerUseSession,
    InMemoryComputerUseSession,
    ComputerUseSessionManager,
    ActionResult,
)
from agentkit.tools.computer_use_recorder import ComputerUseRecorder

logger = logging.getLogger(__name__)

# 支持的操作类型
_VALID_ACTIONS = ("screenshot", "click", "type", "scroll", "drag", "key", "wait")

# Anthropic Computer Use API 端点
_ANTHROPIC_COMPUTER_USE_URL = "https://api.anthropic.com/v1/messages"

# 降级建议映射：操作 → 可能的 Shell 替代
_FALLBACK_SHELL_SUGGESTIONS: dict[str, list[str]] = {
    "click": ["考虑使用 xdotool 或 API 替代点击操作"],
    "type": ["考虑使用 xdotool type 或 API 替代输入操作"],
    "scroll": ["考虑使用 xdotool scroll 或键盘 PageDown/PageUp"],
    "drag": ["考虑使用 xdotool mousemove 替代拖拽操作"],
    "key": ["考虑使用 xdotool key 或 API 替代键盘操作"],
    "screenshot": ["考虑使用 scrot/screenshot 命令或 API 替代截屏"],
}


from agentkit.utils.security import is_safe_url as _is_safe_url


class ComputerUseTool(Tool):
    """Computer Use 工具

    封装 Anthropic Computer Use API 调用，支持截屏识别和 UI 操作。
    支持降级链：API 失败 → Shell 替代建议 → AskHuman。

    Usage:
        tool = ComputerUseTool(api_key="sk-...")
        result = await tool.execute(action="screenshot")
        result = await tool.execute(action="click", x=100, y=200)
        result = await tool.execute(action="type", text="hello")
    """

    def __init__(
        self,
        name: str = "computer_use",
        description: str = "Anthropic Computer Use API 集成，支持截屏识别和 UI 操作",
        input_schema: dict[str, Any] | None = None,
        output_schema: dict[str, Any] | None = None,
        version: str = "1.0.0",
        tags: list[str] | None = None,
        api_key: str | None = None,
        model: str = "claude-sonnet-4-20250514",
        api_base_url: str = _ANTHROPIC_COMPUTER_USE_URL,
        session_factory: type[ComputerUseSession] | None = None,
        recorder: ComputerUseRecorder | None = None,
        fallback_callback: Callable[[str, dict[str, Any]], Awaitable[dict[str, Any]]] | None = None,
        max_retries: int = 1,
        request_timeout: float = 30.0,
    ):
        super().__init__(
            name=name,
            description=description,
            input_schema=input_schema or self._default_input_schema(),
            output_schema=output_schema or self._default_output_schema(),
            version=version,
            tags=tags or ["computer-use", "ui-automation", "anthropic"],
        )
        self._api_key = api_key
        self._model = model
        self._api_base_url = api_base_url
        if not _is_safe_url(self._api_base_url):
            raise ValueError(f"Unsafe api_base_url: {self._api_base_url}. Private/internal URLs are not allowed.")
        self._session_manager = ComputerUseSessionManager(
            session_factory=session_factory or InMemoryComputerUseSession,
        )
        self._recorder = recorder or ComputerUseRecorder()
        self._fallback_callback = fallback_callback
        self._max_retries = max_retries
        self._request_timeout = request_timeout
        self._http_client: httpx.AsyncClient | None = None
        self._client_lock = asyncio.Lock()

    async def _get_http_client(self) -> httpx.AsyncClient:
        """Get or create a persistent httpx.AsyncClient for connection reuse."""
        if self._http_client is None or self._http_client.is_closed:
            async with self._client_lock:
                if self._http_client is None or self._http_client.is_closed:
                    self._http_client = httpx.AsyncClient(timeout=self._request_timeout)
        return self._http_client

    async def close(self) -> None:
        """Close the persistent HTTP client."""
        if self._http_client and not self._http_client.is_closed:
            await self._http_client.aclose()

    @staticmethod
    def _default_input_schema() -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "action": {
                    "type": "string",
                    "enum": list(_VALID_ACTIONS),
                    "description": "操作类型：screenshot/click/type/scroll/drag/key/wait",
                },
                "x": {
                    "type": "integer",
                    "description": "X 坐标（click/drag 操作需要）",
                },
                "y": {
                    "type": "integer",
                    "description": "Y 坐标（click/drag 操作需要）",
                },
                "text": {
                    "type": "string",
                    "description": "输入文本（type 操作需要）",
                },
                "key_name": {
                    "type": "string",
                    "description": "按键名称（key 操作需要），如 Enter、Tab、Escape",
                },
                "direction": {
                    "type": "string",
                    "enum": ["up", "down", "left", "right"],
                    "description": "滚动方向（scroll 操作需要）",
                },
                "amount": {
                    "type": "integer",
                    "description": "滚动量（scroll 操作），默认 3",
                    "default": 3,
                },
                "start_x": {
                    "type": "integer",
                    "description": "拖拽起始 X 坐标（drag 操作需要）",
                },
                "start_y": {
                    "type": "integer",
                    "description": "拖拽起始 Y 坐标（drag 操作需要）",
                },
                "end_x": {
                    "type": "integer",
                    "description": "拖拽结束 X 坐标（drag 操作需要）",
                },
                "end_y": {
                    "type": "integer",
                    "description": "拖拽结束 Y 坐标（drag 操作需要）",
                },
                "duration": {
                    "type": "number",
                    "description": "等待时长（wait 操作），默认 1.0 秒",
                    "default": 1.0,
                },
                "session_id": {
                    "type": "string",
                    "description": "会话 ID，指定后在对应会话中执行操作",
                },
            },
            "required": ["action"],
        }

    @staticmethod
    def _default_output_schema() -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "success": {"type": "boolean", "description": "操作是否成功"},
                "action": {"type": "string", "description": "执行的操作类型"},
                "output": {"type": "string", "description": "操作结果描述"},
                "screenshot_base64": {"type": "string", "description": "截图 base64 编码"},
                "error": {"type": "string", "description": "错误信息"},
                "fallback_suggestion": {
                    "type": "string",
                    "description": "降级建议（操作失败时提供）",
                },
                "session_id": {"type": "string", "description": "会话 ID"},
            },
        }

    async def execute(self, **kwargs) -> dict:
        """执行 Computer Use 操作

        Args:
            action: 操作类型（必需）
            x, y: 坐标（click 操作）
            text: 输入文本（type 操作）
            key_name: 按键名称（key 操作）
            direction, amount: 滚动方向和量（scroll 操作）
            start_x, start_y, end_x, end_y: 拖拽坐标（drag 操作）
            duration: 等待时长（wait 操作）
            session_id: 会话 ID

        Returns:
            包含 success, action, output 等字段的字典
        """
        action = kwargs.get("action")
        if not action:
            return self._error_result(
                action="unknown",
                error="action 参数是必需的",
                fallback="指定操作类型：screenshot/click/type/scroll/drag/key/wait",
            )

        if action not in _VALID_ACTIONS:
            return self._error_result(
                action=action,
                error=f"无效的操作类型: {action}，支持: {', '.join(_VALID_ACTIONS)}",
            )

        # 参数验证
        validation_error = self._validate_params(action, kwargs)
        if validation_error:
            return self._error_result(
                action=action,
                error=validation_error,
            )

        session_id = kwargs.get("session_id")

        # 获取或创建会话
        session = self._session_manager.get_or_create(session_id=session_id)
        if not session.is_started:
            await session.start()

        # 提取操作参数（排除 action 和 session_id）
        params = {k: v for k, v in kwargs.items() if k not in ("action", "session_id") and v is not None}

        # 尝试通过 API 执行
        result = await self._execute_with_fallback(session, action, params)

        # 录制操作
        self._recorder.record(action, params, result)

        return self._format_result(result, session.session_id)

    async def _execute_with_fallback(
        self,
        session: ComputerUseSession,
        action: str,
        params: dict[str, Any],
    ) -> ActionResult:
        """带降级链的操作执行

        降级链：Anthropic API → Session 本地执行 → Shell 替代建议 → AskHuman
        """
        # 1. 尝试 Anthropic Computer Use API
        if self._api_key:
            try:
                result = await self._call_anthropic_api(session, action, params)
                if result.success:
                    return result
                logger.warning(
                    "Anthropic API returned failure for %s: %s",
                    action,
                    result.error,
                )
            except Exception as e:
                logger.warning("Anthropic API call failed for %s: %s", action, e)

        # 2. 降级到 Session 本地执行
        try:
            if action == "screenshot":
                result = await session.screenshot()
            else:
                result = await session.execute_action(action, **params)
            if result.success:
                result.metadata["execution_mode"] = "session_local"
                return result
        except Exception as e:
            logger.warning("Session local execution failed for %s: %s", action, e)

        # 3. 尝试自定义降级回调
        if self._fallback_callback:
            try:
                fallback_result = await self._fallback_callback(action, params)
                if fallback_result.get("success"):
                    return ActionResult(
                        success=True,
                        action=action,
                        output=fallback_result.get("output", ""),
                        metadata={"execution_mode": "fallback_callback"},
                    )
            except Exception as e:
                logger.warning("Fallback callback failed for %s: %s", action, e)

        # 4. 返回降级建议
        suggestions = _FALLBACK_SHELL_SUGGESTIONS.get(action, [])
        fallback_msg = "; ".join(suggestions) if suggestions else "考虑使用 ShellTool 或 AskHumanTool 替代"

        return ActionResult(
            success=False,
            action=action,
            error=f"Computer Use 操作失败，降级建议: {fallback_msg}",
            metadata={
                "execution_mode": "fallback_suggestion",
                "fallback_suggestions": suggestions,
            },
        )

    async def _call_anthropic_api(
        self,
        session: ComputerUseSession,
        action: str,
        params: dict[str, Any],
    ) -> ActionResult:
        """调用 Anthropic Computer Use API

        通过 Anthropic Messages API 发送 computer_use_tool 请求。
        """
        # 构造 computer_use 工具调用参数
        tool_input: dict[str, Any] = {"action": action}
        if action == "click":
            tool_input["coordinate"] = [params.get("x", 0), params.get("y", 0)]
        elif action == "type":
            tool_input["text"] = params.get("text", "")
        elif action == "scroll":
            tool_input["direction"] = params.get("direction", "down")
            tool_input["amount"] = params.get("amount", 3)
        elif action == "drag":
            tool_input["start_coordinate"] = [
                params.get("start_x", 0),
                params.get("start_y", 0),
            ]
            tool_input["coordinate"] = [
                params.get("end_x", 0),
                params.get("end_y", 0),
            ]
        elif action == "key":
            tool_input["key"] = params.get("key_name", "")
        elif action == "wait":
            tool_input["duration"] = params.get("duration", 1.0)
        elif action == "screenshot":
            pass  # screenshot 不需要额外参数

        # 获取当前截图作为上下文
        screenshot_result = await session.screenshot()
        screenshot_b64 = screenshot_result.screenshot_base64

        # 构造 API 请求
        content_blocks: list[dict[str, Any]] = []
        if screenshot_b64:
            content_blocks.append({
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/png",
                    "data": screenshot_b64,
                },
            })

        content_blocks.append({
            "type": "text",
            "text": f"Execute computer use action: {action}",
        })

        request_body = {
            "model": self._model,
            "max_tokens": 1024,
            "tools": [
                {
                    "type": "computer_20250124",
                    "name": "computer",
                    "display_width_px": session.screen.width,
                    "display_height_px": session.screen.height,
                }
            ],
            "messages": [
                {
                    "role": "user",
                    "content": content_blocks,
                }
            ],
        }

        headers = {
            "x-api-key": self._api_key,
            "anthropic-version": "2023-06-01",
            "content-type": "application/json",
        }

        client = await self._get_http_client()
        response = await client.post(
            self._api_base_url,
            json=request_body,
            headers=headers,
        )

        if response.status_code != 200:
            error_detail = response.text[:500]
            return ActionResult(
                success=False,
                action=action,
                error=f"Anthropic API error {response.status_code}: {error_detail}",
            )

        data = response.json()

        # 解析 API 响应中的 tool_use 内容
        for block in data.get("content", []):
            if block.get("type") == "tool_use" and block.get("name") == "computer":
                tool_input_resp = block.get("input", {})
                resp_action = tool_input_resp.get("action", action)
                return ActionResult(
                    success=True,
                    action=resp_action,
                    output=f"API executed: {resp_action}",
                    metadata={"api_response": data},
                )

        # API 没有返回 tool_use，可能是纯文本响应
        text_output = ""
        for block in data.get("content", []):
            if block.get("type") == "text":
                text_output += block.get("text", "")

        return ActionResult(
            success=True,
            action=action,
            output=text_output[:500] if text_output else "API call completed",
            metadata={"api_response": data},
        )

    def _validate_params(self, action: str, kwargs: dict[str, Any]) -> str | None:
        """验证操作参数

        Returns:
            错误信息，None 表示验证通过
        """
        if action == "click":
            if "x" not in kwargs or "y" not in kwargs:
                return "click 操作需要 x 和 y 参数"
        elif action == "type":
            if not kwargs.get("text"):
                return "type 操作需要 text 参数"
        elif action == "key":
            if not kwargs.get("key_name"):
                return "key 操作需要 key_name 参数"
        elif action == "drag":
            required = ("start_x", "start_y", "end_x", "end_y")
            missing = [r for r in required if r not in kwargs]
            if missing:
                return f"drag 操作需要 {', '.join(missing)} 参数"
        return None

    def _format_result(self, result: ActionResult, session_id: str) -> dict[str, Any]:
        """格式化操作结果"""
        formatted: dict[str, Any] = {
            "success": result.success,
            "action": result.action,
            "output": result.output,
            "session_id": session_id,
        }
        if result.screenshot_base64:
            formatted["screenshot_base64"] = result.screenshot_base64
        if result.error:
            formatted["error"] = result.error
        if not result.success and result.metadata.get("fallback_suggestions"):
            formatted["fallback_suggestion"] = "; ".join(
                result.metadata["fallback_suggestions"]
            )
        return formatted

    def _error_result(
        self,
        action: str,
        error: str,
        fallback: str = "",
    ) -> dict[str, Any]:
        """构造错误结果"""
        result: dict[str, Any] = {
            "success": False,
            "action": action,
            "error": error,
            "session_id": None,
        }
        if fallback:
            result["fallback_suggestion"] = fallback
        return result

    @property
    def session_manager(self) -> ComputerUseSessionManager:
        """获取会话管理器"""
        return self._session_manager

    @property
    def recorder(self) -> ComputerUseRecorder:
        """获取录制器"""
        return self._recorder
