"""XBot Agent - 封装 nanobot 核心能力的 Agent"""

import os
from pathlib import Path
from typing import Any, Optional
from datetime import datetime

from .loop import AgentLoop
from .session import SessionManager
from .adapter import XBotLLMAdapter, LLMResponse
from . import config

# 尝试导入 simplemem
try:
    from simplemem import SimpleMemSystem
    HAS_SIMPLEMEM = True
except ImportError:
    HAS_SIMPLEMEM = False


class SimpleToolRegistry:
    """简单的工具注册表"""

    def __init__(self):
        self._tools: dict[str, Any] = {}

    def register(self, name: str, func: Any, description: str = "") -> None:
        """注册一个工具"""
        self._tools[name] = {
            "function": func,
            "description": description,
        }

    def get_definitions(self) -> list[dict]:
        """获取工具定义列表"""
        tools = []
        for name, tool in self._tools.items():
            tools.append({
                "type": "function",
                "function": {
                    "name": name,
                    "description": tool.get("description", ""),
                    "parameters": {
                        "type": "object",
                        "properties": {},
                        "required": [],
                    }
                }
            })
        return tools

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

    async def execute(self, name: str, arguments: dict) -> Any:
        """执行工具"""
        tool = self._tools.get(name)
        if not tool:
            return f"Tool {name} not found"

        func = tool.get("function")
        if not func:
            return f"Tool {name} has no function"

        try:
            if callable(func):
                return await func(**arguments) if hasattr(func, '__await__') else func(**arguments)
            return "Tool function is not callable"
        except Exception as e:
            return f"Tool execution error: {str(e)}"


