# core/chat_service.py - 修复版本
import asyncio
from typing import Dict, Any, List
from langchain_core.messages import HumanMessage
from utils.logger import chat_logger, log_chat_entry
from core.agent_manager import agent_manager


class ChatService:
    """聊天服务 - 支持真正并发的版本"""

    def __init__(self):
        self.agent_manager = agent_manager
        # 创建专用的线程池用于执行同步的Langchain操作
        self._thread_pool = None

    def _get_thread_pool(self):
        """获取或创建线程池"""
        if self._thread_pool is None:
            import concurrent.futures

            # 创建足够大的线程池支持并发
            self._thread_pool = concurrent.futures.ThreadPoolExecutor(
                max_workers=20,  # 根据服务器配置调整
                thread_name_prefix="langchain_worker",
            )
        return self._thread_pool

    async def process_chat_request(
        self, request_data: Dict[str, Any]
    ) -> Dict[str, Any]:
        """异步处理聊天请求 - 真正并发版本"""
        try:
            # 提取请求数据
            message = request_data["message"]
            thread_id = request_data["thread_id"]
            username = request_data["username"]
            backend_url = request_data["backend_url"]
            token = request_data["token"]

            # 生成用户标识符
            # user_id = self.agent_manager._get_user_identifier(username, token)
            user_id = username

            chat_logger.info(
                f"收到请求 - 用户={user_id} , 线程ID={thread_id}, 消息={message[:100]}"
            )

            # 异步获取agent实例
            agent = await self.agent_manager.get_agent_instance(
                thread_id=thread_id,
                username=username,
                backend_url=backend_url,
                token=token,
            )

            # ✅ 修复：在线程池中执行同步的Langchain操作
            result = await self._run_agent_in_threadpool(
                agent, message, thread_id, user_id
            )

            chat_logger.info(f"Agent处理完成 - 用户={user_id}")

            if not isinstance(result, dict) or "messages" not in result:
                raise ValueError(f"Agent返回格式异常: {type(result)}")

            # 处理结果
            return self._process_agent_result(result, user_id, request_data)

        except Exception as e:
            chat_logger.error(f"聊天处理失败: {str(e)}")
            raise

    async def _run_agent_in_threadpool(
        self, agent, message: str, thread_id: str, user_id: str
    ):
        """在线程池中执行Langchain Agent"""
        loop = asyncio.get_event_loop()
        thread_pool = self._get_thread_pool()

        # 准备输入
        inputs = {"messages": [HumanMessage(content=message)]}
        config = {"configurable": {"thread_id": thread_id}}

        chat_logger.info(f"在线程池中执行Agent - 用户={user_id}")

        try:
            # 在线程池中执行同步操作
            result = await loop.run_in_executor(
                thread_pool, lambda: agent.invoke(inputs, config)
            )
            return result
        except Exception as e:
            chat_logger.error(f"Agent执行失败 - 用户={user_id}: {str(e)}")
            raise

    def _process_agent_result(
        self, result: Dict[str, Any], user_id: str, request_data: Dict
    ) -> Dict[str, Any]:
        """处理Agent返回结果"""
        all_messages = result["messages"]
        processed_messages = []
        all_ai_messages = []
        all_tool_calls = []
        final_answer = ""

        for i, msg in enumerate(all_messages):
            msg_data = {
                "index": i,
                "type": getattr(msg, "type", "unknown"),
                "content": "",
            }

            # 获取内容
            if hasattr(msg, "content"):
                content = msg.content
                if isinstance(content, str):
                    msg_data["content"] = content
                else:
                    msg_data["content"] = str(content)

            # 获取工具调用
            if hasattr(msg, "tool_calls") and msg.tool_calls:
                msg_data["tool_calls"] = msg.tool_calls
                all_tool_calls.extend(msg.tool_calls)

                for tool_call in msg.tool_calls:
                    tool_name = tool_call.get("name", "unknown")
                    tool_args = tool_call.get("args", {})
                    chat_logger.info(f"工具调用 - 用户={user_id}, 工具={tool_name}")

            if hasattr(msg, "tool_call_id"):
                msg_data["tool_call_id"] = msg.tool_call_id

            if hasattr(msg, "name"):
                msg_data["name"] = msg.name

            processed_messages.append(msg_data)

            # 收集AI消息
            if msg_data["type"] == "ai":
                all_ai_messages.append(msg_data)
                final_answer = msg_data["content"]

        # 构建响应
        response = {
            "final_answer": final_answer,
            "all_ai_messages": all_ai_messages,
            "all_messages": processed_messages,
            "tool_calls": all_tool_calls,
            "thread_id": request_data["thread_id"],
            "user_identifier": user_id,
            "backend_config": {
                "backend_url": request_data["backend_url"] or "未配置",
                "username": request_data["username"],
                "has_token": bool(request_data["token"]),
            },
            "success": True,
        }

        # 记录日志
        log_chat_entry(user_id, request_data["message"], response)
        chat_logger.info(f"请求处理完成 - 用户={user_id}")

        return response

    async def shutdown(self):
        """关闭线程池"""
        if self._thread_pool:
            self._thread_pool.shutdown(wait=False)
            self._thread_pool = None
            chat_logger.info("聊天服务线程池已关闭")


# 全局实例
chat_service = ChatService()