from langchain_mcp_adapters.client import MultiServerMCPClient
from langchain_mcp_adapters.tools import load_mcp_tools
from langchain.agents import create_agent
from ai_hub_agents.config import settings,Config
import json
from typing import List
from pathlib import Path
from langchain_core.messages import BaseMessage,messages_from_dict,messages_to_dict,HumanMessage,AIMessage
from contextlib import AsyncExitStack
from ai_hub_agents.callback import LoadMCPTools
from pydantic import BaseModel, Field
from langchain_core.tools import BaseTool
from langchain_openai import ChatOpenAI
from .json_store import JsonStore
import logging

logger = logging.getLogger(__name__)

class Agent(BaseModel):
    config: Config = Field(default_factory=Config)
    """配置"""
    thread_id: str
    """线程 ID"""
    mcp_client: MultiServerMCPClient = None
    """MCP 客户端"""
    mcp_server_names: List[str] = None
    """MCP 服务名称列表"""
    prompt: str = None
    """提示词"""
    json_store: JsonStore = None
    """JSON 存储"""

    def model_post_init(self,ctx):
        """初始化"""
        self._load_mcp_client()
        self._load_prompt()
        self._load_json_store()

    def _load_mcp_client(self):
        """加载 MCP 客户端"""
        if Path(self.config.mcp_path).exists():
            with open(self.config.mcp_path, "r") as f:
                mcp_metadata = json.load(f)
            mcp_servers: dict = mcp_metadata.get("mcpServers", {})

            # 默认值
            mcp_servers = {k: v|{
                "transport": "stdio",
            } for k,v in mcp_servers.items()}
        else:
            mcp_servers = {}

        self.mcp_client = MultiServerMCPClient(mcp_servers)
        self.mcp_server_names = mcp_servers.keys()

    def _load_prompt(self):
        """加载提示词"""
        if Path(self.config.prompt_path).exists():
            with open(self.config.prompt_path, "r") as f:
                prompt = f.read()
        else:
            prompt = ""
        self.prompt = prompt

    def _load_json_store(self):
        """加载 JSON 存储"""
        self.json_store = JsonStore(Path(self.config.data_dir) / self.thread_id / self.config.memory_file_name)

    def _load_memory(self) -> List[BaseMessage]:
        """加载记忆"""
        metadata: dict = self.json_store.read()
        try:
            messages = messages_from_dict(metadata)
        except Exception as e:
            logger.error(f"加载记忆失败: {e}")
            return []
        return messages

    def _save_memory(self,messages:List[BaseMessage]):
        """保存记忆"""
        metadata = messages_to_dict(messages)
        self.json_store.write(metadata)

    def _append_memory(self,message:BaseMessage):
        """追加记忆"""
        self._save_memory(self._load_memory()+[message])

    async def run(self,query:str,name:str="user") -> str:
        """运行一轮对话
        Args:
            query: 用户问题
            name: 用户名称
        Returns:
            对话结果
        """

        async with AsyncExitStack() as stack:
            # 同时持有所有服务器的 session
            sessions = {}
            for name in self.mcp_server_names:
                session = await stack.enter_async_context(self.mcp_client.session(name))
                sessions[name] = session

            # 合并所有服务器的 tools，用 server_name 做前缀避免重名
            all_tools: List[BaseTool] = []
            for name, session in sessions.items():
                tools = await load_mcp_tools(
                    session,
                    server_name=name,
                    tool_name_prefix=True  # 如 "mcp-tool-web-search_search"
                )
                all_tools.extend(tools)

            # 加载工具事件
            tool_names = [tool.name for tool in all_tools]
            LoadMCPTools(tool_names=tool_names)

            # 加载 LLM
            llm = ChatOpenAI(
                model=settings.llm_model,
                api_key=settings.llm_api_key,
                base_url=settings.llm_base_url,
            )

            # 记录用户消息
            self._append_memory(HumanMessage(content=query,name=name))

            agent = create_agent(
                model=llm,
                tools=all_tools,
                system_prompt=self.prompt,
            )

            response = await agent.ainvoke(self._load_memory())

            # AI 回答
            self._append_memory(response)

            return response.content