"""中医知识图谱问答系统 —— FastAPI 接口层

按 API_INTERFACE.md 实现三个接口：

    GET  /                健康检查
    POST /process         中医问答（核心接口）
    POST /process_stream  中医问答（流式，逐节点返回思考过程）

这一层很薄：所有业务逻辑都在 `_004_langgraph_more_nodes.langgraph_more_nodes` 编译好的
图里，接口只做三件事 —— 校验入参、把图跑起来、把 `state["output"]` 取出来返回。
`output` 是整张图的统一出口（图谱路和兜底路都会写它），所以这里不用关心答案是查图谱
得来的还是大模型直接给的。

两个和"薄"有关的设计：

- **图是懒加载的模块级单例**（`get_graph()`）。导入这张图会连带加载 76 MB 的 FAISS
  索引和 embedding 模型，几十秒量级，所以不在 import 时做，也不每个请求做一次。
  `lifespan` 在服务启动时提前把它加载好，这样首个请求不会额外背上初始化的时间；
  加载失败只打日志不阻止启动，健康检查照样能用来看服务活着。
- **图是同步阻塞的**（LLM 调用 + Neo4j 查询，单次 5-30 秒），而接口是 async 的。
  直接 `graph.invoke()` 会把事件循环占死，一个请求就让整个服务失去并发能力，
  所以必须丢进 `run_in_threadpool`。

运行（在仓库根目录，导入用的是绝对包名）:

    python -X utf8 -m _005_fastapi.main
"""

import asyncio
import json
import traceback
from contextlib import asynccontextmanager

import uvicorn
from fastapi import FastAPI
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import JSONResponse, StreamingResponse
from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel, Field, field_validator

# 出错时返回给用户的固定文案，和 API_INTERFACE.md 的错误响应示例保持一致
ERROR_OUTPUT = "系统出错了，请重试！"

# 流式接口里，这些键每个节点都会带回来，但对"思考过程"没有信息量，逐个节点重复推送
# 只会把有用的内容挤掉，所以不进 thinking 消息：
#   input                 每条消息都一样，最后的 final 里已经有了
#   history_messages      跨轮记忆，会越来越长，且每轮都变
#   cypher_retry_count    内部计数
#   output/neo4j_answer/direct_out  最终答案，final 消息负责给
STREAM_SKIP_KEYS = frozenset(
    {
        "input",
        "history_messages",
        "cypher_retry_count",
        "output",
        "neo4j_answer",
        "direct_out",
    }
)

# 队列结束标记。用一个独一无二的对象而不是 None，因为 None 本身可能是正常数据
_STREAM_DONE = object()

# 图的模块级缓存。测试可以直接把这个变量替换成假图，绕开 LLM 和 Neo4j
_graph = None


def get_graph():
    """懒加载并缓存编译好的 LangGraph。

    import 放在函数体内是故意的：模块顶层 import 会让 `import main` 就去加载
    FAISS 索引和 embedding 模型。
    """
    global _graph
    if _graph is None:
        from _004_langgraph_more_nodes.langgraph_more_nodes import graph

        _graph = graph
    return _graph


@asynccontextmanager
async def lifespan(app: FastAPI):
    """服务启动时预热图，省掉首个请求的初始化等待。

    加载失败只记日志、不抛异常：让服务照样起来，健康检查能回 200，
    /process 会在真正被调用时再报错，比整个服务起不来好排查。
    """
    try:
        await run_in_threadpool(get_graph)
        print("LangGraph 预加载完成")
    except Exception:
        traceback.print_exc()
        print("LangGraph 预加载失败，服务仍会启动，/process 将返回错误响应")
    yield


app = FastAPI(
    title="中医知识图谱问答系统",
    description="基于 FastAPI + LangGraph + Neo4j + FAISS 的中医知识图谱智能问答接口",
    version="1.0.0",
    lifespan=lifespan,
)


# ======================
# 请求 / 响应模型
# ======================
class ProcessRequest(BaseModel):
    """POST /process 的请求体。两个字段都是必填，见接口文档的请求参数表。"""

    input: str = Field(..., description="用户输入的问题")
    user_id: str = Field(..., description="用户唯一标识，用于隔离不同用户的会话")

    @field_validator("input", "user_id")
    @classmethod
    def not_blank(cls, value: str) -> str:
        """全空格的入参和空字符串一样没有意义，一并挡掉，并顺手去掉首尾空白"""
        value = value.strip()
        if not value:
            raise ValueError("不能为空")
        return value


