from typing import TypedDict, Literal, List
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain_core.messages import HumanMessage, SystemMessage

# ============================================================
# 第一步：定义状态结构
# 状态就像一个"工作台"，记录整个检索过程中的所有信息
# ============================================================
class AgenticRAGState(TypedDict):
    question: str                  # 用户的原始问题
    needs_retrieval: bool          # 是否需要检索
    retrieved_docs: List[str]      # 已检索到的文档内容
    retrieval_quality: str         # 检索质量评估结果: "sufficient" / "insufficient"
    iteration_count: int           # 已迭代次数（防止无限循环）
    search_query: str              # 当前使用的检索关键词
    answer: str                    # 最终生成的答案

# ============================================================
# 第二步：初始化工具
# ============================================================
llm = ChatOpenAI(model="gpt-4o", temperature=0)
embeddings = OpenAIEmbeddings()

# 假设已有一个向量数据库（见第01章：如何构建向量库）
vectorstore = Chroma(
    collection_name="company_docs",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

# ============================================================
# 第三步：定义各个节点函数
# ============================================================

def assess_retrieval_need(state: AgenticRAGState) -> AgenticRAGState:
    """
    节点1：判断是否需要检索
    LLM评估问题类型，决定是否触发检索流程
    """
    response = llm.invoke([
        SystemMessage(content="""你是一个智能路由助手。
        判断用户问题是否需要检索企业知识库。
        需要检索的情况：涉及公司内部政策、产品信息、业务流程、私有数据。
        不需要检索的情况：通用常识、数学计算、编程语法、历史知识。
        只回答 yes 或 no。"""),
        HumanMessage(content=f"问题：{state['question']}")
    ])

    needs_retrieval = response.content.strip().lower() == "yes"
    return {**state, "needs_retrieval": needs_retrieval, "iteration_count": 0}


def generate_search_query(state: AgenticRAGState) -> AgenticRAGState:
    """
    节点2：生成优化的检索关键词
    将用户问题转化为更适合向量检索的查询语句
    """
    response = llm.invoke([
        SystemMessage(content="将用户问题转化为适合文档检索的简洁关键词，只输出关键词，不要解释。"),
        HumanMessage(content=state["question"])
    ])
    search_query = response.content.strip()
    return {**state, "search_query": search_query}


def retrieve_documents(state: AgenticRAGState) -> AgenticRAGState:
    """
    节点3：执行检索
    """
    docs = retriever.invoke(state["search_query"])
    doc_contents = [doc.page_content for doc in docs]
    # 合并新旧检索结果（迭代检索时会累积）
    all_docs = state.get("retrieved_docs", []) + doc_contents
    return {**state, "retrieved_docs": all_docs}


def evaluate_retrieval_quality(state: AgenticRAGState) -> AgenticRAGState:
    """
    节点4：评估检索质量
    判断已有文档是否足以回答用户问题
    """
    docs_text = "\n---\n".join(state["retrieved_docs"])

    response = llm.invoke([
        SystemMessage(content="""评估检索到的文档是否足以回答用户问题。
        如果文档包含足够的相关信息，回答 sufficient。
        如果文档信息不足、不相关或缺少关键内容，回答 insufficient。
        只回答 sufficient 或 insufficient。"""),
        HumanMessage(content=f"用户问题：{state['question']}\n\n检索到的文档：\n{docs_text}")
    ])

    quality = response.content.strip().lower()
    return {**state, "retrieval_quality": quality}


def generate_answer(state: AgenticRAGState) -> AgenticRAGState:
    """
    节点5：生成最终答案
    """
    docs_text = "\n---\n".join(state.get("retrieved_docs", []))
    context_prompt = f"参考文档：\n{docs_text}\n\n" if docs_text else ""

    response = llm.invoke([
        SystemMessage(content=f"你是企业知识库助手。根据提供的文档回答问题，如信息不完整请说明。{context_prompt}"),
        HumanMessage(content=state["question"])
    ])
    return {**state, "answer": response.content}


def refine_search_query(state: AgenticRAGState) -> AgenticRAGState:
    """
    节点6：优化检索策略（用于迭代检索）
    在前一次检索不够时，换一个角度重新生成查询词
    """
    response = llm.invoke([
        SystemMessage(content="前次检索结果不足以回答问题。请从不同角度生成新的检索关键词。只输出关键词。"),
        HumanMessage(content=f"原始问题：{state['question']}\n已用查询词：{state['search_query']}")
    ])
    new_query = response.content.strip()
    return {
        **state,
        "search_query": new_query,
        "iteration_count": state["iteration_count"] + 1
    }

# ============================================================
# 第四步：定义条件边（决策逻辑）
# ============================================================

def route_after_assessment(state: AgenticRAGState) -> Literal["generate_search_query", "generate_answer"]:
    """判断是否跳过检索直接生成"""
    if state["needs_retrieval"]:
        return "generate_search_query"
    return "generate_answer"


def route_after_evaluation(state: AgenticRAGState) -> Literal["generate_answer", "refine_search_query"]:
    """判断是否需要继续检索"""
    MAX_ITERATIONS = 3  # 最多迭代3次，防止无限循环
    if state["retrieval_quality"] == "sufficient" or state["iteration_count"] >= MAX_ITERATIONS:
        return "generate_answer"
    return "refine_search_query"

# ============================================================
# 第五步：构建图结构
# ============================================================
graph_builder = StateGraph(AgenticRAGState)

# 添加节点
graph_builder.add_node("assess_retrieval_need", assess_retrieval_need)
graph_builder.add_node("generate_search_query", generate_search_query)
graph_builder.add_node("retrieve_documents", retrieve_documents)
graph_builder.add_node("evaluate_retrieval_quality", evaluate_retrieval_quality)
graph_builder.add_node("generate_answer", generate_answer)
graph_builder.add_node("refine_search_query", refine_search_query)

# 添加边（流程连接）
graph_builder.add_edge(START, "assess_retrieval_need")
graph_builder.add_conditional_edges("assess_retrieval_need", route_after_assessment)
graph_builder.add_edge("generate_search_query", "retrieve_documents")
graph_builder.add_edge("retrieve_documents", "evaluate_retrieval_quality")
graph_builder.add_conditional_edges("evaluate_retrieval_quality", route_after_evaluation)
graph_builder.add_edge("refine_search_query", "retrieve_documents")  # 迭代回到检索节点
graph_builder.add_edge("generate_answer", END)

# 编译成可执行的图
agentic_rag = graph_builder.compile()

# ============================================================
# 第六步：使用
# ============================================================
result = agentic_rag.invoke({
    "question": "我们公司差旅报销的上限是多少，出租车票需要发票吗？",
    "retrieved_docs": [],
    "search_query": "",
    "answer": ""
})

print(result["answer"])

# 场景模拟：HR 助手，对接公司 HR 知识库
# 展示三种不同问题的处理路径

test_questions = [
    # 场景1：不需要检索（LLM直接回答）
    "Python 列表和元组的区别是什么？",

    # 场景2：简单检索（一次检索即可）
    "公司的年假天数政策是什么？",

    # 场景3：多跳推理（需要多步检索）
    "上海研发中心的试用期考核标准和北京总部一样吗？",
]

for question in test_questions:
    print(f"\n{'='*50}")
    print(f"问题：{question}")

    result = agentic_rag.invoke({
        "question": question,
        "retrieved_docs": [],
        "search_query": "",
        "answer": ""
    })

    print(f"是否触发检索：{result['needs_retrieval']}")
    print(f"检索迭代次数：{result['iteration_count']}")
    print(f"答案：{result['answer'][:200]}...")