# -*- coding: utf-8 -*-
"""开放问答接口测试脚本(流式版,Jupyter / 命令行均可)。

接口只返回「最终结果」:智能体把调研 / 工具调用全部跑完后,只把最后那段
文本(markdown)放在 answer 里返回 —— 不含中间过程、工具调用、思考。

为什么默认用流式(/api/open/chat/stream):
  - 长任务(联网调研、出报告)同步阻塞那版整段没有下行字节,容易被前置 nginx 反代
    的 proxy_read_timeout(默认 60s)掐成 504。流式版每隔几秒就有心跳/增量下行,
    代理不会判定空闲超时,**多久都不会超时**。
  - 客户端只需读到 `event: result` 解析 answer 即可,中途 delta/progress/心跳都忽略。

用法:改下面「配置」段,然后整段运行;answer 变量就是最终 md 文本。
无第三方依赖(纯标准库 urllib)。
"""

import json
import time
import urllib.error
import urllib.request

# ====================== 配置(按需修改这里) ======================
BASE_URL = "http://47.94.209.59:2026"   # 走 nginx 用 :2026;Gateway 直连用 :8001
STREAM = True                            # True=流式(推荐,不超时);False=同步阻塞(短问题用)

MESSAGE = "帮我调研一下 赖清德,输出 markdown 报告。"   # 要问的问题
MODEL_NAME = "deepseek-ai/DeepSeek-V4-Flash"          # 模型名(config.yaml models[].name);None=默认模型
AGENT_NAME = "你好君"                     # 智能体中文名称;None=默认 lead agent。不确定填啥先跑 ping_agents()
THINKING_ENABLED = False                 # False=关闭思考(SiliconFlow/内网模型也会被强制关掉);True=开启
SHOW_LIVE = True                         # 流式时是否实时打印增量正文(看着它一点点出),只看结果可设 False
READ_TIMEOUT = 120                       # 单次 socket 读超时(秒)。服务端每 ~15s 发心跳,故 120s 绰绰有余
TIMEOUT = 6000                           # 同步(非流式)模式的整体超时(秒)
# ===============================================================


def open_chat(message, model_name=None, agent_name=None, thinking_enabled=False,
              base_url=BASE_URL, timeout=TIMEOUT):
    """同步阻塞版:整段跑完才返回。短问题用;长任务请用 open_chat_stream。"""
    url = base_url.rstrip("/") + "/api/open/chat"
    payload = {
        "message": message,
        "model_name": model_name,
        "agent_name": agent_name,   # 按中文名称选智能体
        "thinking_enabled": thinking_enabled,
    }
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url, data=data,
        headers={"Content-Type": "application/json; charset=utf-8"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        result = json.loads(resp.read().decode("utf-8"))
    return result["answer"]   # 只返回最终结果


def open_chat_stream(message, model_name=None, agent_name=None, thinking_enabled=False,
                     base_url=BASE_URL, read_timeout=READ_TIMEOUT, show_live=SHOW_LIVE):
    """流式版:读 SSE,边跑边收心跳/增量(不超时),最终从 event: result 取 answer。"""
    url = base_url.rstrip("/") + "/api/open/chat/stream"
    payload = {
        "message": message,
        "model_name": model_name,
        "agent_name": agent_name,
        "thinking_enabled": thinking_enabled,
    }
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url, data=data,
        headers={"Content-Type": "application/json; charset=utf-8", "Accept": "text/event-stream"},
        method="POST",
    )

    answer = None
    error = None
    event_name = None
    data_lines = []

    with urllib.request.urlopen(req, timeout=read_timeout) as resp:
        for raw in resp:                       # 按行读取,服务端心跳让读取永不空闲超时
            line = raw.decode("utf-8").rstrip("\r\n")
            if line.startswith(":"):           # SSE 注释行(心跳),忽略
                continue
            if line == "":                     # 空行 = 一个事件结束,分发
                if data_lines:
                    payload_str = "\n".join(data_lines)
                    try:
                        obj = json.loads(payload_str)
                    except json.JSONDecodeError:
                        obj = {}
                    if event_name == "result":
                        answer = obj.get("answer")
                    elif event_name == "error":
                        error = obj.get("detail") or payload_str
                    elif event_name == "delta" and show_live:
                        print(obj.get("text", ""), end="", flush=True)
                event_name, data_lines = None, []
                continue
            if line.startswith("event:"):
                event_name = line[len("event:"):].strip()
            elif line.startswith("data:"):
                data_lines.append(line[len("data:"):].lstrip())

    if error:
        raise RuntimeError(f"接口返回错误:{error}")
    if answer is None:
        raise RuntimeError("未收到 event: result(流可能被中途断开)。")
    return answer


def ping_agents(base_url=BASE_URL, timeout=30):
    """探活:GET /api/open/agents。秒回说明链路 OK,顺带看有哪些 agent_name 可填。"""
    with urllib.request.urlopen(base_url.rstrip("/") + "/api/open/agents", timeout=timeout) as r:
        return json.loads(r.read().decode("utf-8"))


# ====================== 调用 ======================
if __name__ == "__main__":
    # 先探活(失败说明是链路/部署问题,而非模型问题)
    try:
        info = ping_agents()
        print(f"[探活OK] 可用智能体 {info.get('count', 0)} 个\n")
    except Exception as exc:                    # noqa: BLE001
        print(f"[探活失败] {exc!r} —— 先排查 BASE_URL / 端口 / 服务是否在跑\n")

    _t0 = time.time()
    if STREAM:
        answer = open_chat_stream(MESSAGE, MODEL_NAME, AGENT_NAME, THINKING_ENABLED)
    else:
        answer = open_chat(MESSAGE, MODEL_NAME, AGENT_NAME, THINKING_ENABLED)
    print(f"\n\n(耗时 {time.time() - _t0:.1f}s)\n")
    print(answer)        # answer 就是最终 markdown 文本

    # ============ (可选)Jupyter 里渲染 markdown ============
    # from IPython.display import Markdown, display
    # display(Markdown(answer))

    # ============ (可选)存成 .md ============
    # with open("result.md", "w", encoding="utf-8") as f:
    #     f.write(answer)