"""
AI 客服核心逻辑 - 基于 LangGraph
"""
from typing import TypedDict, Annotated, Sequence
import operator

from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver

from config import LLM_API_KEY, LLM_BASE_URL, LLM_MODEL, SYSTEM_PROMPT
from knowledge import get_knowledge_base


# 定义状态类型
class ChatState(TypedDict):
    """对话状态"""
    messages: Annotated[Sequence[BaseMessage], operator.add]  # 消息历史
    context: str  # 检索到的知识上下文
    session_id: str  # 会话ID


class AIChatbot:
    """AI 客服机器人"""
    
    def __init__(self):
        # 初始化 LLM
        self.llm = ChatOpenAI(
            model=LLM_MODEL,
            api_key=LLM_API_KEY,
            base_url=LLM_BASE_URL,
            temperature=0.7,
            max_tokens=1000
        )
        
        # 获取知识库
        self.knowledge_base = get_knowledge_base()
        
        # 构建工作流
        self.workflow = self._build_workflow()
    
    def _retrieve_knowledge(self, state: ChatState) -> ChatState:
        """检索相关知识"""
        # 获取最后一条用户消息
        last_message = state["messages"][-1]
        if not isinstance(last_message, HumanMessage):
            return state
        
        query = last_message.content
        
        # 从知识库检索
        docs = self.knowledge_base.search(query)
        
        # 构建上下文
        if docs:
            context = "\n\n".join([f"[知识库内容 {i+1}]\n{doc.page_content}" 
                                   for i, doc in enumerate(docs)])
        else:
            context = "【知识库中没有找到相关信息】"
        
        return {"context": context, "messages": state["messages"], "session_id": state["session_id"]}
    
    def _generate_response(self, state: ChatState) -> ChatState:
        """生成回答"""
        # 构建系统提示
        system_prompt = f"""{SYSTEM_PROMPT}

【知识库信息】
{state['context']}

请基于以上知识库信息回答用户问题。如果知识库中没有相关信息，请诚实地告诉用户你无法回答。
"""
        
        # 构建消息列表
        messages = [SystemMessage(content=system_prompt)] + list(state["messages"])
        
        # 调用 LLM
        response = self.llm.invoke(messages)
        
        # 返回新的消息（会自动添加到历史）
        return {"messages": [response], "context": state["context"], "session_id": state["session_id"]}
    
    def _build_workflow(self) -> StateGraph:
        """构建 LangGraph 工作流"""
        # 创建工作流
        workflow = StateGraph(ChatState)
        
        # 添加节点
        workflow.add_node("retrieve", self._retrieve_knowledge)
        workflow.add_node("generate", self._generate_response)
        
        # 添加边
        workflow.set_entry_point("retrieve")
        workflow.add_edge("retrieve", "generate")
        workflow.add_edge("generate", END)
        
        # 编译（添加内存检查点）
        memory = MemorySaver()
        return workflow.compile(checkpointer=memory)
    
    async def chat(self, message: str, session_id: str) -> str:
        """
        处理用户消息
        
        Args:
            message: 用户消息
            session_id: 会话ID（用于区分不同用户）
        
        Returns:
            AI 回复
        """
        # 构建配置
        config = {"configurable": {"thread_id": session_id}}
        
        # 运行工作流
        result = await self.workflow.ainvoke(
            {
                "messages": [HumanMessage(content=message)],
                "context": "",
                "session_id": session_id
            },
            config
        )
        
        # 返回最后一条 AI 消息
        last_message = result["messages"][-1]
        if isinstance(last_message, AIMessage):
            return last_message.content
        return str(last_message)
    
    def get_history(self, session_id: str) -> list:
        """获取对话历史"""
        config = {"configurable": {"thread_id": session_id}}
        # 从 checkpointer 获取状态
        state = self.workflow.get_state(config)
        if state and "messages" in state.values:
            return [
                {"role": "user" if isinstance(m, HumanMessage) else "assistant",
                 "content": m.content}
                for m in state.values["messages"]
            ]
        return []


# 全局客服实例
_chatbot = None

def get_chatbot() -> AIChatbot:
    """获取客服单例"""
    global _chatbot
    if _chatbot is None:
        _chatbot = AIChatbot()
    return _chatbot