# 安装依赖：pip install langgraph langchain-openai langchain-community faiss-cpu

from typing import TypedDict, List
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate

# 定义状态：所有节点共享这个状态，通过它传递信息
class RAGState(TypedDict):
    question: str               # 用户原始问题
    query: str                  # 当前检索词（可能被优化过）
    documents: List[Document]   # 检索到的文档（累积所有检索结果）
    generation: str             # 最终生成的答案
    retrieval_count: int        # 已检索次数（防止无限循环）
    is_sufficient: bool         # 当前信息是否充分


# 初始化 LLM 和检索器
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
embeddings = OpenAIEmbeddings()

# 实际项目中从文件加载：
# vectorstore = FAISS.load_local("my_index", embeddings)
# retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

MAX_RETRIEVAL = 3  # 最多检索 3 次，防止无限循环


def retrieve(state: RAGState) -> dict:
    """执行一次检索，累积到已有文档列表"""
    docs = retriever.invoke(state["query"])
    return {
        "documents": state["documents"] + docs,  # 累积，不是替换
        "retrieval_count": state["retrieval_count"] + 1
    }


def grade_documents(state: RAGState) -> dict:
    """让 LLM 评估当前文档是否足以回答问题

    这是 Agentic 的关键：用 LLM 替代规则来做判断
    """
    grade_prompt = ChatPromptTemplate.from_template("""
你是一个信息充分性评估员。

用户问题：{question}

当前已检索到的文档内容：
{documents}

判断这些文档是否包含足够的信息来完整回答用户的问题。
评判标准：
1. 文档是否直接包含问题所需的关键信息
2. 信息是否完整（不缺少重要的方面）
3. 信息是否足够具体（不只是泛泛的相关内容）

只回答 "sufficient" 或 "insufficient"，不要任何其他内容。
""")

    docs_text = "\n\n---\n\n".join([
        f"文档 {i+1}:\n{d.page_content}"
        for i, d in enumerate(state["documents"])
    ])

    result = llm.invoke(
        grade_prompt.format_messages(
            question=state["question"],
            documents=docs_text
        )
    )

    is_sufficient = "sufficient" in result.content.lower()
    return {"is_sufficient": is_sufficient}


def refine_query(state: RAGState) -> dict:
    """让 LLM 基于现有文档，生成更精准的查询词

    这是 Agentic 的另一个关键：用 LLM 优化检索策略
    """
    refine_prompt = ChatPromptTemplate.from_template("""
原始问题：{question}
已尝试的查询词：{query}
已检索到的信息摘要：{documents}

已有信息还不够完整，需要继续检索。
分析一下：现在缺少什么信息？用什么关键词检索能找到这些信息？

请生成一个新的、更精准的查询词来补充缺失的信息。
只返回查询词本身，不要解释，不要其他内容。
""")

    docs_text = "\n".join([
        d.page_content[:200] + "..."   # 只取前 200 字作为摘要，避免 token 过多
        for d in state["documents"]
    ])

    result = llm.invoke(
        refine_prompt.format_messages(
            question=state["question"],
            query=state["query"],
            documents=docs_text
        )
    )

    new_query = result.content.strip()
    print(f"[优化查询] 原查询: '{state['query']}' → 新查询: '{new_query}'")
    return {"query": new_query}


def generate(state: RAGState) -> dict:
    """基于所有检索到的文档生成最终答案"""
    generate_prompt = ChatPromptTemplate.from_template("""
基于以下文档，回答用户的问题。如果文档中没有相关信息，请如实说明，不要编造。

文档内容：
{documents}

用户问题：{question}

请提供准确、完整的回答：
""")

    docs_text = "\n\n---\n\n".join([
        f"文档 {i+1}:\n{d.page_content}"
        for i, d in enumerate(state["documents"])
    ])

    result = llm.invoke(
        generate_prompt.format_messages(
            documents=docs_text,
            question=state["question"]
        )
    )

    return {"generation": result.content}


def should_continue_retrieval(state: RAGState) -> str:
    """条件判断：继续检索还是直接生成答案

    这里决定了 Agent 的行为逻辑
    """
    if state["is_sufficient"]:
        print(f"[决策] 信息充分，准备生成答案（共检索 {state['retrieval_count']} 次）")
        return "generate"

    if state["retrieval_count"] >= MAX_RETRIEVAL:
        print(f"[决策] 达到最大检索次数 {MAX_RETRIEVAL}，用现有信息生成答案")
        return "generate"

    print(f"[决策] 信息不足，继续检索（已检索 {state['retrieval_count']} 次）")
    return "refine_query"


# 构建 Agentic RAG 图
workflow = StateGraph(RAGState)

workflow.add_node("retrieve", retrieve)
workflow.add_node("grade_documents", grade_documents)
workflow.add_node("refine_query", refine_query)
workflow.add_node("generate", generate)

# 定义流程
workflow.add_edge(START, "retrieve")                  # 开始就检索
workflow.add_edge("retrieve", "grade_documents")      # 检索后评估
workflow.add_conditional_edges(
    "grade_documents",
    should_continue_retrieval,                         # 根据评估结果决定下一步
    {
        "generate": "generate",                        # 信息充分，生成答案
        "refine_query": "refine_query"                 # 信息不足，优化查询
    }
)
workflow.add_edge("refine_query", "retrieve")          # 优化后重新检索
workflow.add_edge("generate", END)

app = workflow.compile()


def ask(question: str) -> str:
    """对外的问答接口"""
    initial_state = RAGState(
        question=question,
        query=question,      # 初始查询词就是原始问题
        documents=[],
        generation="",
        retrieval_count=0,
        is_sufficient=False
    )

    result = app.invoke(initial_state)
    print(f"\n共检索 {result['retrieval_count']} 次，使用了 {len(result['documents'])} 份文档")
    return result["generation"]


# 使用示例（需要先建立好 retriever）
# answer = ask("张三负责的项目用了哪些技术，各有什么优缺点？")
# print(answer)