"""_005_fastapi 的测试用例

仓库里没有 pytest，测试沿用各模块 `if __name__ == "__main__"` 自测的约定，
用一个极小的断言 + 计数框架，跑完打印 通过/失败 汇总，失败时退出码非 0。

两种模式：

    python -X utf8 -m _005_fastapi.test_api          # 离线（默认）
    python -X utf8 -m _005_fastapi.test_api --live   # 打真实服务

**离线模式是主力**：它把 `main._graph` 换成一张假图，于是不碰 LLM、Neo4j、FAISS，
一秒内跑完，验证的是接口层自己的行为 —— 状态码、字段、校验、异常收敛。
真图跑一遍要 5-30 秒还要花钱，不适合当回归测试。

`--live` 是对着 `python -m _005_fastapi.main` 起好的服务打真实请求，验证的是端到端
真的通（图能加载、答案有内容）。它依赖外部服务，所以是可选的，不进默认流程。
"""

import argparse
import json
import sys

from fastapi.testclient import TestClient

from _005_fastapi import main

BASE_URL = "http://localhost:8000"

# 结果计数，_check 往里累加
passed = 0
failed = 0


def _check(condition: bool, label: str, detail: str = "") -> None:
    """一条断言 = 一行输出。不抛异常，让后面的用例继续跑完"""
    global passed, failed
    if condition:
        passed += 1
        print(f"  [PASS] {label}")
    else:
        failed += 1
        print(f"  [FAIL] {label}" + (f" -> {detail}" if detail else ""))


# ======================
# 假图：替掉真图，测接口层自己
# ======================
class FakeGraph:
    """只实现 invoke，返回预先设定的 state 或抛异常。

    真图的接口就这一个方法（`graph.invoke(dict) -> dict`），所以替换成本极低。
    """

    def __init__(self, state=None, error=None):
        self.state = state
        self.error = error
        self.calls = []
        self.configs = []

    def invoke(self, initial_state, config=None):
        """签名要跟真图对齐。

        加了 checkpointer 之后 main.py 会带上 `config`（里面是 thread_id），
        少这个参数会让每次调用都 TypeError，然后被 /process 的 except 收敛成 500 ——
        测出来是"图挂了"，看不出其实是假图的签名过时了。
        """
        self.calls.append(initial_state)
        self.configs.append(config)
        if self.error is not None:
            raise self.error
        return self.state


def _client(fake_graph) -> TestClient:
    """装上假图，返回 TestClient。

    直接给 `main._graph` 赋值即可绕过懒加载 —— `get_graph()` 见到非 None 就不再 import
    真图。TestClient 会触发 lifespan，预热那一步同样命中缓存，不会去加载 FAISS。
    """
    main._graph = fake_graph
    return TestClient(main.app)


# ======================
# 用例 1：健康检查
# ======================
def test_health_check() -> None:
    print("\n[1] GET / 健康检查")
    with _client(FakeGraph(state={"output": "不该被调用"})) as client:
        response = client.get("/")

    _check(response.status_code == 200, "状态码 200", str(response.status_code))
    body = response.json()
    _check(body.get("status") == "ok", "status == ok", repr(body.get("status")))
    _check(
        body.get("message") == "中医知识图谱问答系统服务运行正常",
        "message 与文档一致",
        repr(body.get("message")),
    )
    _check(set(body) == {"status", "message"}, "字段只有 status/message", str(set(body)))


# ======================
# 用例 2：/process 正常返回
# ======================
def test_process_ok() -> None:
    print("\n[2] POST /process 正常问答（图谱路）")
    answer = "人参具有大补元气、复脉固脱、补脾益肺、生津养血、安神益智等功效。"
    graph = FakeGraph(state={"input": "人参有什么功效？", "output": answer})

    with _client(graph) as client:
        response = client.post(
            "/process", json={"input": "人参有什么功效？", "user_id": "user-123"}
        )

    _check(response.status_code == 200, "状态码 200", str(response.status_code))
    body = response.json()
    _check(body.get("input") == "人参有什么功效？", "input 原样回显", repr(body.get("input")))
    _check(body.get("output") == answer, "output 取自 state['output']", repr(body.get("output")))
    _check(set(body) == {"input", "output"}, "字段只有 input/output", str(set(body)))
    # 只把 input 传进图，user_id 不进 state（AgentState 里没有这个字段）
    _check(graph.calls == [{"input": "人参有什么功效？"}], "只把 input 传给图", str(graph.calls))


