#!/usr/bin/env python3
"""SSE termination probe for DeepSeek-family chat endpoints.

Detects streams that end without the two terminal frames the DeepSeek streaming
contract requires: (a) a final chunk with non-null choices[0].finish_reason
(carrying the aggregated usage), and (b) the `data: [DONE]` sentinel.

Usage:
    DS_API_KEY=sk-... OPENCODE_GO_API_KEY=oc_sk-... \
        python3 ds_stream_probe.py --arm both --mode small --count 10 --concurrency 2
    DS_API_KEY=sk-... python3 ds_stream_probe.py --arm official --mode big --count 3
    DS_API_KEY=sk-... python3 ds_stream_probe.py --arm official --mode reuse --count 6

Notes:
    * mode=small -> short prompt, mode=big -> ~1.7 MB uncached prompt,
      mode=reuse -> one keep-alive connection, growing multi-turn conversation.
    * raw response headers + tail bytes of every request are written to --out.
"""
import argparse, http.client, json, os, threading, time

ARMS = {
    "official": dict(host="api.deepseek.com", path="/chat/completions",
                     model="deepseek-flash", key_env="DS_API_KEY", extra={}),
    "gateway": dict(host="opencode.ai", path="/inference/go/v1/chat/completions",
                    model="deepseek-v4.1-flash", key_env="OPENCODE_GO_API_KEY",
                    extra={"x-stainless-timeout": "300"}),
}
ASK = ("Think step by step at length, then answer: compute 37*941, 128*256, "
       "999*999, 1234*5678 and verify each result by a second method.")
FILLER = "\n".join(
    f"src/app/mod{i%71}/file{i}.ts: export const v{i} = {i}; // component {i}"
    for i in range(18000))
TOOLS = [dict(type="function", function=dict(
    name=f"tool_{i}", description="Probe tool. " + "detail " * 30,
    parameters=dict(type="object", required=["path"], properties=dict(
        path=dict(type="string"), mode=dict(type="string", enum=["r", "w"])))))
    for i in range(12)]
STATS, LOCK = [], threading.Lock()


def body_for(mode, model):
    if mode == "small":
        msgs, mx = [dict(role="user", content=ASK)], 8000
    elif mode == "big":
        msgs, mx = [dict(role="user", content=FILLER + "\n\n" + ASK)], 20000
    else:  # reuse
        msgs, mx = [dict(role="system", content=FILLER)], 4000
    return dict(model=model, messages=msgs, stream=True,
                stream_options=dict(include_usage=True),
                reasoning_effort="max", max_tokens=mx, tools=TOOLS)


def open_stream(arm, mode, model, conc_idx=1):
    cfg = ARMS[arm]
    key = os.environ.get(cfg["key_env"], "")
    assert key, f"missing {cfg['key_env']}"
    conn = http.client.HTTPSConnection(cfg["host"], timeout=900)
    headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json",
               "Accept": "text/event-stream", "User-Agent": "omp/18.12.2",
               "x-opencode-session": f"dsprobe-{mode}-{int(time.time())}-{conc_idx}",
               **{k: v for k, v in cfg["extra"].items() if v is not None}}
    payload = json.dumps(body_for(mode, model)).encode()
    conn.request("POST", cfg["path"], body=payload, headers=headers)
    return conn, conn.getresponse()


def read_stream(resp, t0):
    """Consume SSE; return (chunks, finish_reasons, saw_done, usage, ttft, tail_bytes)."""
    chunks, fin, done, usage, ttft, tail = 0, [], False, None, None, bytearray()
    while True:
        line = resp.readline()
        if not line:
            break
        tail.extend(line)
        del tail[:-2000]
        if not line.endswith(b"\n"):          # response cut inside an SSE frame
            return chunks, fin, done, usage, ttft, bytes(tail)
        s = line.strip()
        if not s.startswith(b"data:"):
            continue
        p = s[5:].strip()
        if p == b"[DONE]":
            done = True
            continue
        try:
            j = json.loads(p)
        except Exception:
            continue
        if j.get("usage"):
            usage = j["usage"]
        ch = (j.get("choices") or [{}])[0]
        if ch.get("finish_reason"):
            fin.append(ch["finish_reason"])
        if ch.get("delta"):
            chunks += 1
            if ttft is None:
                ttft = time.time() - t0
    return chunks, fin, done, usage, ttft, bytes(tail)


def probe(arm, mode, idx, out):
    cfg, t0 = ARMS[arm], time.time()
    rec = dict(arm=arm, mode=mode, idx=idx, t0=time.strftime("%H:%M:%S"))
    try:
        conn, resp = open_stream(arm, mode, cfg["model"], idx)
        rec["status"] = resp.status
        rec["headers"] = {k.lower(): v for k, v in resp.getheaders()}
        if resp.status != 200:
            rec["tail"], rec["err"] = resp.read(500), f"HTTP {resp.status}"
        else:
            (rec["chunks"], rec["finish"], rec["done"], rec["usage"],
             rec["ttft"], rec["tail"]) = read_stream(resp, t0)
        conn.close()
    except Exception as e:  # transport level failure
        rec["err"] = f"{type(e).__name__}: {e}"
    rec["dur"] = time.time() - t0
    rec["truncated"] = bool(not rec.get("done") and not rec.get("finish")
                            and rec.get("chunks"))
    with LOCK:
        STATS.append(rec)
        with open(f"{out}/{arm}-{mode}-{idx:02d}.log", "wb") as fh:
            fh.write(f"# {arm} {mode} #{idx} {rec['t0']} status={rec.get('status')}\n".encode())
            fh.write(json.dumps(rec.get("headers", {}), indent=1).encode() + b"\n")
            fh.write(f"# chunks={rec.get('chunks')} ttft={rec.get('ttft')} "
                     f"finish={rec.get('finish')} DONE={rec.get('done')} "
                     f"err={rec.get('err')}\n".encode())
            fh.write(b"# tail (last 2 KiB):\n" + rec.get("tail", b""))
    print(f"[{arm}/{mode} #{idx}] {rec['dur']:6.1f}s ttft={(rec.get('ttft') or -1):5.1f}s "
          f"chunks={rec.get('chunks')} finish={rec.get('finish')} DONE={rec.get('done')} "
          f"truncated={rec['truncated']} err={rec.get('err')}", flush=True)
    return rec


