"""动态工具加载器

从配置文件加载工具，支持：
- YAML 配置热加载
- 多来源工具发现（内置、本地插件、pip 工具、MCP）
- 插件式架构
"""

import logging
import os
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional

import yaml

from .base import BaseTool, FunctionCallDefinition, ToolCollection, ToolOutput
from .code_generator import ToolCodeGenerator
from .plugin_manager import get_plugin_manager
from .validator import ToolCodeValidator

logger = logging.getLogger("agent")


# 尝试导入新的插件注册表
try:
    from .registry import (
        PluginToolRegistry,
        get_plugin_registry,
        reload_plugin_registry,
    )
    PLUGIN_REGISTRY_AVAILABLE = True
except ImportError:
    PLUGIN_REGISTRY_AVAILABLE = False
    logger.warning("[ToolLoader] 插件注册表不可用")


# ============================================================================
# 配置模型
# ============================================================================


@dataclass
class ToolConfig:
    """单个工具配置"""
    name: str
    enabled: bool = True
    category: str = "general"
    config: Dict[str, Any] = field(default_factory=dict)


@dataclass
class ToolLoaderConfig:
    """工具加载器配置"""
    config_path: str = "tools.yaml"
    builtin_tools: List[str] = field(default_factory=list)
    mcp_servers: Dict[str, Any] = field(default_factory=dict)


# ============================================================================
# 工具注册表
# ============================================================================


class ToolRegistry:
    """工具注册表 - 管理工具定义和实例"""

    def __init__(self):
        self._definitions: Dict[str, FunctionCallDefinition] = {}
        self._executors: Dict[str, Callable] = {}
        self._categories: Dict[str, List[str]] = {}

    def register(
        self,
        name: str,
        description: str,
        input_schema: Dict[str, Any],
        executor: Callable,
        category: str = "general"
    ) -> None:
        """注册工具

        Args:
            name: 工具名称
            description: 工具描述
            input_schema: JSON Schema 参数定义
            executor: 执行函数 (parameters, context) -> result
            category: 工具分类
        """
        # 注册定义
        self._definitions[name] = FunctionCallDefinition(
            name=name,
            description=description,
            parameters=input_schema
        )

        # 注册执行器
        self._executors[name] = executor

        # 注册分类
        if category not in self._categories:
            self._categories[category] = []
        self._categories[category].append(name)

        logger.debug(f"[ToolRegistry] 注册工具: {name} ({category})")

    def get_definition(self, name: str) -> Optional[FunctionCallDefinition]:
        """获取工具定义"""
        return self._definitions.get(name)

    def get_executor(self, name: str) -> Optional[Callable]:
        """获取工具执行器"""
        return self._executors.get(name)

    def list_definitions(self, category: Optional[str] = None) -> List[FunctionCallDefinition]:
        """列出工具定义"""
        if category:
            names = self._categories.get(category, [])
            return [self._definitions[n] for n in names if n in self._definitions]
        return list(self._definitions.values())

    def list_categories(self) -> List[str]:
        """列出所有分类"""
        return list(self._categories.keys())


# ============================================================================
# 动态工具加载器
# ============================================================================