# ======================
# 用例 3：/process 兜底路（非中医提问）
# ======================
def test_process_direct_out() -> None:
    print("\n[3] POST /process 兜底路（非中医提问）")
    answer = "抱歉，我无法提供天气信息。我是一个中医知识图谱问答系统。"
    # 兜底路的 state 里没有 neo4j_answer，但 output 一样被写满，接口层不用区分
    graph = FakeGraph(state={"input": "今天天气怎么样？", "direct_out": answer, "output": answer})

    with _client(graph) as client:
        response = client.post(
            "/process", json={"input": "今天天气怎么样？", "user_id": "user-456"}
        )

    _check(response.status_code == 200, "状态码 200", str(response.status_code))
    _check(response.json().get("output") == answer, "兜底路答案也能取到", repr(response.json()))


# ======================
# 用例 4：入参校验
# ======================
def test_process_validation() -> None:
    print("\n[4] POST /process 入参校验（FastAPI 返回 422）")
    graph = FakeGraph(state={"output": "不该被调用"})

    bad_requests = [
        ({"user_id": "user-123"}, "缺 input"),
        ({"input": "人参有什么功效？"}, "缺 user_id"),
        ({"input": "", "user_id": "user-123"}, "input 为空串"),
        ({"input": "   ", "user_id": "user-123"}, "input 全空格"),
        ({"input": "人参有什么功效？", "user_id": ""}, "user_id 为空串"),
        ({}, "两个字段都缺"),
    ]

    with _client(graph) as client:
        for payload, label in bad_requests:
            response = client.post("/process", json=payload)
            _check(response.status_code == 422, f"{label} -> 422", str(response.status_code))

    # 参数不合法就该被挡在图之前，一次大模型都不该调
    _check(graph.calls == [], "非法请求没有触达图", str(graph.calls))


# ======================
# 用例 5：异常收敛成 500 + 固定文案
# ======================
def test_process_error() -> None:
    print("\n[5] POST /process 图执行异常")
    graph = FakeGraph(error=RuntimeError("Neo4j 连接失败"))

    with _client(graph) as client:
        response = client.post(
            "/process", json={"input": "人参有什么功效？", "user_id": "user-123"}
        )

    _check(response.status_code == 500, "状态码 500", str(response.status_code))
    body = response.json()
    _check(body.get("input") == "", "错误响应 input 为空串", repr(body.get("input")))
    _check(body.get("output") == main.ERROR_OUTPUT, "output 为固定错误文案", repr(body.get("output")))
    # 内部异常信息不该漏给调用方
    _check("Neo4j" not in response.text, "不泄漏内部异常细节", response.text)


# ======================
# 用例 6：图没写 output 时的兜底
# ======================
def test_process_empty_output() -> None:
    print("\n[6] POST /process 图返回空 output")
    for state, label in [
        ({"input": "x"}, "state 里没有 output 键"),
        ({"input": "x", "output": None}, "output 为 None"),
        ({"input": "x", "output": "   "}, "output 全空白"),
    ]:
        with _client(FakeGraph(state=state)) as client:
            response = client.post("/process", json={"input": "x", "user_id": "u"})
        _check(
            response.status_code == 200
            and response.json().get("output") == main.ERROR_OUTPUT,
            f"{label} -> 200 + 错误文案",
            response.text,
        )


# ======================
# 用例 7：并发不被阻塞
# ======================
def test_process_concurrent() -> None:
    """图是同步阻塞的，接口必须把它丢进线程池，否则一个慢请求会卡住整个服务。

    这里让假图 sleep，再用线程并发打两个请求：串行执行耗时会是 2 倍 sleep，
    并发则接近 1 倍。用 1.5 倍做判据，够区分两种情况又不至于被机器抖动误判。
    """
    print("\n[7] POST /process 并发（验证没有阻塞事件循环）")
    import time
    from concurrent.futures import ThreadPoolExecutor

    delay = 0.5

    class SlowGraph(FakeGraph):
        # config 同样要接（见 FakeGraph.invoke）
        def invoke(self, initial_state, config=None):
            time.sleep(delay)
            return {"output": "慢回答"}

    with _client(SlowGraph()) as client:
        def one(_):
            return client.post("/process", json={"input": "人参", "user_id": "u"})

        start = time.perf_counter()
        with ThreadPoolExecutor(max_workers=2) as pool:
            responses = list(pool.map(one, range(2)))
        elapsed = time.perf_counter() - start

    _check(all(r.status_code == 200 for r in responses), "两个请求都成功")
    _check(elapsed < delay * 1.5, f"并发执行（耗时 {elapsed:.2f}s < {delay * 1.5}s）", f"{elapsed:.2f}s")


