# coding=utf-8
"""
工具辅助类和注册表。

提供 BuiltinTool 的注册、查找和管理功能，以及 MCPTool 的平台授权辅助工具。
"""

from __future__ import annotations

import functools
import inspect
from typing import Any, Callable, Dict, List, Optional, Type, Union
from pydantic import BaseModel

from turbo_agent_core.schema.agents import BuiltinTool, MCPTool
from turbo_agent_core.schema.enums import JSON
from turbo_agent_core.schema.external import Platform


# =============================================================================
# BuiltinTool 注册和管理
# =============================================================================

class BuiltinToolRegistration:
    """单个内置工具的注册信息。"""
    
    def __init__(
        self,
        name: str,
        handler: Callable[..., Any],
        description: Optional[str] = None,
        input_schema: Optional[Dict[str, Any]] = None,
        output_schema: Optional[Dict[str, Any]] = None,
    ):
        self.name = name
        self.handler = handler
        self.description = description or handler.__doc__ or ""
        self.input_schema = input_schema or {}
        self.output_schema = output_schema or {}
        
        # 自动从函数签名提取参数信息
        if not self.input_schema:
            self.input_schema = self._extract_schema_from_signature(handler)
    
    @staticmethod
    def _extract_schema_from_signature(func: Callable) -> Dict[str, Any]:
        """从函数签名提取 JSON Schema。"""
        sig = inspect.signature(func)
        properties = {}
        required = []
        
        for param_name, param in sig.parameters.items():
            if param_name == 'self':
                continue
                
            param_info = {"type": "string"}  # 默认类型
            
            # 尝试获取类型注解
            if param.annotation != inspect.Parameter.empty:
                if param.annotation == str:
                    param_info["type"] = "string"
                elif param.annotation == int:
                    param_info["type"] = "integer"
                elif param.annotation == float:
                    param_info["type"] = "number"
                elif param.annotation == bool:
                    param_info["type"] = "boolean"
                elif param.annotation == dict or param.annotation == Dict:
                    param_info["type"] = "object"
                elif param.annotation == list or param.annotation == List:
                    param_info["type"] = "array"
            
            # 检查是否有默认值
            if param.default == inspect.Parameter.empty:
                required.append(param_name)
            else:
                param_info["default"] = param.default
                
            properties[param_name] = param_info
        
        schema = {
            "type": "object",
            "properties": properties,
        }
        if required:
            schema["required"] = required
            
        return schema


class BuiltinToolRegistry:
    """内置工具注册表。
    
    用于在代码层面注册和管理内置工具，支持：
    - 注册/注销内置工具
    - 根据 builtin_name 查找工具
    - 列出所有已注册的工具
    - 装饰器方式注册
    
    示例:
        registry = BuiltinToolRegistry()
        
        # 装饰器方式注册
        @registry.register("calculator")
        def calculator(a: int, b: int, operation: str = "add") -> int:
            if operation == "add":
                return a + b
            elif operation == "multiply":
                return a * b
            ...
        
        # 手动注册
        registry.register_handler("echo", echo_handler)
        
        # 查找和执行
        tool_def = registry.get_tool_def("calculator")
        result = registry.execute("calculator", {"a": 1, "b": 2, "operation": "add"})
    """
    
    _instance: Optional["BuiltinToolRegistry"] = None
    
    def __new__(cls):
        """单例模式，确保全局只有一个注册表实例。"""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        self._initialized = True
        self._tools: Dict[str, BuiltinToolRegistration] = {}
    
    def register(
        self,
        name: Optional[str] = None,
        description: Optional[str] = None,
        input_schema: Optional[Dict[str, Any]] = None,
        output_schema: Optional[Dict[str, Any]] = None,
    ) -> Callable:
        """装饰器方式注册内置工具。
        
        Args:
            name: 工具名称（默认为函数名）
            description: 工具描述（默认为函数 docstring）
            input_schema: 输入参数 schema（自动从函数签名提取）
            output_schema: 输出参数 schema
            
        Returns:
            装饰器函数
        """
        def decorator(func: Callable) -> Callable:
            tool_name = name or func.__name__
            self.register_handler(
                tool_name,
                func,
                description=description,
                input_schema=input_schema,
                output_schema=output_schema,
            )
            return func
        return decorator
    
    def register_handler(
        self,
        name: str,
        handler: Callable[..., Any],
        description: Optional[str] = None,
        input_schema: Optional[Dict[str, Any]] = None,
        output_schema: Optional[Dict[str, Any]] = None,
    ) -> None:
        """手动注册内置工具处理器。
        
        Args:
            name: 工具名称
            handler: 处理函数
            description: 工具描述
            input_schema: 输入参数 schema
            output_schema: 输出参数 schema
        """
        self._tools[name] = BuiltinToolRegistration(
            name=name,
            handler=handler,
            description=description,
            input_schema=input_schema,
            output_schema=output_schema,
        )
    
    def unregister(self, name: str) -> bool:
        """注销内置工具。
        
        Args:
            name: 工具名称
            
        Returns:
            是否成功注销
        """
        if name in self._tools:
            del self._tools[name]
            return True
        return False
    
    def get(self, name: str) -> Optional[BuiltinToolRegistration]:
        """获取工具注册信息。
        
        Args:
            name: 工具名称
            
        Returns:
            工具注册信息，不存在则返回 None
        """
        return self._tools.get(name)
    
    def get_tool_def(
        self,
        name: str,
        tool_id: Optional[str] = None,
        belong_to_project_id: str = "system",
        name_id: Optional[str] = None,
    ) -> Optional[BuiltinTool]:
        """根据名称创建 BuiltinTool 定义对象。
        
        Args:
            name: 工具名称
            tool_id: 工具 ID（默认自动生成）
            belong_to_project_id: 所属项目 ID
            name_id: 工具 name_id（默认使用 name）
            
        Returns:
            BuiltinTool 对象，不存在则返回 None
        """
        reg = self._tools.get(name)
        if not reg:
            return None
        
        import uuid
        return BuiltinTool(
            id=tool_id or f"builtin_{uuid.uuid4().hex[:8]}",
            name=name,
            belongToProjectId=belong_to_project_id,
            name_id=name_id or name,
            description=reg.description,
            builtin_name=name,
            input_schema=reg.input_schema,
            output_schema=reg.output_schema,
        )
    
    def execute(self, name: str, input_data: JSON, **kwargs) -> Any:
        """执行内置工具。
        
        Args:
            name: 工具名称
            input_data: 输入参数
            **kwargs: 额外参数
            
        Returns:
            工具执行结果
            
        Raises:
            KeyError: 工具不存在
            Exception: 执行异常
        """
        reg = self._tools.get(name)
        if not reg:
            raise KeyError(f"Builtin tool '{name}' not found")
        
        return reg.handler(**input_data)
    
    async def aexecute(self, name: str, input_data: JSON, **kwargs) -> Any:
        """异步执行内置工具。
        
        Args:
            name: 工具名称
            input_data: 输入参数
            **kwargs: 额外参数
            
        Returns:
            工具执行结果
        """
        import asyncio
        reg = self._tools.get(name)
        if not reg:
            raise KeyError(f"Builtin tool '{name}' not found")
        
        handler = reg.handler
        # 如果是协程函数，直接调用
        if asyncio.iscoroutinefunction(handler):
            return await handler(**input_data)
        # 否则在线程池中执行
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, functools.partial(handler, **input_data))
    
    def list_tools(self) -> List[str]:
        """列出所有已注册的工具名称。
        
        Returns:
            工具名称列表
        """
        return list(self._tools.keys())
    
    def clear(self) -> None:
        """清空所有注册的工具。"""
        self._tools.clear()


