"""
LangChain + LangGraph 集成示例
数据源：docs/菌群知识库/_chunks.jsonl（22,781 个 chunk）

两种用法：
  1. RAG 检索增强（推荐 Qdrant/Pinecone）
  2. LangGraph 多步 Agent（带分类路由）
"""

# ============================================================
# 方式一：RAG 检索增强（最小路径）
# ============================================================

from langchain_community.document_loaders import JSONLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OpenAIEmbeddings  # 或 BGE / text-embedding-3-large
from langchain_community.vectorstores import QdrantVectorStore
from langchain_community.vectorstores.utils import DistanceStrategy
from langchain_core.runnables import RunnableParallel, RunnableLambda

# 1. 加载 JSONL
loader = JSONLoader(
    file_path=r"C:\code\cfc\docs\菌群知识库\_chunks.jsonl",
    jq_schema=".[] | {page_content: .content, metadata: .metadata}",
    content_key="page_content",
    metadata_key="metadata",
)
docs = loader.load()

# 2. Embedding — 中文推荐 bge-m3 或 text-embedding-3-large
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# 3. 向量化存储（Qdrant 本地示例）
vector_store = QdrantVectorStore.from_documents(
    documents=docs,
    embedding=embeddings,
    collection_name="guhejk_kb",
    url="http://localhost:6333",
    distance=DistanceStrategy.COSINE,
)

# 4. RAG 检索器（支持 metadata 过滤）
retriever = vector_store.as_retriever(
    search_type="mmr",
    search_kwargs={
        "k": 5,
        "filter": {
            "must": [{"key": "category", "match": {"value": "科普"}}]  # 可选：限定分类
        }
    }
)

# 5. 简单问答链
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是菌群健康领域的专家助手。请根据以下资料回答用户问题，引用具体来源。"),
    ("human", "问题：{question}\n\n参考资料：{context}"),
])
llm = ChatOpenAI(model="gpt-4o", temperature=0)
rag_chain = prompt | llm

# 使用
result = rag_chain.invoke({
    "question": "肠道菌群和抑郁有什么关系？",
    "context": retriever.invoke("肠道菌群 抑郁")
})
print(result.content)


# ============================================================
# 方式二：LangGraph 多步 Agent（带分类路由 + 总结）
# ============================================================

from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
from langgraph.checkpoint.memory import MemorySaver

class AgentState(TypedDict):
    question: str
    category: str           # 自动推断的分类
    retrieved_chunks: list  # 检索到的 chunk
    summary: str
    final_answer: str

# 节点 1：分类路由
def classify_question(state: AgentState) -> AgentState:
    """根据用户问题自动推断应检索哪个分类"""
    q = state["question"]
    routing = {
        ("检测", "报告", "测序"): "技术服务",
        ("临床", "患者", "诊断"): "科普",
        ("技术", "方法", "测序"): "技术",
        ("公告", "新闻"): "公告",
        ("报告解读"): "报告解读",
        ("视频"): "视频",
    }
    for kws, cat in routing.items():
        if all(kw in q for kw in kws):
            return {"category": cat}
    return {"category": "科普"}  # 默认科普

# 节点 2：检索
def retrieve(state: AgentState) -> AgentState:
    cat = state.get("category", "科普")
    chunks = retriever.invoke(state["question"])
    return {"retrieved_chunks": chunks}

# 节点 3：总结
def summarize(state: AgentState) -> AgentState:
    texts = "\n\n".join([c.page_content for c in state["retrieved_chunks"]])
    prompt_summary = ChatPromptTemplate.from_messages([
        ("system", "你是一个科学文章总结助手。请总结以下内容，突出关键发现和数据。"),
        ("human", "{texts}"),
    ])
    summary = (prompt_summary | ChatOpenAI(model="gpt-4o-mini", temperature=0)).invoke({"texts": texts})
    return {"summary": summary.content}

# 节点 4：生成最终回答
def answer(state: AgentState) -> AgentState:
    ctx = "\n\n---\n\n".join([c.page_content for c in state["retrieved_chunks"]])
    final = rag_chain.invoke({"question": state["question"], "context": ctx})
    return {"final_answer": final.content}

# 构建图
graph = StateGraph(AgentState)
graph.add_node("classify", classify_question)
graph.add_node("retrieve", retrieve)
graph.add_node("summarize", summarize)
graph.add_node("answer", answer)
graph.add_edge("classify", "retrieve")
graph.add_edge("retrieve", "summarize")
graph.add_edge("summarize", "answer")
graph.set_entry_point("classify")
graph.set_finish_point("answer")

memory = MemorySaver()
app = graph.compile(checkpointer=memory)

# 使用
config = {"configurable": {"thread_id": "user-001"}}
result = app.invoke(
    {"question": "肠道菌群和抑郁有什么关系？宏基因组检测能发现什么？"},
    config=config,
)
print(result["final_answer"])


# ============================================================
# 方式三：在 cfc-backend (Java) 中调用 LangGraph API
# ============================================================
"""
// 假设 LangGraph 服务运行在 localhost:8123
POST /invoke
{
  "question": "肠道菌群与糖尿病的关系",
  "category": "科普"
}

Response:
{
  "final_answer": "...",
  "summary": "...",
  "sources": [
    {"postId": "1234", "title": "...", "category": "科普"},
    ...
  ]
}
"""

# ============================================================
# 常用搜索示例
# ============================================================

# 按分类过滤检索
retriever_science = vector_store.as_retriever(
    search_kwargs={"filter": {"must": [{"key": "category", "match": {"value": "科普"}}]}}
)

# 按主题过滤（topics metadata）
retriever_depression = vector_store.as_retriever(
    search_kwargs={"filter": {"must": [{"key": "topics", "match": {"any": ["抑郁", "焦虑"]}}]}}
)

# 多路检索合并
retrievers = RunnableParallel({
    "science": retriever_science,
    "clinical": vector_store.as_retriever(search_kwargs={"filter": {"must": [{"key": "topics", "match": {"any": ["临床", "诊断"]}}]}}),
})
merged = retrievers | (lambda x: sorted(set([c for chunk_list in x.values() for c in chunk_list]), key=lambda c: c.metadata["pub_date"] or ""))