# ======================
# 用例 8：流式接口
# ======================
class FakeStreamGraph(FakeGraph):
    """假图的流式版本：按给定顺序吐出 {节点名: 状态}。

    真图的 stream 是"每个节点执行完产出一条"，而且因为我们的节点是改完整个 state
    再返回，产出的是**全量状态**不是增量。这里照着这个形状造数据，才测得到
    main._state_delta 做差那一步。
    """

    def __init__(self, chunks=None, error=None):
        super().__init__(state=None, error=error)
        self.chunks = chunks or []

    def stream(self, initial_state, config=None, stream_mode=None):
        self.calls.append(initial_state)
        self.configs.append(config)
        if self.error is not None:
            raise self.error
        for chunk in self.chunks:
            yield chunk


def _read_stream(response) -> list:
    """把换行分隔的响应体解析成消息列表"""
    return [json.loads(line) for line in response.text.splitlines() if line.strip()]


def test_process_stream_ok() -> None:
    """逐节点产出 thinking，最后一条 final，且 data 是增量不是全量"""
    print("\n[8] POST /process_stream 正常流式问答")

    # 模拟三个节点依次执行，状态逐步累积（真图就是这个形状）
    chunks = [
        {"语义转写节点": {"input": "人参有什么功效？", "input_semantic_trans": "人参有什么功效？"}},
        {
            "中医意图识别节点": {
                "input": "人参有什么功效？",
                "input_semantic_trans": "人参有什么功效？",
                "is_zhongyi_intent": True,
            }
        },
        {
            "图谱答案节点": {
                "input": "人参有什么功效？",
                "input_semantic_trans": "人参有什么功效？",
                "is_zhongyi_intent": True,
                "neo4j_answer": "人参能大补元气。",
                "output": "人参能大补元气。",
            }
        },
    ]

    with _client(FakeStreamGraph(chunks)) as client:
        response = client.post(
            "/process_stream", json={"input": "人参有什么功效？", "user_id": "user-123"}
        )

    _check(response.status_code == 200, "状态码 200", str(response.status_code))
    messages = _read_stream(response)

    thinking = [m for m in messages if m["type"] == "thinking"]
    finals = [m for m in messages if m["type"] == "final"]

    _check(len(finals) == 1, "有且只有一条 final", f"{len(finals)} 条")
    _check(messages[-1]["type"] == "final", "final 是最后一条", messages[-1]["type"])
    _check(len(thinking) >= 2, f"有多条 thinking（{len(thinking)} 条）", str(len(thinking)))

    # 关键：data 必须是这一步新增的键，不能把全量状态原样推出去
    intent_msg = next((m for m in thinking if m["node"] == "中医意图识别节点"), None)
    _check(intent_msg is not None, "能找到意图识别节点的消息")
    if intent_msg:
        _check(
            intent_msg["data"] == {"is_zhongyi_intent": True},
            "data 只含这一步的增量",
            repr(intent_msg["data"]),
        )

    _check(
        finals[0]["output"] == "人参能大补元气。", "final 取到 output", repr(finals[0]["output"])
    )
    _check(finals[0]["input"] == "人参有什么功效？", "final 回显 input", repr(finals[0]["input"]))

    # thinking 里不该出现 output/history_messages 这类噪音键
    noisy = [k for m in thinking for k in m["data"] if k in main.STREAM_SKIP_KEYS]
    _check(not noisy, "thinking 不含被过滤的键", repr(noisy))


def test_process_stream_error() -> None:
    """图抛异常时，流以 error 消息收尾，而不是把堆栈漏出去或让请求挂住"""
    print("\n[9] POST /process_stream 图执行异常")

    with _client(FakeStreamGraph(error=RuntimeError("Neo4j 连接失败"))) as client:
        response = client.post(
            "/process_stream", json={"input": "人参", "user_id": "u"}
        )

    messages = _read_stream(response)
    _check(response.status_code == 200, "状态码 200（头已发出，改不了）", str(response.status_code))
    _check(len(messages) == 1, "只有一条消息", f"{len(messages)} 条")
    _check(messages[-1]["type"] == "error", "最后一条是 error", messages[-1]["type"])
    _check(
        messages[-1]["output"] == main.ERROR_OUTPUT, "错误文案固定", repr(messages[-1]["output"])
    )
    _check("Neo4j" not in response.text, "堆栈/内部信息没有漏出去")


