#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""test_openai_api.py — openai_api.py 的 mock 全端点联调 (本地可跑, 无需 GPU/transformers)
覆盖: /health /v1/models /v1/chat/completions (非流式 bare / think_off / 模型名后缀 /
      raw_content / 流式SSE含usage) /v1/completions / 错误分支 (401/400)
"""
import json
import subprocess
import sys
import time
import urllib.request
import urllib.error

PORT = 8399
BASE = f"http://127.0.0.1:{PORT}"
PY = sys.executable


def req(method, path, body=None, headers=None, timeout=15):
    data = json.dumps(body).encode() if body is not None else None
    r = urllib.request.Request(BASE + path, data=data, method=method,
                               headers={"Content-Type": "application/json",
                                        **(headers or {})})
    try:
        with urllib.request.urlopen(r, timeout=timeout) as resp:
            return resp.status, resp.read().decode()
    except urllib.error.HTTPError as e:
        return e.code, e.read().decode()


def raw_stream(body):
    r = urllib.request.Request(BASE + "/v1/chat/completions",
                               data=json.dumps(body).encode(), method="POST",
                               headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(r, timeout=20) as resp:
        return resp.read().decode()


fails = []


def check(name, cond, detail=""):
    print(("PASS " if cond else "FAIL ") + name + ("" if cond else f"  <- {detail}"))
    if not cond:
        fails.append(name)


def main():
    proc = subprocess.Popen([PY, "/home/z/my-project/scripts/openai_api.py",
                             "--mock", "--port", str(PORT)],
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        for _ in range(40):
            try:
                s, b = req("GET", "/health")
                if s == 200 and json.loads(b).get("loaded"):
                    break
            except Exception:
                pass
            time.sleep(0.25)
        else:
            print("server not up")
            sys.exit(1)

        s, b = req("GET", "/health")
        h = json.loads(b)
        check("health.loaded", s == 200 and h["mock"] is True and h["loaded"] is True, b[:120])

        s, b = req("GET", "/v1/models")
        m = json.loads(b)
        check("models.list", s == 200 and m["object"] == "list"
              and any(x["id"] == "samai-2b" for x in m["data"]), b[:120])

        # bare 协议 (默认): content 剥 think + reasoning_content + raw
        s, b = req("POST", "/v1/chat/completions", {
            "model": "samai-2b",
            "messages": [{"role": "user", "content": "你好"}],
            "max_tokens": 100})
        d = json.loads(b)
        ch = d["choices"][0]
        check("chat.bare.status", s == 200 and d["object"] == "chat.completion", b[:200])
        check("chat.bare.content_stripped", "<think>" not in ch["message"]["content"]
              and ch["message"]["content"].startswith("[mock]"), ch["message"]["content"][:80])
        check("chat.bare.reasoning", "reasoning_content" in ch["message"]
              and "mock" in ch["message"]["reasoning_content"], str(ch["message"])[:150])
        check("chat.bare.raw_kept", ch["message"]["raw"].startswith("<think>"), "")
        check("chat.bare.usage", d["usage"]["total_tokens"] ==
              d["usage"]["prompt_tokens"] + d["usage"]["completion_tokens"], str(d["usage"]))
        check("chat.bare.finish", ch["finish_reason"] == "stop", ch["finish_reason"])
        check("chat.bare.ponder", isinstance(d.get("ponder"), dict)
              and "steps_mean" in d["ponder"], str(d.get("ponder")))
        check("chat.bare.ids", d["id"].startswith("chatcmpl-") and d["model"] == "samai-2b", d["id"])

        # think_off via enable_thinking=false
        s, b = req("POST", "/v1/chat/completions", {
            "model": "samai-2b", "enable_thinking": False,
            "messages": [{"role": "user", "content": "12x8?"}]})
        d = json.loads(b)
        check("chat.thinkoff.no_think", "<think>" not in d["choices"][0]["message"]["content"]
              and "reasoning_content" not in d["choices"][0]["message"], b[:200])

        # think_off via chat_template_kwargs
        s, b = req("POST", "/v1/chat/completions", {
            "model": "samai-2b", "chat_template_kwargs": {"enable_thinking": False},
            "messages": [{"role": "user", "content": "12x8?"}]})
        d = json.loads(b)
        check("chat.ctk.thinkoff", "<think>" not in d["choices"][0]["message"]["content"], b[:150])

        # 模型名后缀 -think-off
        s, b = req("POST", "/v1/chat/completions", {
            "model": "samai-2b-think-off",
            "messages": [{"role": "user", "content": "hi"}]})
        d = json.loads(b)
        check("chat.suffix.thinkoff", "<think>" not in d["choices"][0]["message"]["content"]
              and d["model"] == "samai-2b", b[:150])

        # raw_content=true → content=原始全文
        s, b = req("POST", "/v1/chat/completions", {
            "model": "samai-2b", "raw_content": True,
            "messages": [{"role": "user", "content": "hi"}]})
        d = json.loads(b)
        check("chat.rawcontent", d["choices"][0]["message"]["content"].startswith("<think>"), b[:150])

        # 流式 SSE
        body = {"model": "samai-2b",
                "messages": [{"role": "user", "content": "流式测试"}],
                "stream": True, "stream_options": {"include_usage": True}}
        sse = raw_stream(body)
        lines = [l for l in sse.split("\n\n") if l.startswith("data: ")]
        objs = [json.loads(l[6:]) for l in lines if l[6:] != "[DONE]"]
        ids = {o["id"] for o in objs if o["choices"]}
        contents = "".join((o["choices"][0]["delta"].get("content") or "")
                           for o in objs if o["choices"])
        finishes = [o["choices"][0]["finish_reason"] for o in objs
                    if o["choices"] and o["choices"][0]["finish_reason"]]
        has_usage = any(o.get("usage") for o in objs)
        has_done = "[DONE]" in sse
        check("stream.ids_consistent", len(ids) == 1, str(ids))
        check("stream.content_no_think", "<think>" not in contents
              and contents.startswith("[mock]"), contents[:80])
        check("stream.reasoning_streamed", any((o["choices"][0]["delta"].get("reasoning_content"))
                                               for o in objs if o["choices"]), "")
        check("stream.finish_and_done", finishes == ["stop"] and has_done, f"{finishes} done={has_done}")
        check("stream.usage", has_usage, "")

        # /v1/completions
        s, b = req("POST", "/v1/completions", {"model": "samai-2b", "prompt": "续写测试", "max_tokens": 32})
        d = json.loads(b)
        check("completions.basic", s == 200 and "text" in d["choices"][0]
              and d["usage"]["completion_tokens"] > 0, b[:150])

        # 错误分支
        s, b = req("POST", "/v1/chat/completions", {"messages": []})
        check("err.empty_messages", s == 400 and "error" in json.loads(b), b[:120])
        s, b = req("POST", "/v1/chat/completions", {"model": "x"})
        check("err.bad_json_shape", s == 400, b[:120])

        print("\n%s  (fails=%d)" % ("ALL_PASS" if not fails else "FAILED", len(fails)))
        sys.exit(0 if not fails else 1)
    finally:
        proc.terminate()


if __name__ == "__main__":
    main()