"""Agent 编排层：OpenAI 兼容模型（DeepSeek 等）经 MCP 工具控制手机。

工作流：
  1. 启动时从 MCP Server 拉工具列表 → 转 OpenAI function schema
  2. run(prompt)：循环 chat/completions
     - 模型返回 tool_calls → 依次执行（经 MCP）→ 结果回喂
     - de_screenshot 的返回图像转为 image_url 追加为下一轮 user 消息（多模态看图）
     - 无 tool_calls → 返回最终文本
"""
import base64
import json
import logging

import httpx
from fastmcp import Client

from mcp_agent.config import AgentSettings

_log = logging.getLogger("agent")

S = AgentSettings()

SYSTEM_PROMPT = """你是手机自动化控制助手。你通过工具实时操作 Android 手机。

工作规范：
1. 先 de_list_devices 确定目标设备（在线才可操作）
2. 观察屏幕：先 de_screenshot 获取截图（图像会随后给你），基于截图理解当前界面
3. 操作：de_tap/de_swipe 的坐标必须与最近一次 de_screenshot 图像一致（直接看图给坐标，服务器自动换算）
4. 每次关键操作后再次 de_screenshot 验证结果，直到完成用户目标
5. 完成或失败时用中文总结：做了什么、当前状态、需要用户注意的事项
6. 设备不可用/操作失败时如实报告错误，不要臆测成功

可用工具清单将由系统提供。"""


class Agent:
    def __init__(self, settings: AgentSettings = None):
        self.s = settings or S
        self.tools_schema = []      # OpenAI function schema
        self._tool_exec = {}        # name -> callable
        self.messages = []
        self.on_step = None         # 可选回调 fn(step_dict)，web 展示进度用

    # ---------- MCP 工具桥 ----------
    async def _load_tools(self):
        """从 MCP Server 拉工具，转 OpenAI function schema。"""
        # 短连接超时：MCP 不可用时快速失败（默认会无限重试卡死线程）
        self._mcp = Client(self.s.mcp_url, timeout=10.0, init_timeout=10.0)
        await self._mcp.__aenter__()
        tools = await self._mcp.list_tools()
        self.tools_schema = []
        for t in tools:
            # MCP SDK v2 改名 input_schema，兼容新旧字段
            schema = getattr(t, "input_schema", None) or getattr(t, "inputSchema", {})
            # fastmcp Tool 属性兼容：name/description/inputSchema
            name = getattr(t, "name", "")
            desc = getattr(t, "description", "") or ""
            self.tools_schema.append({
                "type": "function",
                "function": {"name": name, "description": desc,
                             "parameters": schema}})
            self._tool_exec[name] = t
        _log.info("MCP 工具已加载: %s", [s["function"]["name"] for s in self.tools_schema])

    async def close(self):
        if getattr(self, "_mcp", None):
            await self._mcp.__aexit__(None, None, None)

    # ---------- 模型调用 ----------
    async def _chat(self):
        """调用 OpenAI 兼容 chat/completions，返回完整 response JSON。"""
        body = {
            "model": self.s.model,
            "messages": self.messages,
            "tools": self.tools_schema if self.tools_schema else None,
            "max_tokens": 4096,
        }
        headers = {"Authorization": f"Bearer {self.s.api_key}",
                   "Content-Type": "application/json"}
        async with httpx.AsyncClient(timeout=self.s.request_timeout) as client:
            r = await client.post(f"{self.s.api_base.rstrip('/')}/chat/completions",
                                  json=body, headers=headers)
        if r.status_code != 200:
            raise RuntimeError(f"模型 API HTTP {r.status_code}: {r.text[:300]}")
        return r.json()

    # ---------- 工具执行 ----------
    async def _execute_tool(self, name, arguments):
        """执行 MCP 工具，返回 (文本结果, image_data_or_None)。"""
        args = json.loads(arguments) if isinstance(arguments, str) else (arguments or {})
        _log.info("执行工具 %s %s", name, args)
        try:
            result = await self._mcp.call_tool(name, args)
            data = getattr(result, "data", result)
        except Exception as e:
            return {"ok": False, "error": f"工具执行失败: {e}"}, None
        # de_screenshot：图像分离（作为 image_url 追加给模型看）
        image_b64 = None
        if name == "de_screenshot" and isinstance(data, dict) and data.get("ok"):
            img = (data.get("data") or {}).get("image") or {}
            if img.get("data"):
                text_result = {k: v for k, v in (data.get("data") or {}).items()
                               if k != "image"}
                image_b64 = img["data"]
        if self.on_step:
            try:
                self.on_step({"tool": name, "args": args,
                              "result": text_result if image_b64 else data,
                              "image_b64": image_b64})
            except Exception:
                pass
        if image_b64:
            return text_result, image_b64
        return data, None

    # ---------- 主循环 ----------
    async def run(self, prompt: str, serial: str = "") -> str:
        """执行一轮指令，返回最终回答文本。"""
        target = serial or self.s.default_serial
        sys_txt = SYSTEM_PROMPT
        if target:
            sys_txt += f"\n\n本次默认目标设备 serial：{target}（未指定设备时用它）。"
        self.messages = [{"role": "system", "content": sys_txt},
                         {"role": "user", "content": prompt}]

        for step in range(self.s.max_steps):
            resp = await self._chat()
            choice = (resp.get("choices") or [{}])[0]
            msg = choice.get("message") or {}

            # 1) 工具调用
            tool_calls = msg.get("tool_calls")
            if tool_calls:
                self.messages.append({
                    "role": "assistant",
                    "content": msg.get("content") or "",
                    "tool_calls": tool_calls})
                for tc in tool_calls:
                    fn = tc.get("function") or {}
                    name = fn.get("name", "")
                    text_result, image_b64 = await self._execute_tool(
                        name, fn.get("arguments", "{}"))
                    self.messages.append({
                        "role": "tool",
                        "tool_call_id": tc.get("id", ""),
                        "content": json.dumps(text_result, ensure_ascii=False)[:4000]})
                    # 截图图像：作为下一轮 user 图像内容（OpenAI 协议 tool 结果只能文本）
                    if image_b64:
                        self.messages.append({
                            "role": "user",
                            "content": [{"type": "text",
                                         "text": "这是最新屏幕截图，请基于它继续判断"},
                                        {"type": "image_url",
                                         "image_url": {"url":
                                                       f"data:image/jpeg;base64,{image_b64}"}}]})
                continue

            # 2) 最终回答
            return msg.get("content") or "（模型无输出）"

        return "（达到最大步骤数未完成，请检查操作是否卡在循环）"