class ToolLoader:
    """动态工具加载器

    支持:
    - 从 YAML 配置文件加载工具配置
    - 动态注册和卸载工具
    - 工具配置热加载
    - 插件式工具来源（内置、本地、pip、MCP）
    """

    DEFAULT_CONFIG = """
# Agent 工具配置
# 所有启用的工具将加载到 Agent 中

# 工具来源配置
sources:
  builtin: true    # 内置工具
  local: true     # 本地插件 (~/.jude_agent/plugins/)
  pip: true       # pip 安装的工具
  mcp: false      # MCP 服务器

# 启用/禁用特定工具
tools:
  calculator: true
  search_documents: true
  web_search: false
  get_current_time: false
  read_file: false
  write_file: false
  list_directory: false
  fetch_url: false
  execute_command: false
  create_tool: true
  list_plugins: true
  delete_plugin: false

# MCP 服务器配置
mcp_servers: {}

# 工具分类配置
categories:
  search:
    - web_search
    - search_documents
  datetime:
    - get_current_time
  file:
    - read_file
    - write_file
    - list_directory
  network:
    - fetch_url
  system:
    - execute_command
  general:
    - calculator
"""

    def __init__(self, config_path: Optional[str] = None):
        global _logged_messages

        # 从全局获取日志去重集合
        self._logged_messages = _logged_messages

        self.config_path = config_path or self._find_config_path()
        self._config: Dict[str, Any] = {}
        self._registry = ToolRegistry()
        self._tools: Dict[str, BaseTool] = {}
        self._langchain_tools: List = []

        # 插件注册表（如果可用）
        self._plugin_registry: Optional[PluginToolRegistry] = None

        # 工具版本号（用于检测动态添加工具）
        self._tool_version: int = 0

        # 加载配置
        self._load_config()

        # 初始化插件注册表（如果可用）
        if PLUGIN_REGISTRY_AVAILABLE:
            self._init_plugin_registry()

        # 注册内置工具
        self._register_builtin_tools()

    def _find_config_path(self) -> str:
        """查找配置文件"""
        # 检查项目根目录
        for path in ["tools.yaml", "config/tools.yaml", ".config/tools.yaml"]:
            if os.path.exists(path):
                return path
        return "tools.yaml"

    def _log_once(self, key: str, message: str):
        """只输出一次日志"""
        global _logged_messages
        if key not in _logged_messages:
            _logged_messages.add(key)
            logger.debug(message)

    def _load_config(self) -> None:
        """加载配置文件"""
        # 避免重复日志
        if hasattr(self, '_config_loaded') and self._config_loaded:
            return
        self._config_loaded = True

        if os.path.exists(self.config_path):
            try:
                with open(self.config_path, 'r', encoding='utf-8') as f:
                    self._config = yaml.safe_load(f) or {}
                self._log_once('load_config', f"[ToolLoader] 加载配置: {self.config_path}")
            except Exception as e:
                logger.warning(f"[ToolLoader] 配置加载失败: {e}，使用默认配置")
                self._config = yaml.safe_load(self.DEFAULT_CONFIG) or {}
        else:
            self._log_once('load_config_default', "[ToolLoader] 配置文件不存在，使用默认配置")
            self._config = yaml.safe_load(self.DEFAULT_CONFIG) or {}

    def _register_builtin_tools(self) -> None:
        """注册内置工具"""
        # 避免重复注册
        if hasattr(self, '_builtin_registered') and self._builtin_registered:
            return
        self._builtin_registered = True

        from .builtin import (
            CalculatorTool,
            CreateToolTool,
            DeletePluginTool,
            ExecuteCommandTool,
            FetchUrlTool,
            GetCurrentTimeTool,
            ListDirectoryTool,
            ListPluginsTool,
            ReadFileTool,
            SearchDocumentsTool,
            WebSearchTool,
            WriteFileTool,
        )

        # 工具类映射
        tool_classes = {
            "calculator": CalculatorTool,
            "search_documents": SearchDocumentsTool,
            "web_search": WebSearchTool,
            "get_current_time": GetCurrentTimeTool,
            "read_file": ReadFileTool,
            "write_file": WriteFileTool,
            "list_directory": ListDirectoryTool,
            "fetch_url": FetchUrlTool,
            "execute_command": ExecuteCommandTool,
            "create_tool": CreateToolTool,
            "list_plugins": ListPluginsTool,
            "delete_plugin": DeletePluginTool,
        }

        # 从配置获取启用状态 - 支持新旧两种配置格式
        # 新格式: tools.calculator
        # 旧格式: builtin.calculator
        tools_config = self._config.get("tools", {})
        builtin_config = self._config.get("builtin", {})

        for tool_name, tool_class in tool_classes.items():
            # 优先使用新格式，其次使用旧格式
            enabled = tools_config.get(tool_name, builtin_config.get(tool_name, True))
            if enabled:
                try:
                    tool = tool_class()
                    self._tools[tool_name] = tool
                    self._registry.register(
                        name=tool.name,
                        description=tool.description,
                        input_schema=tool.input_schema,
                        executor=lambda p, c, t=tool: t.execute(p, c),
                        category=tool.category
                    )
                    self._log_once(f'tool_{tool_name}', f"[ToolLoader] 加载内置工具: {tool_name}")
                except Exception as e:
                    logger.warning(f"[ToolLoader] 加载工具 {tool_name} 失败: {e}")

    def _init_plugin_registry(self) -> None:
        """初始化插件注册表"""
        # 避免重复初始化
        if hasattr(self, '_plugin_registry_initialized') and self._plugin_registry_initialized:
            return
        self._plugin_registry_initialized = True

        # 获取来源配置
        sources_config = self._config.get("sources", {
            "builtin": True,
            "local": True,
            "pip": True,
            "mcp": False,
        })

        # 创建插件注册表
        self._plugin_registry = get_plugin_registry()

        # 根据配置启用/禁用来源
        for source, enabled in sources_config.items():
            if enabled:
                self._plugin_registry.enable_source(source)
            else:
                self._plugin_registry.disable_source(source)

        self._log_once('plugin_registry', f"[ToolLoader] 插件注册表已初始化，启用的来源: {list(sources_config.keys())}")

    def discover_plugins(self) -> List[str]:
        """发现可用的插件"""
        if not self._plugin_registry:
            if PLUGIN_REGISTRY_AVAILABLE:
                self._init_plugin_registry()
            else:
                return []

        packages = self._plugin_registry.discover_all()
        return [pkg.name for pkg in packages]

    def load_plugins(self) -> int:
        """加载所有启用的插件工具，返回加载数量"""
        if not self._plugin_registry:
            return 0

        # 从配置获取启用的工具列表
        tools_config = self._config.get("tools", {})
        enabled_tools = [name for name, enabled in tools_config.items() if enabled]

        # 加载插件工具
        collection = self._plugin_registry.load_all(enabled_tools)

        # 将加载的工具添加到当前加载器
        for tool in collection.list_tools():
            self._tools[tool.name] = tool
            self._registry.register(
                name=tool.name,
                description=tool.description,
                input_schema=tool.input_schema,
                executor=lambda p, c, t=tool: t.execute(p, c),
                category=tool.category
            )

        return len(collection)

    def reload(self) -> None:
        """重新加载配置和工具"""
        # 重置去重标志，允许重新加载时输出日志
        self._config_loaded = False
        self._plugin_registry_initialized = False
        self._builtin_registered = False

        logger.debug("[ToolLoader] 重新加载配置...")
        self._load_config()

        # 清空并重新注册
        self._tools.clear()
        self._registry = ToolRegistry()
        self._langchain_tools = []

        # 重新初始化插件注册表
        if PLUGIN_REGISTRY_AVAILABLE:
            reload_plugin_registry()
            self._init_plugin_registry()

        # 注册内置工具
        self._register_builtin_tools()

        # 加载插件工具
        if self._plugin_registry:
            self.load_plugins()

        # 重建 LangChain 工具列表
        self._langchain_tools = self.get_langchain_tools()

    def generate_tool(
        self,
        name: str,
        description: str,
        input_schema: Dict[str, Any],
        category: str = "general",
        code: Optional[str] = None,
        validate: bool = True,
        auto_load: bool = True
    ) -> ToolOutput:
        """动态生成并加载工具

        Args:
            name: 工具名称
            description: 工具描述
            input_schema: 输入参数 schema
            category: 工具分类
            code: 自定义代码（可选，如果不提供则使用模板生成）
            validate: 是否验证代码
            auto_load: 是否自动加载到工具集合

        Returns:
            ToolOutput: 操作结果
        """
        try:
            # 1. 生成代码
            generator = ToolCodeGenerator()
            if code:
                tool_code = generator.generate_with_code(
                    name=name,
                    description=description,
                    input_schema=input_schema,
                    code=code,
                    category=category
                )
            else:
                tool_code = generator.generate(
                    name=name,
                    description=description,
                    input_schema=input_schema,
                    category=category
                )

            # 2. 验证代码
            if validate:
                validator = ToolCodeValidator()
                result = validator.validate_full(tool_code, name)
                if not result.valid:
                    return ToolOutput(
                        success=False,
                        result=None,
                        error="代码验证失败: " + "; ".join(result.errors)
                    )

            # 3. 保存插件
            plugin_manager = get_plugin_manager()
            plugin_path = plugin_manager.save_plugin(
                name=name,
                code=tool_code,
                description=description,
                category=category
            )

            # 4. 加载工具
            if auto_load:
                tool = plugin_manager.load_plugin(name)
                if tool:
                    self._tools[tool.name] = tool
                    self._registry.register(
                        name=tool.name,
                        description=tool.description,
                        input_schema=tool.input_schema,
                        executor=lambda p, c, t=tool: t.execute(p, c),
                        category=tool.category
                    )
                    # 清空 LangChain 缓存
                    self._langchain_tools = []
                    # 设置版本为 0，强制 Agent 刷新工具列表
                    self._tool_version = 0
                    logger.debug(f"[ToolLoader] 动态工具已加载: {name}")

                    # 刷新全局 Agent（如果存在），使其在当前会话中立即可以使用新工具
                    self._notify_agent_refresh()
                else:
                    return ToolOutput(
                        success=False,
                        result=None,
                        error="工具加载失败"
                    )

            return ToolOutput(
                success=True,
                result=f"工具已创建并保存到 {plugin_path}",
                metadata={"name": name, "path": str(plugin_path)}
            )

        except Exception as e:
            logger.error(f"[ToolLoader] 生成工具失败: {e}")
            return ToolOutput(success=False, result=None, error=str(e))

    def get_tool(self, name: str) -> Optional[BaseTool]:
        """获取工具实例"""
        return self._tools.get(name)

    def get_langchain_tools(self) -> List:
        """获取 LangChain 格式工具"""
        if not self._langchain_tools:
            self._langchain_tools = [t.to_langchain_tool() for t in self._tools.values()]
        return self._langchain_tools

    def get_function_calls(self) -> List[FunctionCallDefinition]:
        """获取 Function Call 格式定义"""
        return [t.to_function_call() for t in self._tools.values()]

    def get_openai_functions(self) -> List[Dict[str, Any]]:
        """获取 OpenAI Function Call 格式"""
        return [t.to_function_call().to_openai_format() for t in self._tools.values()]

    def get_mcp_schemas(self) -> List[Dict[str, Any]]:
        """获取 MCP 格式定义"""
        return [t.to_mcp_schema() for t in self._tools.values()]

    def get_collection(self) -> ToolCollection:
        """获取工具集合"""
        return ToolCollection(list(self._tools.values()))

    @property
    def registry(self) -> ToolRegistry:
        """获取工具注册表"""
        return self._registry

    @property
    def tools(self) -> Dict[str, BaseTool]:
        """获取所有工具"""
        return self._tools.copy()

    @property
    def config(self) -> Dict[str, Any]:
        """获取当前配置"""
        return self._config.copy()

    @property
    def plugin_registry(self) -> Optional[PluginToolRegistry]:
        """获取插件注册表"""
        return self._plugin_registry

    def _notify_agent_refresh(self):
        """通知全局 Agent 刷新工具列表

        在动态创建工具后调用，使新工具在当前会话中立即可用
        """
        try:
            # 尝试导入并刷新全局 Agent
            # _global_agent 可能是包装后的 Agent (agent/agent.py) 或 CoreAgent (core/agent.py)
            import src.agent.agent as agent_module
            if agent_module._global_agent is not None:
                global_agent = agent_module._global_agent
                # 包装后的 Agent 有 _core_agent 属性
                if hasattr(global_agent, '_core_agent'):
                    global_agent._core_agent.refresh_tools()
                    # 同时设置 _tool_version 为 0，强制下一次调用时重新创建 Agent
                    self._tool_version = 0
                    logger.debug("[ToolLoader] 已通知包装 Agent 刷新工具列表")
                # 直接使用 CoreAgent 的情况
                elif hasattr(global_agent, 'refresh_tools'):
                    global_agent.refresh_tools()
                    self._tool_version = 0
                    logger.debug("[ToolLoader] 已通知 CoreAgent 刷新工具列表")
        except Exception as e:
            logger.warning(f"[ToolLoader] 通知 Agent 刷新失败: {e}")

    @property
    def sources(self) -> List[str]:
        """获取启用的工具来源"""
        if self._plugin_registry:
            return self._plugin_registry._enabled_sources.copy()
        return ["builtin"]


# ============================================================================
# 全局加载器实例
# ============================================================================


_loader: Optional[ToolLoader] = None
_loader_initialized = False  # 标记是否已初始化
_logged_messages: set = set()  # 全局日志去重集合


def get_tool_loader(config_path: Optional[str] = None) -> ToolLoader:
    """获取全局工具加载器实例"""
    global _loader, _loader_initialized, _logged_messages

    if _loader is None:
        _loader = ToolLoader(config_path)
        _loader_initialized = True

    # 将日志去重集合传递给 ToolLoader 实例
    _loader._logged_messages = _logged_messages

    return _loader


def reload_tools() -> None:
    """重新加载工具配置"""
    global _loader
    if _loader:
        _loader.reload()