# 全局注册表实例
def get_builtin_tool_registry() -> BuiltinToolRegistry:
    """获取全局内置工具注册表实例。"""
    return BuiltinToolRegistry()


# =============================================================================
# MCPTool 平台授权辅助
# =============================================================================

class MCPToolAuthHelper:
    """MCP 工具平台授权辅助类。
    
    用于检查和验证 MCPTool 的平台授权配置。
    """
    
    @staticmethod
    def requires_platform_auth(tool: MCPTool) -> bool:
        """检查 MCP 工具是否需要平台授权。
        
        Args:
            tool: MCP 工具对象
            
        Returns:
            是否需要平台授权
        """
        return tool.require_platform_auth
    
    @staticmethod
    def get_platform(tool: MCPTool) -> Optional[Platform]:
        """获取 MCP 工具关联的平台。
        
        Args:
            tool: MCP 工具对象
            
        Returns:
            平台对象，如果没有则返回 None
        """
        if not tool.require_platform_auth:
            return None
        return tool.platform
    
    @staticmethod
    def validate_auth_config(tool: MCPTool) -> List[str]:
        """验证 MCP 工具的授权配置是否有效。
        
        Args:
            tool: MCP 工具对象
            
        Returns:
            错误信息列表，空列表表示验证通过
        """
        errors = []
        
        if tool.require_platform_auth:
            if not tool.platform:
                errors.append("Platform is required when require_platform_auth is True")
            else:
                # 检查平台是否有授权方法
                if not tool.platform.authMethods:
                    errors.append(f"Platform '{tool.platform.name}' has no auth methods")
        
        # 检查 server 配置
        if not tool.server:
            errors.append("MCP server configuration is required")
        else:
            # 检查 transport 和对应配置的匹配
            from turbo_agent_core.schema.external import McpTransport
            
            if tool.server.transport in (McpTransport.http, McpTransport.ws):
                if not tool.server.endpoint:
                    errors.append(f"Endpoint is required for {tool.server.transport} transport")
            elif tool.server.transport == McpTransport.stdio:
                if not tool.server.command:
                    errors.append("Command is required for stdio transport")
        
        return errors


# =============================================================================
# 便捷函数
# =============================================================================

def register_builtin_tool(
    name: Optional[str] = None,
    description: Optional[str] = None,
    input_schema: Optional[Dict[str, Any]] = None,
    output_schema: Optional[Dict[str, Any]] = None,
) -> Callable:
    """便捷函数：使用全局注册表注册内置工具。
    
    示例:
        @register_builtin_tool("calculator")
        def calculator(a: int, b: int) -> int:
            return a + b
    """
    registry = get_builtin_tool_registry()
    return registry.register(
        name=name,
        description=description,
        input_schema=input_schema,
        output_schema=output_schema,
    )


def create_builtin_tool_def(
    name: str,
    tool_id: Optional[str] = None,
    belong_to_project_id: str = "system",
) -> Optional[BuiltinTool]:
    """便捷函数：根据已注册的 builtin_name 创建 BuiltinTool 定义。
    
    Args:
        name: 已注册的工具名称
        tool_id: 工具 ID
        belong_to_project_id: 所属项目 ID
        
    Returns:
        BuiltinTool 对象，不存在则返回 None
    """
    registry = get_builtin_tool_registry()
    return registry.get_tool_def(name, tool_id, belong_to_project_id)