class ProcessResponse(BaseModel):
    """POST /process 的响应体：原始输入 + 系统回答"""

    input: str = Field(..., description="用户原始输入")
    output: str = Field(..., description="系统生成的回答")


class HealthResponse(BaseModel):
    """GET / 的响应体"""

    status: str = Field(..., description="服务状态，正常为 ok")
    message: str = Field(..., description="状态说明")


# ======================
# 流式接口的辅助函数
# ======================
def _json_safe(value):
    """把状态值转成能 json.dumps 的形状。

    状态里绝大多数是 str/list/bool，但 cypher_results 里是 Neo4j 驱动返回的行，
    真出现不认识的类型时不能让整条流断掉，退化成 str() 至少还能看。
    """
    if value is None or isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, dict):
        return {str(k): _json_safe(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_json_safe(v) for v in value]
    return str(value)


def _state_delta(previous: dict, current: dict) -> dict:
    """算出这个节点相比上一个节点新写了什么。

    我们的节点都是"改完整个 state 再返回"，所以 LangGraph 的 updates 模式给回来的是
    全量状态，不是增量 —— 直接推给前端每条消息都有二十多个键。这里跟上一步的状态做差，
    才是接口文档里 data 想要的"该节点这一步产生的状态增量"。
    """
    delta = {}
    for key, value in current.items():
        if key in STREAM_SKIP_KEYS:
            continue
        # 语义转写节点会把上一轮的残留键清成 None，这种"清空"不算这一步的产出
        if value is None or value == [] or value == {}:
            continue
        if previous.get(key) != value:
            delta[key] = _json_safe(value)
    return delta


# ======================
# 接口 1：健康检查
# ======================
@app.get("/", response_model=HealthResponse, summary="健康检查")
async def health_check() -> HealthResponse:
    """检查服务是否正常运行。无请求参数，永远返回 200。

    这里刻意不去探测 Neo4j / LLM / FAISS：健康检查要快且不花钱，
    它回答的是"HTTP 服务活着吗"，而不是"所有依赖都好吗"。
    """
    return HealthResponse(status="ok", message="中医知识图谱问答系统服务运行正常")


# ======================
# 接口 2：中医问答（核心）
# ======================
@app.post("/process", response_model=ProcessResponse, summary="中医问答接口")
async def process(request: ProcessRequest) -> ProcessResponse:
    """跑一遍 LangGraph，把最终回答返回给调用方。


    graph.invoke 是同步阻塞的（LLM + Neo4j，单次 5-30 秒），所以扔到线程池里，
    否则事件循环被占死，服务就退化成一次只能处理一个请求。

    任何异常都收敛成 500 + 固定文案：内部报错细节（cypher、堆栈）不应该漏给调用方，
    完整堆栈打在服务端日志里。
    """
    print(f"[/process] user_id={request.user_id} input={request.input!r}")
    try:
        graph = get_graph()
        # thread_id="user_0001"
        thread_id=request.user_id
        config = RunnableConfig(configurable={
            "thread_id": thread_id
        })
        state = await run_in_threadpool(graph.invoke, {"input": request.input},config)
        print("===================")
        print(state)
        print("===================")
        # output 是图的统一出口。理论上每条路径都会写它，兜底一层防止图改坏了返回 None
        output = (state.get("output") or "").strip() or ERROR_OUTPUT
        return ProcessResponse(input=request.input, output=output)
    except Exception:
        traceback.print_exc()
        # 错误响应的 input 按文档示例返回空串
        return JSONResponse(
            status_code=500,
            content={"input": "", "output": ERROR_OUTPUT},
        )


# ======================
# 接口 3：流式中医问答
# ======================
async def _stream_messages(user_input: str, user_id: str):
    """生产者-消费者：把图的执行过程变成一行行 JSON。

    生产者是图里的各个节点 —— 每个节点执行完，LangGraph 的 stream(updates) 就吐出
    一条它的状态，producer 把它塞进队列。消费者是这个异步生成器，从队列取出来
    yield 给 StreamingResponse。

    为什么要队列：`graph.stream()` 是**同步**生成器，里面是 LLM 调用和 Neo4j 查询，
    一次 5-30 秒。直接在 async 函数里 for 它，事件循环会被卡住，服务退化成单并发
    （和 /process 用 run_in_threadpool 是同一个理由）。所以让它在工作线程里跑，
    通过队列把结果交回事件循环 —— 队列是这两个世界之间唯一的通道。

    队列有界（maxsize=32）：万一客户端读得慢，节点又一直产出，无界队列会一直吃内存。
    """
    queue: asyncio.Queue = asyncio.Queue(maxsize=32)
    loop = asyncio.get_running_loop()

    def producer():
        """在工作线程里跑图，每个节点产出一条消息。

        这个函数整体在 run_in_threadpool 的线程里执行，不能直接 await，
        所以往队列里放东西要用 run_coroutine_threadsafe 跨回事件循环。
        """

        def put(item):
            # .result() 是必须的：它让这里
            # 在队列满时阻塞住，形成反压。
            # 不等的话，慢客户端场景下 put 会无声堆积，maxsize 就白设了。
            asyncio.run_coroutine_threadsafe(queue.put(item), loop).result()

        try:
            graph = get_graph()
            config = RunnableConfig(configurable={"thread_id": user_id})
            previous: dict = {}

            for chunk in graph.stream(
                {"input": user_input}, config, stream_mode="updates"
            ):
                # chunk 形如 {节点名: 该节点返回的状态}
                for node_name, node_state in chunk.items():
                    if not isinstance(node_state, dict):
                        continue
                    delta = _state_delta(previous, node_state)
                    previous = {**previous, **node_state}
                    # 没有新增内容的节点（比如纯路由）不推空消息，前端不用过滤
                    if delta:
                        put({"type": "thinking", "node": node_name, "data": delta})

            # 图跑完了，最后一条给最终答案。previous 累积了全程状态，output 在里面
            output = (previous.get("output") or "").strip() or ERROR_OUTPUT
            put({"type": "final", "input": user_input, "output": output})
        except Exception:
            # 堆栈只打在服务端，返回给调用方的是固定文案，和 /process 保持一致
            traceback.print_exc()
            put({"type": "error", "input": user_input, "output": ERROR_OUTPUT})
        finally:
            # 无论成功失败都要放结束标记，否则消费者会一直等下去，请求挂住
            put(_STREAM_DONE)

    # 不 await：让它在后台线程跑，下面一边跑一边从队列取
    task = asyncio.create_task(run_in_threadpool(producer))
    try:
        while True:
            message = await queue.get()
            if message is _STREAM_DONE:
                break
            # 换行分隔的 JSON（不是严格 SSE），对应文档里的响应格式。
            # ensure_ascii=False 保留中文原样，前端按行 json.loads 即可
            yield json.dumps(message, ensure_ascii=False) + "\n"
    finally:
        # 客户端提前断开时，生成器会被关掉。这里把线程任务收干净，
        # 不然异常会以 "Task exception was never retrieved" 的形式漏出来
        if not task.done():
            task.cancel()
        try:
            await task
        except (asyncio.CancelledError, Exception):
            pass


@app.post("/process_stream", summary="中医问答接口（流式）")
async def process_stream(request: ProcessRequest) -> StreamingResponse:
    """和 /process 同样的入参，但每个节点执行完就先返回一条"思考过程"。

    一次问答要 5-30 秒，一次性接口在这段时间里什么都没有，用户只能干等；
    流式把中间过程吐出来，前端可以显示"正在抽取实体""正在查图谱"。

    响应体不是严格 SSE，是换行分隔的 JSON（text/plain），见接口文档。
    这里不套 try：异常在 producer 内部就收敛成 error 消息了，
    HTTP 状态码已经随第一个字节发出去，改不了了。
    """
    print(f"[/process_stream] user_id={request.user_id} input={request.input!r}")
    return StreamingResponse(
        _stream_messages(request.input, request.user_id),
        media_type="text/plain; charset=utf-8",
        # 关掉 nginx 一类反向代理的缓冲，否则它会攒够一批再转发，流式就没了
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8001)