"""用 requests 手工测 POST /process_stream。

和 test_api.py 的区别：test_api.py 默认拿 FakeGraph 跑离线断言，这个脚本是**打真实服务**，
用来肉眼看流式效果 —— 每个节点的思考过程是不是一条条到达，而不是最后一次性糊过来。

前提：先在仓库根目录起服务（main.py 里监听的是 8001，不是文档里的 8000）

    python -X utf8 -m _005_fastapi.main

运行：

    python -X utf8 -m _005_fastapi.my_test_api_stream
    python -X utf8 -m _005_fastapi.my_test_api_stream --base-url http://127.0.0.1:8000
"""

import argparse
import json
import time

import requests

BASE_URL = "http://127.0.0.1:8001"

# 两条问题走两条不同的路：第一条是闲聊，图会直接跳到兜底节点；
# 第二条要查图谱，节点多，流式的效果才看得出来
QUESTIONS = [
    "你好",
    "你好,我气虚应该吃人参吗?",
]


def query_stream(question: str, base_url: str = BASE_URL, user_id: str = "user_id_1") -> str:
    """打一次流式接口，边收边打印，返回最终答案。

    stream=True 是关键：不加的话 requests 会把整个响应体读完才返回，
    客户端这边就退化成一次性接口，测不出流式。
    iter_lines 正好匹配服务端"换行分隔 JSON"的格式。
    """
    url = f"{base_url}/process_stream"
    payload = {"input": question, "user_id": user_id}

    start = time.time()
    first_at = None
    final_output = ""

    with requests.post(url, json=payload, stream=True, timeout=180) as res:
        print(f"HTTP {res.status_code}  Content-Type: {res.headers.get('Content-Type')}")
        res.raise_for_status()

        for line in res.iter_lines(decode_unicode=True):
            # 空行是心跳/分隔，跳过
            if not line or not line.strip():
                continue

            elapsed = time.time() - start
            if first_at is None:
                first_at = elapsed

            message = json.loads(line)
            kind = message.get("type")

            if kind == "thinking":
                # data 是这个节点这一步产生的状态增量，可能挺长，截断了看
                data = json.dumps(message.get("data"), ensure_ascii=False)
                if len(data) > 160:
                    data = data[:160] + "..."
                print(f"  [{elapsed:6.2f}s] 思考中 · {message.get('node')}: {data}")
            elif kind == "final":
                final_output = message.get("output", "")
                print(f"  [{elapsed:6.2f}s] 最终答案: {final_output}")
            elif kind == "error":
                final_output = message.get("output", "")
                print(f"  [{elapsed:6.2f}s] 出错: {final_output}")
            else:
                print(f"  [{elapsed:6.2f}s] 未知消息类型: {message}")

    total = time.time() - start
    # 首条消息明显早于总耗时，才说明是真流式；两个数字贴在一起就是被缓冲了
    print(f"  首条 {first_at if first_at is None else round(first_at, 2)}s / 总计 {total:.2f}s")
    return final_output


def main() -> None:
    parser = argparse.ArgumentParser(description="requests 手测 /process_stream")
    parser.add_argument("--base-url", default=BASE_URL, help=f"服务地址，默认 {BASE_URL}")
    parser.add_argument("--input", default=None, help="只测这一个问题")
    args = parser.parse_args()

    questions = [args.input] if args.input else QUESTIONS

    for question in questions:
        print(f"\n=== 提问: {question} ===")
        try:
            query_stream(question, base_url=args.base_url)
        except requests.RequestException as exc:
            print(f"  请求失败: {exc}（服务起了吗？python -X utf8 -m _005_fastapi.main）")

    # 顺手验一下入参校验：流式接口和 /process 共用 ProcessRequest，空 input 应该 422
    print("\n=== 入参校验: 空 input ===")
    try:
        res = requests.post(
            f"{args.base_url}/process_stream", json={"input": "", "user_id": "u"}, timeout=10
        )
        print(f"HTTP {res.status_code}（期望 422） {res.text[:120]}")
    except requests.RequestException as exc:
        print(f"  请求失败: {exc}")


if __name__ == "__main__":
    main()