from __future__ import annotations

import inspect
import json
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Mapping

from aicd.agent.handlers.registry import tool_handlers
from aicd.config import AppConfig
from aicd.tasks.manager import TaskManager
from aicd.tools.browser_cdp import BrowserCDP
from aicd.tools.docs_convert import DocumentConverter
from aicd.tools.fs_tool import FSTool
from aicd.tools.http_tool import HttpTool
from aicd.tools.path_policy import PathPolicy
from aicd.tools.process_tool import ProcessTool
from aicd.tools.script_tool import ScriptTool
from aicd.tools.text_tool import TextTool

if TYPE_CHECKING:
    from aicd.session_log import SessionLogger

EnvCapabilitiesCallback = Callable[[dict[str, Any], Path], None]
AiWorkspaceSetter = Callable[[str | None], tuple[bool, str]]


class ToolRouter:
    def __init__(
        self,
        cfg: AppConfig,
        *,
        layout: dict[str, Path],
        task_manager: TaskManager,
        cwd: Path | None = None,
        path_policy: PathPolicy | None = None,
        session_log: SessionLogger | None = None,
        chat_session_id: str | None = None,
        chat_session_switch: Callable[[str], None] | None = None,
    ) -> None:
        self._cfg = cfg
        self._layout = layout
        self._tasks = task_manager
        self._session_log = session_log
        self._chat_session_id = chat_session_id
        self._chat_session_switch = chat_session_switch
        self._cwd = cwd or Path.cwd()
        self._policy = path_policy if path_policy is not None else PathPolicy(cfg.paths)
        self._policy.set_workspace_root(self._cwd)
        self._fs = FSTool(self._policy, self._cwd)
        self._text = TextTool(self._fs)
        self._process = ProcessTool(self._policy, self._cwd)
        self._http = HttpTool(layout["cache"])
        self._scripts = ScriptTool(layout["tmp_scripts"], self._process)
        self._docs = DocumentConverter(lambda p: self._policy.check_allowed(p))
        self._browser = BrowserCDP()
        self._ai_output = (self._cwd / "aicd").resolve()
        self._ai_tmp = (self._ai_output / "tmp").resolve()
        self._handlers = tool_handlers()
        self._on_env_capabilities_updated: EnvCapabilitiesCallback | None = None
        self._ai_workspace_setter: AiWorkspaceSetter | None = None
        self._vision_llm: Any = None

    def set_env_capabilities_callback(
        self, fn: EnvCapabilitiesCallback | None
    ) -> None:
        """主 Agent 在 runtime_env_refresh 后更新系统提示中的能力块。"""
        self._on_env_capabilities_updated = fn

    def notify_env_capabilities_updated(
        self, data: Mapping[str, Any], yaml_path: Path
    ) -> None:
        fn = self._on_env_capabilities_updated
        if fn:
            fn(dict(data), yaml_path)

    def set_chat_session_switch(self, fn: Callable[[str], None] | None) -> None:
        """主 Agent 切换会话 id 时同步清空内存上下文（由 AgentLoop 注册）。"""
        self._chat_session_switch = fn

    def set_chat_session_id(self, sid: str) -> None:
        self._chat_session_id = sid

    def set_ai_workspace_setter(self, fn: AiWorkspaceSetter | None) -> None:
        """由 RuntimeContext 注册，使主会话中 ai_workspace_set 工具可切换工作目录。"""
        self._ai_workspace_setter = fn

    def set_vision_llm(self, client: Any | None) -> None:
        """注册带 ``async chat_with_images(prompt, paths) -> dict`` 的 LLM 客户端（多模态）。"""
        self._vision_llm = client

    @property
    def cwd(self) -> Path:
        return self._cwd

    @property
    def path_policy(self) -> PathPolicy:
        return self._policy

    def get_ai_paths(self) -> tuple[Path, Path]:
        return self._ai_output, self._ai_tmp

    def set_ai_paths(self, output_dir: Path, tmp_dir: Path) -> None:
        """AI 产出根目录与其下 tmp（中间数据）；与 fs 的 cwd 独立。"""
        self._ai_output = output_dir.resolve()
        self._ai_tmp = tmp_dir.resolve()

    def set_cwd(self, cwd: Path) -> None:
        """切换主 Agent 工具使用的当前工作目录（相对路径、默认 shell cwd）。"""
        self._cwd = cwd.resolve()
        self._policy.set_workspace_root(self._cwd)
        self._fs = FSTool(self._policy, self._cwd)
        self._text = TextTool(self._fs)
        self._process = ProcessTool(self._policy, self._cwd)
        self._scripts = ScriptTool(self._layout["tmp_scripts"], self._process)

    def _resolve_path(self, path_str: str) -> Path:
        p = Path(path_str).expanduser()
        if not p.is_absolute():
            p = (self._cwd / p).resolve()
        else:
            p = p.resolve()
        self._policy.check_allowed(p)
        return p

    async def run(self, name: str, arguments: str | dict[str, Any]) -> dict[str, Any]:
        slog = self._session_log
        try:
            args = (
                json.loads(arguments) if isinstance(arguments, str) else dict(arguments)
            )
        except json.JSONDecodeError as e:
            if slog:
                raw = arguments if isinstance(arguments, str) else json.dumps(arguments)
                slog.write_event(
                    "tool",
                    "ToolRouter",
                    name,
                    str(raw)[:800],
                    f"bad_json err={e}",
                )
            return {"ok": False, "error": f"bad json: {e}"}
        try:
            result = await self._dispatch(name, args)
            if slog:
                from aicd.session_log import summarize_tool_call, summarize_tool_result

                try:
                    sin = summarize_tool_call(name, args)
                    if isinstance(result, dict):
                        sout = summarize_tool_result(name, result)
                    else:
                        raw = json.dumps(result, ensure_ascii=False, default=str)
                        sout = raw if len(raw) <= 4000 else raw[:3997] + "..."
                except Exception as ex:  # noqa: BLE001
                    sin = repr(args)[:2000]
                    sout = f"summarize_error={ex!r} raw={str(result)[:1500]}"
                slog.write_event("tool", "ToolRouter", name, sin, sout)
            return result
        except PermissionError as e:
            if slog:
                slog.write_exception("tool", "ToolRouter", name, e)
            return {"ok": False, "error": str(e)}
        except Exception as e:  # noqa: BLE001
            if slog:
                slog.write_exception("tool", "ToolRouter", name, e)
            return {"ok": False, "error": repr(e)}

    async def _dispatch(self, name: str, args: dict[str, Any]) -> dict[str, Any]:
        fn = self._handlers.get(name)
        if fn is None:
            return {"ok": False, "error": f"unknown tool: {name}"}
        result = fn(self, args)
        if inspect.isawaitable(result):
            return await result
        return result