def run_reuse(arm, count, out):
    """Sequential requests over ONE keep-alive connection; the conversation grows."""
    cfg = ARMS[arm]
    conn = http.client.HTTPSConnection(cfg["host"], timeout=900)
    msgs = [dict(role="system", content=FILLER),
            dict(role="user", content=ASK + " Use a tool if you need more data.")]
    for i in range(1, count + 1):
        t0 = time.time()
        rec = dict(arm=arm, mode="reuse", idx=i, t0=time.strftime("%H:%M:%S"))
        try:
            key = os.environ[cfg["key_env"]]
            payload = dict(body_for("reuse", cfg["model"]), messages=msgs)
            conn.request("POST", cfg["path"], body=json.dumps(payload).encode(), headers={
                "Authorization": f"Bearer {key}", "Content-Type": "application/json",
                "Accept": "text/event-stream", "User-Agent": "omp/18.12.2",
                "x-opencode-session": f"dsprobe-reuse-{int(time.time())}-{i}",
                **{k: v for k, v in cfg["extra"].items() if v is not None}})
            resp = conn.getresponse()
            rec["status"] = resp.status
            rec["headers"] = {k.lower(): v for k, v in resp.getheaders()}
            (rec["chunks"], rec["finish"], rec["done"], rec["usage"],
             rec["ttft"], rec["tail"]) = read_stream(resp, t0)
        except Exception as e:
            rec["err"] = f"{type(e).__name__}: {e}"
        rec["dur"] = time.time() - t0
        rec["truncated"] = bool(not rec.get("done") and not rec.get("finish")
                                and rec.get("chunks"))
        with LOCK:
            STATS.append(rec)
            with open(f"{out}/{arm}-reuse-{i:02d}.log", "wb") as fh:
                fh.write(f"# {arm} reuse #{i} {rec['t0']} status={rec.get('status')}\n".encode())
                fh.write(json.dumps(rec.get("headers", {}), indent=1).encode() + b"\n")
                fh.write(f"# chunks={rec.get('chunks')} ttft={rec.get('ttft')} "
                         f"finish={rec.get('finish')} DONE={rec.get('done')} "
                         f"usage={json.dumps(rec.get('usage'))}\n".encode())
                fh.write(b"# tail (last 2 KiB):\n" + rec.get("tail", b""))
        print(f"[{arm}/reuse #{i}] {rec['dur']:6.1f}s ttft={(rec.get('ttft') or -1):5.1f}s "
              f"chunks={rec.get('chunks')} finish={rec.get('finish')} DONE={rec.get('done')} "
              f"truncated={rec['truncated']} err={rec.get('err')}", flush=True)
        if rec.get("done"):
            msgs += [dict(role="assistant", content=f"Turn {i}: noted."),
                     dict(role="user", content=f"Continue: verify item {i} briefly.")]


def worker(arm, mode, queue, out, qlock):
    while True:
        with qlock:
            if not queue:
                return
            idx = queue.pop(0)
        probe(arm, mode, idx, out)
        time.sleep(1)


def run_arm(arm, mode, count, conc, out):
    q, qlock = list(range(1, count + 1)), threading.Lock()
    ts = [threading.Thread(target=worker, args=(arm, mode, q, out, qlock))
          for _ in range(conc)]
    [t.start() for t in ts]
    [t.join() for t in ts]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", default="both", choices=["official", "gateway", "both"])
    ap.add_argument("--mode", default="small", choices=["small", "big", "reuse"])
    ap.add_argument("--count", type=int, default=10)
    ap.add_argument("--concurrency", type=int, default=2)
    ap.add_argument("--out", default="./dsprobe-out")
    a = ap.parse_args()
    os.makedirs(a.out, exist_ok=True)
    arms = list(ARMS) if a.arm == "both" else [a.arm]
    if a.mode == "reuse":
        th = [threading.Thread(target=run_reuse, args=(x, a.count, a.out)) for x in arms]
    else:
        th = [threading.Thread(target=run_arm,
                               args=(x, a.mode, a.count, a.concurrency, a.out)) for x in arms]
    [t.start() for t in th]
    [t.join() for t in th]
    print("\n=== SUMMARY (truncated = content streamed, no finish_reason, no [DONE]) ===")
    for x in arms:
        rows = [r for r in STATS if r["arm"] == x]
        bad = [r for r in rows if r["truncated"]]
        print(f"{x}: {len(rows)} requests, truncated={len(bad)}, "
              f"transport_errors={sum(1 for r in rows if r.get('err'))}")
        for r in bad:
            print(f"   !! #{r['idx']} {r['t0']} dur={r['dur']:.1f}s chunks={r['chunks']} "
                  f"trace={r['headers'].get('x-ds-trace-id') or r['headers'].get('cf-ray')} "
                  f"log={r['headers'].get('x-opencode-log-id')}")


if __name__ == "__main__":
    main()