#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
本地RAG Agent - 使用本地文本数据库
不依赖外部API的向量搜索，使用简单文本匹配
"""

from __future__ import annotations
import os
import json
import asyncio
from typing import Literal, List, Dict, Any
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_core.documents import Document

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

class LocalTextRetriever(BaseTool):
    """本地文本检索工具"""
    name: str = "retrieve_telco_handbook"
    description: str = "Search and return relevant sections from the Telco Customer Churn Analytics Handbook."
    
    def _run(self, query: str, run_manager: CallbackManagerForToolRun = None) -> str:
        """检索相关文档内容"""
        try:
            # 加载本地文本数据库
            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)
            
            # 执行搜索
            results = self._simple_search(sections, query, top_k=3)
            
            if not results:
                return f"未找到与'{query}'相关的信息。"
            
            # 格式化结果
            formatted_results = []
            for i, result in enumerate(results, 1):
                formatted_results.append(f"【相关内容 {i}】\n{result['title']}\n\n{result['content'][:500]}...")
            
            return "\n\n".join(formatted_results)
            
        except Exception as e:
            return f"检索时发生错误：{str(e)}"
    
    def _simple_search(self, sections: List[Dict], query: str, top_k: int = 3) -> List[Dict]:
        """简单的文本搜索"""
        results = []
        query_lower = query.lower()
        
        # 扩展查询词
        query_terms = self._expand_query(query_lower)
        
        for section in sections:
            score = 0
            content_lower = section['content'].lower()
            title_lower = section['title'].lower()
            
            # 计算匹配得分
            for term in query_terms:
                # 标题匹配得分更高
                if term in title_lower:
                    score += 10
                
                # 内容匹配
                score += content_lower.count(term) * 2
                
                # 关键词匹配
                for keyword in section.get('keywords', []):
                    if term in keyword.lower():
                        score += 5
            
            if score > 0:
                results.append({
                    'section': section,
                    'score': score,
                    'title': section['title'],
                    'content': section['content']
                })
        
        # 按得分排序
        results.sort(key=lambda x: x['score'], reverse=True)
        return results[:top_k]
    
    def _expand_query(self, query: str) -> List[str]:
        """扩展查询词"""
        terms = [query]
        
        # 添加相关术语映射
        term_mapping = {
            '流失': ['churn', '客户流失', '用户流失', '离网'],
            '合同': ['contract', '合约', '签约'],
            '账单': ['charges', '费用', '付费', '缴费'],
            '服务': ['service', '业务', '功能'],
            '预测': ['prediction', '建模', '模型'],
            '特征': ['feature', '字段', '变量'],
            '数据': ['data', '信息', '记录'],
            '客户': ['customer', '用户', '客人'],
            '分析': ['analysis', '统计', '分析']
        }
        
        for key, synonyms in term_mapping.items():
            if key in query:
                terms.extend(synonyms)
            for syn in synonyms:
                if syn in query:
                    terms.append(key)
                    terms.extend([s for s in synonyms if s != syn])
                    break
        
        return list(set(terms))

# 创建检索工具实例    
retriever_tool = LocalTextRetriever()

# Prompt定义
SYSTEM_INSTRUCTION = (
    "你是一个电信数据分析助手。你可以回答与电信客户流失分析手册相关的问题，"
    "包括客户流失预测、数据字典、建模策略、可视化分析、业务洞察和特征工程等主题。"
    "如果用户的问题不相关，请回复：'我只能回答与电信客户流失分析手册相关的问题。'"
    "当你需要更多上下文信息时，可以调用工具 `retrieve_telco_handbook`。"
)

GRADE_PROMPT = (
    "你是一个评估检索文档与用户问题相关性的评分员。\n"
    "检索到的文档：\n{context}\n\n用户问题：{question}\n"
    "如果相关，返回 'yes'，否则返回 'no'。"
)

REWRITE_PROMPT = (
    "你是一个电信领域专家助手。\n"
    "给定一个关于电信客户流失分析手册的用户问题，"
    "你的任务是重写或扩展它，以便更清楚地捕获潜在意图，"
    "同时提高基于文档检索的召回能力。使用正式和明确的措辞。\n\n"
    "原始问题：\n{question}\n\n"
    "改进版本："
)

ANSWER_PROMPT = (
    "你是一个训练于电信客户流失分析手册的专家助手。\n"
    "请使用提供的上下文以详细、结构化的方式回答问题。\n"
    "你的答案应该包括：\n"
    "- 相关概念或发现的简洁总结\n"
    "- 其背景或基本原理的解释\n"
    "- 实用建议或使用建议\n"
    "- 文档中任何支持的统计数据或示例\n\n"
    "如果上下文缺乏相关信息，请回复：'很抱歉，我无法从手册中找到相关信息。'\n\n"
    "问题：\n{question}\n\n"
    "上下文：\n{context}\n\n"
    "答案："
)

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

class GradeDoc(BaseModel):
    binary_score: str = Field(description="Relevance score 'yes' or 'no'.")

async def grade_documents(state: MessagesState) -> Literal["generate_answer", "rewrite_question"]:
    """评估文档相关性"""
    question = state["messages"][0].content  # 原始用户问题
    ctx = state["messages"][-1].content      # 检索器输出
    prompt = GRADE_PROMPT.format(question=question, context=ctx)
    result = await grader_model.with_structured_output(GradeDoc).ainvoke([
        {"role": "user", "content": prompt}
    ])
    return "generate_answer" if result.binary_score.lower().startswith("y") else "rewrite_question"

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

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

# 工具节点
async def 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 = retriever_tool._run(query)
        
        return {"messages": [{"role": "tool", "content": result, "tool_call_id": tool_call['id']}]}
    else:
        return {"messages": [{"role": "tool", "content": "未找到有效的检索查询"}]}

def 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

# 构建工作流
workflow = StateGraph(MessagesState)
workflow.add_node("generate_query_or_respond", generate_query_or_respond)
workflow.add_node("retrieve", retrieve_node)
workflow.add_node("rewrite_question", rewrite_question)
workflow.add_node("generate_answer", generate_answer)

workflow.add_edge(START, "generate_query_or_respond")
workflow.add_conditional_edges(
    "generate_query_or_respond", tools_condition, {"tools": "retrieve", END: END}
)
workflow.add_conditional_edges("retrieve", grade_documents)
workflow.add_edge("generate_answer", END)
workflow.add_edge("rewrite_question", "generate_query_or_respond")

# 编译本地RAG Agent
local_rag_agent = workflow.compile(name="local_rag_agent")

# 测试函数
async def test_local_rag_agent():
    """测试本地RAG Agent"""
    print("🧪 测试本地RAG Agent...")
    
    test_queries = [
        "什么是客户流失？",
        "影响客户流失的主要因素有哪些？",
        "数据集包含哪些字段？",
        "如何进行数据预处理？"
    ]
    
    for query in test_queries:
        print(f"\n📝 查询: {query}")
        try:
            result = await local_rag_agent.ainvoke({
                "messages": [{"role": "user", "content": query}]
            })
            print(f"🤖 回答: {result['messages'][-1].content}")
        except Exception as e:
            print(f"❌ 错误: {str(e)}")
        print("-" * 50)

if __name__ == "__main__":
    # 运行测试
    asyncio.run(test_local_rag_agent())