import json
from collections.abc import Callable
from typing import Any, Literal

from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
from langchain_core.tools import StructuredTool
from pydantic import BaseModel, Field, create_model

from yuxi.channel.protocols import AgentToolProtocol
from yuxi.channel.plugins.registry import ChannelPluginRegistry
from yuxi.utils.logging_config import logger

_TYPE_MAP: dict[str, type] = {
    "string": str,
    "integer": int,
    "number": float,
    "boolean": bool,
    "array": list,
    "object": dict,
}


def _schema_to_pydantic_model(tool_name: str, parameters: dict) -> type[BaseModel]:
    properties = parameters.get("properties", {})
    required = set(parameters.get("required", []))

    fields_def: dict[str, Any] = {}
    for param_name, param_info in properties.items():
        json_type = param_info.get("type", "string")
        enum_values = param_info.get("enum")

        if enum_values:
            param_type = Literal[tuple(enum_values)]  # type: ignore[valid-type]
        else:
            param_type = _TYPE_MAP.get(json_type, str)

        field_kwargs: dict[str, Any] = {"description": param_info.get("description", "")}
        if param_name not in required:
            field_kwargs["default"] = None

        fields_def[param_name] = (param_type, Field(**field_kwargs))

    model_name = f"{tool_name}_args"
    return create_model(model_name, **fields_def)


def _make_channel_tool_executor(channel_type: str, tool_name: str):
    async def _execute(**kwargs):
        plugin = ChannelPluginRegistry.get(channel_type)
        if plugin is None:
            return f"Error: 渠道 '{channel_type}' 未注册"
        if not isinstance(plugin, AgentToolProtocol):
            return f"Error: 渠道 '{channel_type}' 不支持 Agent 工具协议"
        try:
            result = await plugin.execute_agent_tool(tool_name, kwargs, {})
        except Exception:
            logger.exception("Channel tool execution failed: %s.%s", channel_type, tool_name)
            return f"Error: 工具 '{tool_name}' 执行异常"
        if result.get("success"):
            return json.dumps(result.get("result", ""), ensure_ascii=False)
        return f"Error: {result.get('error', 'Unknown error')}"

    return _execute


class ChannelToolMiddleware(AgentMiddleware):
    def __init__(self):
        super().__init__()
        self._cache: dict[tuple[str, int], list[StructuredTool]] = {}

    def _build_tools(self, channel_type: str, tool_schemas: list[dict]) -> list[StructuredTool]:
        schema_hash = hash(json.dumps(tool_schemas, sort_keys=True, ensure_ascii=False))
        cache_key = (channel_type, schema_hash)
        if cache_key in self._cache:
            return self._cache[cache_key]

        tools: list[StructuredTool] = []
        for schema in tool_schemas:
            func_def = schema.get("function", {})
            tool_name = func_def.get("name", "")
            if not tool_name:
                continue
            description = func_def.get("description", "")
            parameters = func_def.get("parameters", {})

            try:
                args_model = _schema_to_pydantic_model(tool_name, parameters)
                executor = _make_channel_tool_executor(channel_type, tool_name)
                tool = StructuredTool.from_function(
                    name=tool_name,
                    description=description,
                    args_schema=args_model,
                    coroutine=executor,
                )
                tools.append(tool)
            except Exception:
                logger.exception("Failed to create channel tool: %s.%s", channel_type, tool_name)

        self._cache[cache_key] = tools
        logger.debug("ChannelToolMiddleware built %d tools for channel %s", len(tools), channel_type)
        return tools

    async def awrap_model_call(
        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
    ) -> ModelResponse:
        runtime_context = request.runtime.context
        channel_type = getattr(runtime_context, "channel_type", None) or ""
        channel_tools = getattr(runtime_context, "channel_tools", None) or []

        if channel_type and channel_tools:
            langchain_tools = self._build_tools(channel_type, channel_tools)
            if langchain_tools:
                tools_by_name = {t.name: t for t in request.tools or []}
                for t in langchain_tools:
                    tools_by_name[t.name] = t
                request = request.override(tools=list(tools_by_name.values()))

        return await handler(request)