"""
混合模式客服系统 - 支持业务知识+通用知识+实时工具
"""
from typing import TypedDict, Literal
from langchain_core.messages import 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 HybridState(TypedDict):
    """混合模式状态"""
    messages: list
    intent: Literal["medical", "general", "realtime", "greeting"]
    context: str
    tool_result: str
    session_id: str


class HybridChatbot:
    """混合客服：垂直领域 + 通用知识 + 实时工具"""
    
    def __init__(self):
        self.llm = ChatOpenAI(
            model=LLM_MODEL,
            api_key=LLM_API_KEY,
            base_url=LLM_BASE_URL,
            temperature=0.7
        )
        self.knowledge_base = get_knowledge_base()
        self.workflow = self._build_workflow()
    
    def _classify_intent(self, state: HybridState) -> HybridState:
        """意图识别 - 判断问题类型"""
        query = state["messages"][-1].content
        
        # 用LLM快速判断意图（也可以用关键词规则）
        prompt = f"""判断用户意图，只能返回以下之一：
- medical: 医疗/健康相关问题
- general: 通用知识（历史、科学、生活常识等）
- realtime: 需要实时数据（天气、股价、新闻等）
- greeting: 问候/闲聊

用户问题：{query}
意图："""
        
        response = self.llm.invoke([HumanMessage(content=prompt)])
        intent = response.content.strip().lower()
        
        # 归一化
        valid_intents = ["medical", "general", "realtime", "greeting"]
        intent = intent if intent in valid_intents else "general"
        
        print(f"🎯 意图识别: {query} -> {intent}")
        return {"intent": intent, "messages": state["messages"], "session_id": state["session_id"]}
    
    def _route_by_intent(self, state: HybridState) -> str:
        """根据意图路由到不同节点"""
        intent = state.get("intent", "general")
        
        routing_map = {
            "medical": "retrieve_knowledge",  # 查医疗知识库
            "general": "general_answer",       # 直接问大模型
            "realtime": "call_tool",           # 调用实时工具
            "greeting": "general_answer"       # 闲聊也走通用
        }
        
        return routing_map.get(intent, "general_answer")
    
    def _retrieve_medical(self, state: HybridState) -> HybridState:
        """检索医疗知识库"""
        query = state["messages"][-1].content
        docs = self.knowledge_base.search(query)
        
        if docs:
            context = "\n".join([doc.page_content for doc in docs])
            return {"context": context, "messages": state["messages"], 
                    "intent": state["intent"], "session_id": state["session_id"]}
        else:
            # 医疗问题但知识库没有 → 转人工
            return {"context": "【无相关医疗知识，需转人工】", 
                    "messages": state["messages"],
                    "intent": state["intent"], "session_id": state["session_id"]}
    
    def _call_realtime_tool(self, state: HybridState) -> HybridState:
        """调用实时工具（MCP/Function Call）"""
        query = state["messages"][-1].content
        
        # 示例：天气查询（实际可用 MCP/天气API）
        if "天气" in query:
            # 这里可以调用 MCP 或天气API
            # result = call_mcp("weather", location="北京")
            tool_result = "【实时天气工具】：北京今天晴，25°C，适合出行"
        elif "股价" in query or "股票" in query:
            tool_result = "【实时股价工具】：需要调用股票API获取实时数据"
        else:
            tool_result = "【无可用工具】"
        
        return {"tool_result": tool_result, "messages": state["messages"],
                "intent": state["intent"], "session_id": state["session_id"]}
    
    def _generate_answer(self, state: HybridState) -> HybridState:
        """生成最终回答"""
        intent = state["intent"]
        query = state["messages"][-1].content
        
        if intent == "medical":
            # 医疗问题：基于知识库回答
            if "无相关医疗知识" in state.get("context", ""):
                response = "抱歉，这个问题我暂时无法回答，为您转接专业医生..."
            else:
                system_msg = f"你是医疗助手，基于以下知识回答：\n{state['context']}"
                messages = [SystemMessage(content=system_msg)] + state["messages"]
                response = self.llm.invoke(messages).content
        
        elif intent == "realtime":
            # 实时信息：基于工具结果回答
            system_msg = f"基于以下实时信息回答用户：\n{state.get('tool_result', '')}"
            messages = [SystemMessage(content=system_msg)] + state["messages"]
            response = self.llm.invoke(messages).content
        
        else:  # general / greeting
            # 通用问题：直接回答
            system_msg = "你是友好的助手，回答用户的通用问题。"
            messages = [SystemMessage(content=system_msg)] + state["messages"]
            response = self.llm.invoke(messages).content
        
        return {"messages": state["messages"] + [AIMessage(content=response)],
                "intent": intent, "session_id": state["session_id"]}
    
    def _build_workflow(self):
        """构建工作流"""
        workflow = StateGraph(HybridState)
        
        # 添加节点
        workflow.add_node("classify", self._classify_intent)
        workflow.add_node("retrieve_knowledge", self._retrieve_medical)
        workflow.add_node("call_tool", self._call_realtime_tool)
        workflow.add_node("general_answer", self._generate_answer)
        workflow.add_node("generate", self._generate_answer)
        
        # 条件路由
        workflow.set_entry_point("classify")
        workflow.add_conditional_edges(
            "classify",
            self._route_by_intent,
            {
                "retrieve_knowledge": "retrieve_knowledge",
                "call_tool": "call_tool",
                "general_answer": "general_answer"
            }
        )
        
        # 汇聚到生成节点
        workflow.add_edge("retrieve_knowledge", "generate")
        workflow.add_edge("call_tool", "generate")
        workflow.add_edge("generate", END)
        workflow.add_edge("general_answer", END)
        
        return workflow.compile(checkpointer=MemorySaver())
    
    async def chat(self, message: str, session_id: str) -> str:
        """处理消息"""
        config = {"configurable": {"thread_id": session_id}}
        result = await self.workflow.ainvoke(
            {"messages": [HumanMessage(content=message)],
             "intent": "general",
             "context": "",
             "tool_result": "",
             "session_id": session_id},
            config
        )
        return result["messages"][-1].content


# 单例
_hybrid_chatbot = None

def get_hybrid_chatbot():
    global _hybrid_chatbot
    if _hybrid_chatbot is None:
        _hybrid_chatbot = HybridChatbot()
    return _hybrid_chatbot