from langchain_core.documents import Document
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from typing import TypedDict, List
from langgraph.graph import StateGraph, END
import sys
from elasticsearch import Elasticsearch
from dotenv import load_dotenv

load_dotenv()

# 配置es客户端
es_client = None
es_index = None

try:
    es_client = Elasticsearch(es_url, basic_auth=(es_user, es_password))
except Exception as e:
    import traceback
    traceback.print_exc()
    sys.exit(1)

# 这里使用腾讯云DeepSeek-v3模型
model = ChatOpenAI(
    model="deepseek-v3",
    base_url="https://api.lkeap.cloud.tencent.com/v1")

def preprocess_query1(state: dict) -> dict:
    """模块A: Query预处理1"""
    print(f"--- 模块A: 预处理Query1 ---")
    state['preprocessed_question'] = state['question'].strip().lower()
    return state

def preprocess_query2(state: dict) -> dict:
    """模块B: Query预处理2"""
    print(f"--- 模块B: 预处理Query2 ---")
    state['preprocessed_question'] = state['preprocessed_question'].replace('\n', ' ')
    return state

# 可以增加其他预处理函数...

def retrieve_documents(state: dict) -> dict:
    """模块C&D: 检索模块 & 获取相关文档"""
    print(f"--- 模块C&D: 检索文档 ---")
    
    try:
        query = {
            "query": {"match": {"content": state['preprocessed_question']}},
            "size": 2
        }
        results = es_client.search(index=es_index, body=query)
        
        documents = []
        for hit in results['hits']['hits']:
            source = hit['_source']
            doc = Document(
                page_content=source.get('content', ''),
                metadata={'score': hit['_score'], 'id': source.get('id', '')}
            )
            documents.append(doc)
        state['documents'] = documents
    except Exception as e:
        import traceback
        traceback.print_exc()
        state['documents'] = []
        
    return state

def rerank_and_filter(state: dict) -> dict:
    """模块E: 重排 & 过滤模块"""
    print(f"--- 模块E: 重排与过滤 ---")
    # E:通过es score或者LLM模型重排
    # 此处省略代码...
    # F:通过阈值过滤
    threshold = 0.7
    
    filtered_docs = []
    for doc in state['documents']:
        score = doc.metadata.get('score', 0)
        score = float(score)
        if score >= threshold:
            filtered_docs.append(doc)
    
    state['filtered_documents'] = filtered_docs if filtered_docs else state['documents']
    return state

def build_prompt(state: dict) -> dict:
    """模块F: Prompt构建模块"""
    print(f"--- 模块F: 构建Prompt ---")
    template = """仅根据以下上下文回答问题:
    ---
    {context}
    ---
    问题: {question}
    """
    prompt_template = ChatPromptTemplate.from_template(template)
    
    context_str = "\n\n".join([doc.page_content for doc in state['filtered_documents']])
    state['prompt'] = prompt_template.format(context=context_str, question=state['question'])
    return state

def llm_call(state: dict) -> str:
    """模块G: LLM调用"""
    print(f"--- 模块G: 调用腾讯云DeepSeek-v3模型 ---")
    response = model.invoke(state['prompt'])
    return response.content

def post_process_answer(state: dict) -> dict:
    """模块H: 生成答案 & 后处理"""
    print(f"--- 模块H: 后处理答案 ---")
    state['final_answer'] = state['generation'].strip()
    return state

lcel_chain = (
    RunnableLambda(preprocess_query1)
    | RunnableLambda(preprocess_query2)
    | RunnableLambda(retrieve_documents)
    | RunnableLambda(rerank_and_filter)
    | RunnableLambda(build_prompt)
    | {
        "generation": RunnableLambda(llm_call),  # 模块H: LLM调用
        "final_state": RunnablePassthrough()    # 将整个状态传递下去
    }
    | {
        "generation": lambda x: x['generation'],  # 提取内容
        "final_state": lambda x: x['final_state']
    }
    | RunnableLambda(lambda state_dict: post_process_answer({**state_dict['final_state'], "generation": state_dict['generation']}))
)

==================== 通过LCEL 实现 RAG ====================
[DEBUG] 构建LCEL处理链...
[DEBUG] LCEL链构建完成
==================== 执行LCEL链 ====================
[DEBUG] 开始执行LCEL链，问题: '什么是langgraph?'
--- 模块A: 预处理Query1 ---
--- 模块B: 预处理Query2 ---
--- 模块C&D: 检索文档 ---
--- 模块E: 重排与过滤 ---
--- 模块F: 构建Prompt ---
--- 模块G: 调用腾讯云DeepSeek-v3模型 ---
--- 模块H: 后处理答案 ---
==================== [LCEL] 最终结果 ====================
根据提供的上下文，没有提到与"langgraph"相关的任何信息。因此，我无法回答这个问题。