def test_process_stream_validation() -> None:
    """入参校验和 /process 共用同一个模型，行为要一致"""
    print("\n[10] POST /process_stream 入参校验")

    with _client(FakeStreamGraph([])) as client:
        _check(
            client.post("/process_stream", json={"input": "", "user_id": "u"}).status_code
            == 422,
            "空 input -> 422",
        )
        _check(
            client.post("/process_stream", json={"input": "人参"}).status_code == 422,
            "缺 user_id -> 422",
        )
        _check(
            client.post("/process_stream", json={"input": "   ", "user_id": "u"}).status_code
            == 422,
            "全空格 input -> 422",
        )


def test_process_stream_thread_id() -> None:
    """user_id 要作为 thread_id 传给图，否则多轮记忆按用户隔离就失效了"""
    print("\n[11] POST /process_stream 透传 user_id 作为 thread_id")

    fake = FakeStreamGraph([{"节点": {"output": "答案"}}])
    with _client(fake) as client:
        client.post("/process_stream", json={"input": "人参", "user_id": "user-abc"})

    _check(fake.calls == [{"input": "人参"}], "只把 input 传给图", repr(fake.calls))
    thread_id = (fake.configs[0] or {}).get("configurable", {}).get("thread_id")
    _check(thread_id == "user-abc", "thread_id == user_id", repr(thread_id))


# ======================
# --live：打真实服务
# ======================
LIVE_QUESTIONS = [
    # 图谱路：药材属性
    "人参有什么功效？",
    # 图谱路：方剂 → 组成
    "四君子汤由哪些药材组成？",
    # 兜底路：非中医提问
    "帮我写一段 Python 快排代码",
]


def run_live(base_url: str) -> None:
    """对着已启动的服务打真实请求。需要先 python -m _005_fastapi.main"""
    import requests

    print(f"\n=== live 模式：{base_url} ===")

    print("\n[live-1] GET /")
    try:
        response = requests.get(f"{base_url}/", timeout=10)
    except requests.RequestException as exc:
        _check(False, "连接服务", f"{exc}（服务没启动？先跑 python -m _005_fastapi.main）")
        return
    _check(response.status_code == 200, "状态码 200", str(response.status_code))
    _check(response.json().get("status") == "ok", "status == ok", response.text)

    for i, question in enumerate(LIVE_QUESTIONS, start=2):
        print(f"\n[live-{i}] POST /process {question!r}")
        try:
            # 文档建议 120 秒超时：单次问答 5-30 秒，复杂问题更久
            response = requests.post(
                f"{base_url}/process",
                json={"input": question, "user_id": "test-user-001"},
                timeout=120,
            )
        except requests.RequestException as exc:
            _check(False, "请求成功", str(exc))
            continue

        _check(response.status_code == 200, "状态码 200", str(response.status_code))
        body = response.json()
        _check(body.get("input") == question, "input 回显", repr(body.get("input")))
        answer = body.get("output") or ""
        _check(bool(answer.strip()), "output 非空", repr(answer))
        _check(answer != main.ERROR_OUTPUT, "不是错误文案", repr(answer))
        print(f"    回答: {answer[:120]}")

    print("\n[live-last] POST /process 入参校验")
    response = requests.post(f"{base_url}/process", json={"input": ""}, timeout=10)
    _check(response.status_code == 422, "空 input + 缺 user_id -> 422", str(response.status_code))


OFFLINE_TESTS = [
    test_health_check,
    test_process_ok,
    test_process_direct_out,
    test_process_validation,
    test_process_error,
    test_process_empty_output,
    test_process_concurrent,
    test_process_stream_ok,
    test_process_stream_error,
    test_process_stream_validation,
    test_process_stream_thread_id,
]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="_005_fastapi 接口测试")
    parser.add_argument(
        "--live",
        action="store_true",
        help="对已启动的服务打真实请求（需要 LLM / Neo4j / FAISS 都可用）",
    )
    parser.add_argument("--base-url", default=BASE_URL, help=f"live 模式的服务地址，默认 {BASE_URL}")
    args = parser.parse_args()

    if args.live:
        run_live(args.base_url.rstrip("/"))
    else:
        print("=== 离线模式：用假图替掉真图，不调用 LLM / Neo4j / FAISS ===")
        for test in OFFLINE_TESTS:
            test()

    print(f"\n{'=' * 70}")
    print(f"通过 {passed} 条，失败 {failed} 条")
    sys.exit(1 if failed else 0)