class XBotAgent:
    """
    XBot Agent - 基于 nanobot 核心的 Agent 实现

    特性：
    - 多轮 tool-calling 对话
    - 自动内存压缩
    - 会话历史持久化
    """

    def __init__(
        self,
        name: str,
        role_description: str,
        provider: str = "openai",
        model: str = "gpt-4",
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        workspace: Optional[Path] = None,
        context_window_tokens: int = 200000,
        embedding_model: Optional[str] = None,
        embedding_base_url: Optional[str] = None,
    ):
        """
        初始化 XBot Agent

        Args:
            name: Agent 名称
            role_description: Agent 角色描述
            provider: LLM 提供商
            model: 模型名称
            api_key: API Key
            base_url: Base URL
            workspace: 工作目录（用于存储会话和记忆）
            context_window_tokens: 上下文窗口大小
        """
        self.name = name
        self.role_description = role_description

        # 使用配置文件的默认值
        if api_key is None:
            api_key = config.API_KEY
        if base_url is None:
            base_url = config.BASE_URL
        if workspace is None:
            workspace = Path(config.WORKSPACE)

        # 创建工作目录
        self.workspace = workspace
        self.workspace.mkdir(parents=True, exist_ok=True)

        # 创建 LLM 适配器
        self.provider = XBotLLMAdapter(
            provider=provider,
            model_name=model,
            api_key=api_key,
            base_url=base_url,
        )

        # 创建工具注册表
        self.tools = SimpleToolRegistry()
        self._register_default_tools()

        # 创建 Agent Loop
        self.agent_loop = AgentLoop(
            provider=self.provider,
            model=model,
            tools=self.tools,
            max_iterations=50,
        )

        # 创建会话管理器
        self.sessions = SessionManager(self.workspace)

        # 创建 SimpleMem 记忆系统
        if HAS_SIMPLEMEM and api_key and config.ENABLE_SIMPLEMEM:
            # 使用配置文件的 embedding 设置
            emb_model = embedding_model or config.EMBEDDING_MODEL
            emb_base = embedding_base_url or config.EMBEDDING_BASE_URL or base_url

            self.memory = SimpleMemSystem(
                api_key=api_key,
                base_url=emb_base,
                model=model,
                embedding_model=emb_model,
                db_path=str(self.workspace / "memory_db"),
                clear_db=False,
                # 并行处理配置
                enable_parallel_processing=config.ENABLE_PARALLEL_PROCESSING,
                max_parallel_workers=config.MAX_PARALLEL_WORKERS,
                enable_parallel_retrieval=config.ENABLE_PARALLEL_RETRIEVAL,
                max_retrieval_workers=config.MAX_RETRIEVAL_WORKERS,
                enable_planning=config.ENABLE_PLANNING,
                enable_reflection=config.ENABLE_REFLECTION,
                max_reflection_rounds=config.MAX_REFLECTION_ROUNDS,
            )
            self._use_simplemem = True
            print(f"SimpleMem initialized with embedding: {emb_model}, base_url: {emb_base}")
        else:
            self.memory = None
            self._use_simplemem = False
            if not api_key:
                print("Warning: No API key provided, SimpleMem will be disabled")

    def _register_default_tools(self) -> None:
        """注册默认工具"""
        # 可以在这里添加默认工具
        pass

    def register_tool(
        self,
        name: str,
        func: Any,
        description: str = "",
        parameters: Optional[dict] = None,
    ) -> None:
        """注册自定义工具"""
        tool_def = {
            "type": "function",
            "function": {
                "name": name,
                "description": description,
                "parameters": parameters or {
                    "type": "object",
                    "properties": {},
                    "required": [],
                }
            }
        }
        # 存储在 tools 中
        self.tools.register(name, func, description)

    async def run(
        self,
        user_input: str,
        session_id: str = "default",
    ) -> dict[str, Any]:
        """
        运行 Agent 对话

        Args:
            user_input: 用户输入
            session_id: 会话 ID

        Returns:
            dict: 包含 content, tool_calls 等
        """
        # 获取或创建会话
        session = self.sessions.get_or_create(session_id)

        # 构建系统提示
        system_prompt = f"""你是 {self.name}。
{self.role_description}

请根据用户的问题回答，并使用 Markdown 格式输出。"""

        # 如果使用 SimpleMem，检索相关记忆
        memory_context = ""
        if self._use_simplemem and self.memory:
            try:
                memory_context = self.memory.ask(user_input)
            except Exception as e:
                print(f"Memory retrieval error: {e}")

        if memory_context:
            system_prompt += f"\n\n相关记忆:\n{memory_context}"

        # 获取历史消息
        history = session.get_history(max_messages=50)

        # 构建初始消息
        initial_messages = history + [
            {"role": "user", "content": user_input}
        ]

        # 运行 agent loop
        final_content, tools_used, all_messages = await self.agent_loop.run_loop(
            initial_messages=initial_messages,
            system_prompt=system_prompt,
        )

        # 保存到会话
        for m in all_messages[len(history):]:
            session.messages.append(m)
        self.sessions.save(session)

        # 保存到 SimpleMem 记忆
        if self._use_simplemem and self.memory and final_content:
            try:
                self.memory.add_dialogue("User", user_input, datetime.now().isoformat())
                self.memory.add_dialogue(self.name, final_content, datetime.now().isoformat())
                self.memory.finalize()
            except Exception as e:
                print(f"Memory save error: {e}")

        return {
            "content": final_content or "No response",
            "tool_calls": tools_used,
            "session_id": session_id,
        }

    async def run_stream(
        self,
        user_input: str,
        session_id: str = "default",
    ):
        """
        运行 Agent 对话（流式输出）

        先完整执行 agent loop，最后流式输出结果

        Args:
            user_input: 用户输入
            session_id: 会话 ID

        Yields:
            str: 流式回复片段
        """
        # 先完整执行 agent loop（包含 tool-calling）
        result = await self.run(user_input, session_id)
        content = result["content"]

        # 流式输出结果
        for char in content:
            yield char

    def clear_session(self, session_id: str) -> None:
        """清除会话"""
        session = self.sessions.get_or_create(session_id)
        session.clear()
        self.sessions.save(session)
        self.sessions.invalidate(session_id)

    def list_sessions(self) -> list[dict]:
        """列出所有会话"""
        return self.sessions.list_sessions()
