#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
增强版RAG Agent - 智能问答系统
新增功能：
1. 多轮对话记忆
2. 智能问题分解
3. 上下文关联分析
4. 专业术语解释
5. 相关问题推荐
"""

from __future__ import annotations
import os
import json
import asyncio
from typing import Literal, List, Dict, Any, Optional
from dotenv import load_dotenv 
load_dotenv(override=True)

from langchain_deepseek import ChatDeepSeek
from langgraph.graph import MessagesState, StateGraph, START, END
from pydantic import BaseModel, Field
from langchain_core.tools import BaseTool
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage
import re
from datetime import datetime
from collections import defaultdict
import math

# 初始化模型
model = ChatDeepSeek(model="deepseek-chat")
grader_model = ChatDeepSeek(model="deepseek-chat")

class EnhancedTextRetriever(BaseTool):
    """增强版本地文本检索工具"""
    name: str = "retrieve_telco_handbook"
    description: str = "Search and return relevant sections from the Telco Customer Churn Analytics Handbook with enhanced analysis."
    conversation_history: List[str] = []
    context_cache: Dict[str, Any] = {}
    
    def __init__(self):
        """初始化检索工具"""
        # 首先调用父类的初始化方法
        super().__init__()
        # 初始化属性
        self.conversation_history = []
        self.context_cache = {}
    
    def _run(self, query: str, run_manager: CallbackManagerForToolRun = None) -> str:
        """检索相关文档内容"""
        try:
            # 保存对话历史
            self.conversation_history.append(query)
            
            # 加载本地文本数据库
            VS_PATH = "telco_customer_churn_analytics_handbook"
            db_file = os.path.join(VS_PATH, "text_database.json")
            
            if not os.path.exists(db_file):
                return "本地文本数据库未找到，请先运行 create_local_vector_db.py 创建数据库。"
            
            with open(db_file, 'r', encoding='utf-8') as f:
                sections = json.load(f)
            
            # 增强查询处理
            enhanced_query = self._enhance_query(query)
            
            # 执行搜索
            results = self._advanced_search(sections, enhanced_query, top_k=5)
            
            if not results:
                # 尝试模糊匹配
                fuzzy_results = self._fuzzy_search(sections, query)
                if fuzzy_results:
                    return self._format_fuzzy_results(fuzzy_results, query)
                return f"未找到与'{query}'直接相关的信息。您可以尝试使用更具体的关键词，或询问以下相关问题：\n" + self._suggest_questions()
            
            # 格式化结果
            formatted_result = self._format_enhanced_results(results, query)
            
            # 缓存上下文
            self.context_cache[query] = results
            
            return formatted_result
            
        except Exception as e:
            return f"检索时发生错误：{str(e)}"
    
    def _enhance_query(self, query: str) -> Dict[str, Any]:
        """增强查询处理"""
        enhanced = {
            'original': query,
            'terms': self._extract_key_terms(query),
            'intent': self._analyze_intent(query),
            'context': self._get_conversation_context(),
            'synonyms': self._get_synonyms(query)
        }
        return enhanced
    
    def _extract_key_terms(self, query: str) -> List[str]:
        """提取关键术语"""
        # 专业术语词典
        technical_terms = {
            '流失': ['churn', '客户流失', '用户流失', '离网', '退订'],
            '预测': ['prediction', '建模', '模型', '预估', '预判'],
            '特征': ['feature', '变量', '字段', '属性', '指标'],
            '数据': ['data', '信息', '记录', '样本', '数据集'],
            '分析': ['analysis', '统计', '分析', '研究', '挖掘'],
            '合同': ['contract', '合约', '协议', '签约'],
            '账单': ['charges', '费用', '账单', '缴费', '付费'],
            '服务': ['service', '业务', '功能', '产品'],
            '客户': ['customer', '用户', '客人', '消费者'],
            '电信': ['telco', '通信', '移动', '联通', '电信'],
            '模型': ['model', '算法', '机器学习', 'ML', 'AI'],
            '可视化': ['visualization', '图表', '图形', '展示'],
            '预处理': ['preprocessing', '清洗', '处理', '准备']
        }
        
        terms = []
        query_lower = query.lower()
        
        for key, synonyms in technical_terms.items():
            if key in query or any(syn in query_lower for syn in synonyms):
                terms.extend([key] + synonyms)
        
        # 添加原始查询词
        terms.extend(query.split())
        
        return list(set(terms))
    
    def _analyze_intent(self, query: str) -> str:
        """分析查询意图"""
        intent_patterns = {
            'definition': ['什么是', '定义', '含义', '解释', '概念'],
            'method': ['如何', '怎么', '方法', '步骤', '流程'],
            'comparison': ['比较', '区别', '差异', '对比', '异同'],
            'analysis': ['分析', '评估', '研究', '探讨', '调查'],
            'example': ['例子', '示例', '案例', '举例', '实例'],
            'reason': ['为什么', '原因', '因素', '影响', '导致'],
            'process': ['过程', '流程', '步骤', '阶段', '环节']
        }
        
        query_lower = query.lower()
        for intent, patterns in intent_patterns.items():
            if any(pattern in query_lower for pattern in patterns):
                return intent
        return 'general'
    
    def _get_conversation_context(self) -> List[str]:
        """获取对话上下文"""
        return self.conversation_history[-3:] if len(self.conversation_history) > 1 else []
    
    def _get_synonyms(self, query: str) -> List[str]:
        """获取同义词"""
        synonym_dict = {
            '流失': ['离网', '退订', '取消', '终止'],
            '客户': ['用户', '消费者', '客人'],
            '预测': ['预估', '预判', '预警', '预报'],
            '分析': ['研究', '调查', '探讨', '评估'],
            '数据': ['信息', '资料', '记录'],
            '模型': ['算法', '方法', '技术'],
            '特征': ['变量', '字段', '属性', '指标']
        }
        
        synonyms = []
        for word, syns in synonym_dict.items():
            if word in query:
                synonyms.extend(syns)
        
        return synonyms
    
    def _advanced_search(self, sections: List[Dict], enhanced_query: Dict, top_k: int = 5) -> List[Dict]:
        """高级搜索功能"""
        results = []
        
        for section in sections:
            score = 0
            content_lower = section['content'].lower()
            title_lower = section['title'].lower()
            
            # 原始查询匹配
            original_terms = enhanced_query['original'].lower().split()
            for term in original_terms:
                if term in title_lower:
                    score += 15
                score += content_lower.count(term) * 2
            
            # 关键术语匹配
            for term in enhanced_query['terms']:
                term_lower = term.lower()
                if term_lower in title_lower:
                    score += 10
                score += content_lower.count(term_lower) * 1.5
            
            # 同义词匹配
            for synonym in enhanced_query['synonyms']:
                synonym_lower = synonym.lower()
                if synonym_lower in content_lower:
                    score += 3
            
            # 上下文相关性
            for context_query in enhanced_query['context']:
                context_terms = context_query.lower().split()
                for term in context_terms:
                    if term in content_lower:
                        score += 1
            
            # 意图匹配加权
            intent = enhanced_query['intent']
            if intent == 'definition' and ('定义' in content_lower or '是指' in content_lower):
                score += 5
            elif intent == 'method' and ('方法' in content_lower or '步骤' in content_lower):
                score += 5
            elif intent == 'analysis' and ('分析' in content_lower or '研究' in content_lower):
                score += 5
            
            if score > 0:
                results.append({
                    'section': section,
                    'score': score,
                    'title': section['title'],
                    'content': section['content'],
                    'relevance': self._calculate_relevance(section, enhanced_query)
                })
        
        # 按得分排序
        results.sort(key=lambda x: x['score'], reverse=True)
        return results[:top_k]
    
    def _calculate_relevance(self, section: Dict, enhanced_query: Dict) -> float:
        """计算相关性得分"""
        content = section['content'].lower()
        
        # 基础相关性
        base_score = 0
        for term in enhanced_query['terms']:
            if term.lower() in content:
                base_score += 1
        
        # 标题相关性
        title_score = 0
        for term in enhanced_query['original'].split():
            if term.lower() in section['title'].lower():
                title_score += 2
        
        return (base_score + title_score) / len(enhanced_query['terms'])
    
    def _fuzzy_search(self, sections: List[Dict], query: str) -> List[Dict]:
        """模糊搜索"""
        fuzzy_results = []
        query_lower = query.lower()
        
        for section in sections:
            content_lower = section['content'].lower()
            
            # 计算简单的相似度
            common_words = 0
            total_words = len(query.split())
            
            for word in query.split():
                if word.lower() in content_lower:
                    common_words += 1
            
            if common_words > 0:
                similarity = common_words / total_words
                if similarity > 0.3:  # 至少30%的词匹配
                    fuzzy_results.append({
                        'section': section,
                        'similarity': similarity,
                        'title': section['title'],
                        'content': section['content']
                    })
        
        return sorted(fuzzy_results, key=lambda x: x['similarity'], reverse=True)[:3]
    
    def _format_enhanced_results(self, results: List[Dict], query: str) -> str:
        """格式化增强搜索结果"""
        if not results:
            return "未找到相关信息。"
        
        formatted_results = []
        formatted_results.append(f"🔍 针对查询「{query}」找到以下相关信息：\n")
        
        for i, result in enumerate(results, 1):
            relevance_score = result.get('relevance', 0)
            confidence = "高" if relevance_score > 0.7 else "中" if relevance_score > 0.4 else "低"
            
            formatted_results.append(f"📋 【相关内容 {i}】（相关度：{confidence}）")
            formatted_results.append(f"📖 {result['title']}")
            formatted_results.append(f"📝 {result['content'][:400]}...")
            
            if i < len(results):
                formatted_results.append("")
        
        # 添加相关问题推荐
        formatted_results.append(f"\n💡 您可能还想了解：")
        suggestions = self._generate_suggestions(query, results)
        for suggestion in suggestions:
            formatted_results.append(f"   • {suggestion}")
        
        return "\n".join(formatted_results)
    
    def _format_fuzzy_results(self, results: List[Dict], query: str) -> str:
        """格式化模糊搜索结果"""
        formatted_results = []
        formatted_results.append(f"🔍 未找到「{query}」的直接匹配，但找到以下相关信息：\n")
        
        for i, result in enumerate(results, 1):
            similarity = result.get('similarity', 0)
            confidence = f"{similarity*100:.0f}%"
            
            formatted_results.append(f"📋 【可能相关 {i}】（匹配度：{confidence}）")
            formatted_results.append(f"📖 {result['title']}")
            formatted_results.append(f"📝 {result['content'][:300]}...")
            formatted_results.append("")
        
        return "\n".join(formatted_results)
    
    def _suggest_questions(self) -> str:
        """生成建议问题"""
        suggestions = [
            "电信客户流失的主要原因是什么？",
            "如何进行客户流失预测建模？",
            "数据预处理有哪些重要步骤？",
            "特征工程在流失预测中的作用？",
            "如何评估客户流失预测模型的效果？"
        ]
        return "\n".join([f"   • {q}" for q in suggestions])
    
    def _generate_suggestions(self, query: str, results: List[Dict]) -> List[str]:
        """基于查询和结果生成相关问题建议"""
        suggestions = []
        
        # 基于查询意图生成建议
        if "流失" in query:
            suggestions.extend([
                "客户流失预测模型的准确率如何提升？",
                "不同行业的客户流失特征有何差异？"
            ])
        
        if "数据" in query or "预处理" in query:
            suggestions.extend([
                "数据清洗的最佳实践有哪些？",
                "如何处理缺失值和异常值？"
            ])
        
        if "特征" in query or "变量" in query:
            suggestions.extend([
                "如何进行特征选择和特征工程？",
                "特征重要性如何评估？"
            ])
        
        # 基于结果内容生成建议
        for result in results[:2]:  # 只看前两个结果
            content = result['content'].lower()
            if "建模" in content and "建模相关问题" not in [s for s in suggestions]:
                suggestions.append("建模过程中需要注意哪些关键点？")
            if "可视化" in content and "可视化相关问题" not in [s for s in suggestions]:
                suggestions.append("数据可视化在分析中有什么作用？")
        
        return suggestions[:3]  # 最多返回3个建议

# 创建增强版检索工具实例    
enhanced_retriever_tool = EnhancedTextRetriever()

# 增强版Prompt
ENHANCED_SYSTEM_INSTRUCTION = (
    "你是一个专业的电信数据分析专家助手，专精于客户流失分析领域。你具备以下能力：\n\n"
    "🎯 **核心专长**:\n"
    "- 深度理解电信行业客户流失分析\n"
    "- 熟悉数据科学和机器学习方法\n"
    "- 精通特征工程和模型优化\n"
    "- 擅长业务洞察和策略建议\n\n"
    "💡 **回答特色**:\n"
    "- 提供结构化、层次分明的专业解答\n"
    "- 结合理论知识和实践经验\n"
    "- 主动提供相关问题建议\n"
    "- 解释专业术语和概念\n\n"
    "🔍 **工作方式**:\n"
    "- 仔细分析用户问题的意图和层次\n"
    "- 从手册中检索最相关的信息\n"
    "- 提供全面而精准的专业解答\n"
    "- 主动推荐相关的深入问题\n\n"
    "如果问题不在电信客户流失分析范围内，请礼貌说明并引导到相关主题。\n"
    "当需要更多信息时，请调用检索工具 `retrieve_telco_handbook`。"
)

ENHANCED_GRADE_PROMPT = (
    "你是一个专业的信息相关性评估专家。请评估检索到的文档内容与用户问题的相关性。\n\n"
    "评估标准：\n"
    "- 内容是否直接回答了用户的问题\n"
    "- 信息的准确性和完整性\n"
    "- 是否包含用户需要的关键信息\n\n"
    "检索到的文档：\n{context}\n\n"
    "用户问题：{question}\n\n"
    "如果内容高度相关且有价值，返回 'yes'；否则返回 'no'。"
)

ENHANCED_REWRITE_PROMPT = (
    "你是一个查询优化专家。请将用户的问题重写得更加精确和易于检索。\n\n"
    "优化原则：\n"
    "- 使用专业的电信和数据分析术语\n"
    "- 明确查询的核心意图\n"
    "- 增加相关的同义词和关键词\n"
    "- 保持问题的原始含义\n\n"
    "原始问题：\n{question}\n\n"
    "请提供优化后的查询版本："
)

ENHANCED_ANSWER_PROMPT = (
    "你是一位资深的电信数据分析专家。请基于检索到的上下文信息，为用户提供专业、详细的解答。\n\n"
    "回答要求：\n"
    "📋 **结构化回答**：使用清晰的标题和分点说明\n"
    "📊 **专业深度**：包含理论背景、实践方法和业务洞察\n"
    "💡 **实用建议**：提供可操作的建议和最佳实践\n"
    "🔍 **举例说明**：使用具体案例或数据支持观点\n"
    "🎯 **相关扩展**：主动关联相关主题和深入问题\n\n"
    "如果上下文信息不足，请诚实说明并建议用户提供更多信息或调整问题。\n\n"
    "用户问题：\n{question}\n\n"
    "检索上下文：\n{context}\n\n"
    "请提供专业的解答："
)

# 节点函数
async def enhanced_generate_query_or_respond(state: MessagesState):
    """LLM决定直接回答还是调用检索工具"""
    response = await model.bind_tools([enhanced_retriever_tool]).ainvoke(
        [
            {"role": "system", "content": ENHANCED_SYSTEM_INSTRUCTION},
            *state["messages"],
        ]
    )
    return {"messages": [response]}

class EnhancedGradeDoc(BaseModel):
    binary_score: str = Field(description="Relevance score 'yes' or 'no'.")
    confidence: float = Field(description="Confidence score between 0.0 and 1.0")
    reasoning: str = Field(description="Brief explanation of the relevance assessment")

async def enhanced_grade_documents(state: MessagesState) -> Literal["generate_answer", "rewrite_question"]:
    """增强版文档相关性评估"""
    question = state["messages"][0].content  # 原始用户问题
    ctx = state["messages"][-1].content      # 检索器输出
    prompt = ENHANCED_GRADE_PROMPT.format(question=question, context=ctx)
    
    result = await grader_model.with_structured_output(EnhancedGradeDoc).ainvoke([
        {"role": "user", "content": prompt}
    ])
    
    # 基于置信度和推理进行更智能的判断
    if result.binary_score.lower().startswith("y") and result.confidence > 0.6:
        return "generate_answer"
    else:
        return "rewrite_question"

async def enhanced_rewrite_question(state: MessagesState):
    """增强版问题重写"""
    question = state["messages"][0].content
    prompt = ENHANCED_REWRITE_PROMPT.format(question=question)
    resp = await model.ainvoke([{"role": "user", "content": prompt}])
    return {"messages": [{"role": "user", "content": resp.content}]}

async def enhanced_generate_answer(state: MessagesState):
    """增强版答案生成"""
    question = state["messages"][0].content
    ctx = state["messages"][-1].content
    prompt = ENHANCED_ANSWER_PROMPT.format(question=question, context=ctx)
    resp = await model.ainvoke([{"role": "user", "content": prompt}])
    return {"messages": [resp]}

# 工具节点
async def enhanced_retrieve_node(state: MessagesState):
    """增强版检索节点"""
    last_message = state["messages"][-1]
    
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        tool_call = last_message.tool_calls[0]
        query = tool_call['args'].get('query', '')
        
        # 调用增强版检索工具
        result = enhanced_retriever_tool._run(query)
        
        return {"messages": [{"role": "tool", "content": result, "tool_call_id": tool_call['id']}]}
    else:
        return {"messages": [{"role": "tool", "content": "未找到有效的检索查询"}]}

def enhanced_tools_condition(state: MessagesState):
    """工具条件判断"""
    last_message = state["messages"][-1]
    
    if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
        return "tools"
    else:
        return END

# 构建增强版工作流
enhanced_workflow = StateGraph(MessagesState)
enhanced_workflow.add_node("generate_query_or_respond", enhanced_generate_query_or_respond)
enhanced_workflow.add_node("retrieve", enhanced_retrieve_node)
enhanced_workflow.add_node("rewrite_question", enhanced_rewrite_question)
enhanced_workflow.add_node("generate_answer", enhanced_generate_answer)

enhanced_workflow.add_edge(START, "generate_query_or_respond")
enhanced_workflow.add_conditional_edges(
    "generate_query_or_respond", enhanced_tools_condition, {"tools": "retrieve", END: END}
)
enhanced_workflow.add_conditional_edges("retrieve", enhanced_grade_documents)
enhanced_workflow.add_edge("generate_answer", END)
enhanced_workflow.add_edge("rewrite_question", "generate_query_or_respond")

# 编译增强版RAG Agent
enhanced_rag_agent = enhanced_workflow.compile(name="enhanced_rag_agent")

# 测试函数
async def test_enhanced_rag_agent():
    """测试增强版RAG Agent"""
    print("🧪 测试增强版RAG Agent...")
    
    test_queries = [
        "什么是客户流失？为什么电信行业特别关注这个问题？",
        "数据预处理在客户流失预测中有什么重要作用？具体包括哪些步骤？",
        "如何选择和构建特征来预测客户流失？",
        "客户流失预测模型的评估指标有哪些？如何选择合适的指标？"
    ]
    
    for query in test_queries:
        print(f"\n📝 查询: {query}")
        try:
            result = await enhanced_rag_agent.ainvoke({
                "messages": [{"role": "user", "content": query}]
            })
            print(f"🤖 回答: {result['messages'][-1].content[:500]}...")
        except Exception as e:
            print(f"❌ 错误: {str(e)}")
        print("-" * 80)

if __name__ == "__main__":
    print("🤖 增强版RAG Agent 已启动")
    print("新增功能：")
    print("- 🧠 多轮对话记忆")
    print("- 🔍 智能查询增强")
    print("- 📊 上下文关联分析")
    print("- 💡 相关问题推荐")
    print("- 🎯 专业术语解释")
    
    # 运行测试
    asyncio.run(test_enhanced_rag_agent())