# 1. 定义中心化的状态
class AdvancedRAGState(TypedDict):
    question: str
    preprocessed_question: str
    rewritten_question: str
    documents: List[Document]
    filtered_documents: List[Document]
    prompt: str
    generation: str
    final_answer: str

# 2. 将业务逻辑函数适配为LangGraph节点
# 节点的签名是 (State) -> PartialStateUpdate
def preprocess_query1_node(state: AdvancedRAGState):
    return preprocess_query1(state)

def preprocess_query2_node(state: AdvancedRAGState):
    return preprocess_query2(state)

def retrieve_documents_node(state: AdvancedRAGState):
    return retrieve_documents(state)

def rerank_and_filter_node(state: AdvancedRAGState):
    return rerank_and_filter(state)

def build_prompt_node(state: AdvancedRAGState):
    return build_prompt(state)

def generate_answer_node(state: AdvancedRAGState):
    # 模块H: LLM接口调用
    print(f"--- 模块H: LLM接口调用 ---")
    generation = model.invoke(state['prompt'])
    return {"generation": generation.content}

def post_process_node(state: AdvancedRAGState):
    return post_process_answer(state)

# 3. 构建图并定义流程
workflow = StateGraph(AdvancedRAGState)

# 添加节点
workflow.add_node("preprocess1", preprocess_query1_node)
workflow.add_node("preprocess2", preprocess_query2_node)
workflow.add_node("retrieve", retrieve_documents_node)
workflow.add_node("rerank", rerank_and_filter_node)
workflow.add_node("build_prompt", build_prompt_node)
workflow.add_node("generate", generate_answer_node)
workflow.add_node("post_process", post_process_node)

# 设置入口点并定义边的连接
workflow.set_entry_point("preprocess1")
workflow.add_edge("preprocess1", "preprocess2")
workflow.add_edge("preprocess2", "retrieve")
workflow.add_edge("retrieve", "rerank")
workflow.add_edge("rerank", "build_prompt")
workflow.add_edge("build_prompt", "generate")
workflow.add_edge("generate", "post_process")
workflow.add_edge("post_process", END)

app = workflow.compile()

# 执行LangGraph
langgraph_result = app.invoke({"question": "什么是langgraph? "})
print("--- [LangGraph] 最终结果 ---")
print(langgraph_result)

==================== 通过LangGraph 实现 RAG ====================
--- 模块A: 预处理Query1 ---
--- 模块B: 预处理Query2 ---
--- 模块C&D: 检索文档 ---
--- 模块E: 重排与过滤 ---
--- 模块F: 构建Prompt ---
--- 模块H: LLM接口调用 ---
--- 模块H: 后处理答案 ---
--- [LangGraph] 最终结果 ---
{'question': '什么是langgraph? ', 
 'preprocessed_question': '什么是langgraph?',
 'documents':[Document(metadata={'score': 15.526812, 'id': '2'}, page_content='In the Launch state, the weapon will continue to navigate based on IMU ... during this state.')], 
 'filtered_documents': [Document(metadata={'score': 15.773178, 'id': '3'}, page_content='These factors affect weapon performance, the degree to which depends on the flight characteristics... performance of the weapon (e.g.'), Document(metadata={'score': 15.526812, 'id': '4'}, page_content='In the Launch state, the ... during this state.')], 
 'prompt': 'Human: 仅根据以下上下文回答问题:\n These factors affect weapon performance, the degree to which depends on the flight characteristics... during this state.\n 问题: 什么是langgraph? \n', 
 'generation': '根据提供的上下文，没有提到"langgraph"这一概念或相关信息，因此无法回答这个问题。上下文主要讨论了武器性能的影响因素、内部组件特性以及武器在发射状态下的操作（如导航、通信和遥测数据发送），但未涉及任何与"langgraph"相关的内容。', 
 'final_answer': '根据提供的上下文，没有提到"langgraph"这一概念或相关信息，因此无法回答这个问题。上下文主要讨论了武器性能的影响因素、内部组件特性以及武器在发射状态下的操作（如导航、通信和遥测数据发送），但未涉及任何与"langgraph"相关的内容。'}