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.plugins.registry import PluginRegistry
    from aicd.session_log import SessionLogger

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


def invoke_plugin_tool_callable(
    fn: Callable[..., Any],
    router: ToolRouter,
    args: dict[str, Any],
) -> Any:
    """按签名调用插件 HANDLERS 中的函数，兼容只写 ``(args,)`` 或未写参数的常见笔误。

    约定优先：``def foo(router, args): ...``；亦接受 ``def foo(args):``、``def foo():``。
    仅一个形参且名为 *router* / *tool_router* / *r* 时视为只收 router（少见）。
    """

    try:
        sig = inspect.signature(fn)
    except (ValueError, TypeError):
        return _try_plugin_call_variants(fn, router, args)

    pos = [
        p
        for p in sig.parameters.values()
        if p.kind
        in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
    ]
    var_pos = any(
        p.kind == inspect.Parameter.VAR_POSITIONAL for p in sig.parameters.values()
    )

    if not pos and not var_pos:
        return fn()
    if len(pos) == 1 and not var_pos:
        name = pos[0].name
        if name in ("router", "tool_router", "r"):
            return fn(router)
        return fn(args)
    return _try_plugin_call_variants(fn, router, args)


def _try_plugin_call_variants(
    fn: Callable[..., Any],
    router: ToolRouter,
    args: dict[str, Any],
) -> Any:
    try:
        return fn(router, args)
    except TypeError:
        try:
            return fn(args)
        except TypeError:
            return fn()


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,
        plugin_handlers: dict[str, Any] | 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"],
            block_external_network=bool(cfg.agent.block_external_network),
        )
        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._base_handlers = tool_handlers()
        self._plugin_handlers: dict[str, Any] = dict(plugin_handlers or {})
        self._handlers: dict[str, Any] = {
            **self._base_handlers,
            **self._plugin_handlers,
        }
        self._plugin_registry: PluginRegistry | None = None
        self._plugin_reload_hook: Callable[[], None] | None = None
        self._on_env_capabilities_updated: EnvCapabilitiesCallback | None = None
        self._ai_workspace_setter: AiWorkspaceSetter | None = None
        self._vision_llm: Any = None
        self._llm_client: 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

    def attach_llm_client(self, client: Any | None) -> None:
        """主会话注册 OpenAI 兼容客户端，用于 **chat_session_info** 的有效上下文预算与模型窗探测。"""
        self._llm_client = client

    def effective_context_budget_tokens(self) -> int:
        cap = self._cfg.agent.context_budget_tokens
        lc = self._llm_client
        if lc is not None and hasattr(lc, "effective_context_budget_tokens"):
            try:
                return int(lc.effective_context_budget_tokens(cap))
            except Exception:  # noqa: BLE001
                pass
        return int(cap)

    def model_context_meta(self) -> dict[str, Any]:
        lc = self._llm_client
        if lc is not None and hasattr(lc, "model_context_meta"):
            try:
                return dict(lc.model_context_meta())
            except Exception:  # noqa: BLE001
                pass
        return {
            "model_context_window_tokens": None,
            "context_budget_tokens_configured": self._cfg.agent.context_budget_tokens,
            "context_budget_tokens_effective": self.effective_context_budget_tokens(),
        }

    def set_plugin_registry(self, reg: Any | None) -> None:
        """绑定 HOME 插件注册表（供 plugin_* 工具与热更新）。"""
        self._plugin_registry = reg

    @property
    def plugin_registry(self) -> Any | None:
        return self._plugin_registry

    def plugin_handlers_snapshot(self) -> dict[str, Any]:
        """子 Agent 等与主会话共享的插件处理器副本。"""
        return dict(self._plugin_handlers)

    def refresh_plugin_handlers(self, handlers: dict[str, Any] | None) -> None:
        """热加载后替换插件 handlers（内置不变）。"""
        self._plugin_handlers = dict(handlers or {})
        self._handlers = {**self._base_handlers, **self._plugin_handlers}

    def set_plugin_reload_hook(self, fn: Callable[[], None] | None) -> None:
        """插件注册表变更后同步提示词等（由 runtime 注册）。"""
        self._plugin_reload_hook = fn

    def after_plugin_registry_changed(self) -> None:
        """在 PluginRegistry.reload / set_enabled / restore 之后刷新本路由器与钩子。"""
        reg = self._plugin_registry
        if reg:
            self.refresh_plugin_handlers(reg.plugin_handlers())
        fn = self._plugin_reload_hook
        if fn:
            fn()

    @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}",
                )
            self._maybe_log_plugin_tool_line(
                name,
                level="ERROR",
                source="bad_json",
                message=str(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)
            self._maybe_log_plugin_tool_failure(name, result)
            return result
        except PermissionError as e:
            if slog:
                slog.write_exception("tool", "ToolRouter", name, e)
            self._maybe_log_plugin_tool_line(
                name,
                level="ERROR",
                source="permission",
                message=str(e),
            )
            return {"ok": False, "error": str(e)}
        except Exception as e:  # noqa: BLE001
            if slog:
                slog.write_exception("tool", "ToolRouter", name, e)
            self._maybe_log_plugin_tool_line(
                name,
                level="ERROR",
                source="exception",
                message=repr(e),
            )
            return {"ok": False, "error": repr(e)}

    def _maybe_log_plugin_tool_line(
        self,
        tool_name: str,
        *,
        level: str,
        source: str,
        message: str,
        detail: str | None = None,
    ) -> None:
        from aicd.plugins.plugin_log import plugin_id_from_prefixed_tool, write_plugin_log

        pid = plugin_id_from_prefixed_tool(tool_name)
        if not pid:
            return
        root = self._layout["plugins"] / pid
        write_plugin_log(
            root,
            level=level,
            source=source,
            message=message,
            detail=detail,
            retention_days=self._cfg.plugins.log_retention_days,
        )

    def _maybe_log_plugin_tool_failure(self, tool_name: str, result: Any) -> None:
        from aicd.plugins.plugin_log import plugin_id_from_prefixed_tool, write_plugin_log

        pid = plugin_id_from_prefixed_tool(tool_name)
        if not pid or not isinstance(result, dict) or result.get("ok") is not False:
            return
        root = self._layout["plugins"] / pid
        err = result.get("error", result)
        write_plugin_log(
            root,
            level="ERROR",
            source="tool",
            message=tool_name,
            detail=str(err)[:8000],
            retention_days=self._cfg.plugins.log_retention_days,
        )

    async def _dispatch(self, name: str, args: dict[str, Any]) -> dict[str, Any]:
        fn = self._base_handlers.get(name) or self._plugin_handlers.get(name)
        if fn is None:
            return {"ok": False, "error": f"unknown tool: {name}"}
        is_plugin = name in self._plugin_handlers
        try:
            if is_plugin:
                result = invoke_plugin_tool_callable(fn, self, args)
            else:
                result = fn(self, args)
            if inspect.isawaitable(result):
                result = await result
            return result
        except Exception as e:  # noqa: BLE001
            if not is_plugin:
                raise
            from aicd.plugins.plugin_log import (
                plugin_id_from_prefixed_tool,
                write_plugin_log,
            )

            detail = repr(e)
            pid = plugin_id_from_prefixed_tool(name)
            if pid:
                write_plugin_log(
                    self._layout["plugins"] / pid,
                    level="ERROR",
                    source="handler_exception",
                    message=name,
                    detail=detail[:8000],
                    retention_days=self._cfg.plugins.log_retention_days,
                )
            return {"ok": False, "error": f"plugin handler crashed: {e!r}"}
