#!/usr/bin/env python3
"""Accurate-size stream test for frank-default-large (DSpark, PARALLEL=1).

Fires one prompt and reads back the server's OWN reported prompt_tokens so
the actual stream size is never mislabeled.  Then, with parallel=1, fires a
second concurrent stream to confirm it QUEUES behind the single slot and
completes without crashing the ROCm backend.

The prose is flowing/realistic, so a completion_tokens < max_tokens with
finish_reason "stop" is NOT a degenerate-prompt artifact.
"""
import json, sys, threading, time, urllib.request

HOST = sys.argv[1] if len(sys.argv) > 1 else "127.0.0.1"
PORT = sys.argv[2] if len(sys.argv) > 2 else "7770"
MAX_TOKENS = int(sys.argv[3]) if len(sys.argv) > 3 else 256
TARGET_TOKENS = int(sys.argv[4]) if len(sys.argv) > 4 else 16000

VAR = (
    "narrative nonfiction blends observation with reflection, ",
    "and a careful writer lets each clause earn its place ",
    "in the sentence's rising line, so meaning accumulates slowly. ",
    "The reader is not hurried; the reader is led. Loose and periodic, ",
    "the style turns on concrete nouns and verbs that carry weight, ",
    "and it resists the abstract noun that hides a vague claim. ",
    "A paragraph is not a list; it is a single turn of thought that ",
    "opens, develops, and lands. This is the craft of prose: choose, ",
    "arrange, and above all omit. Omission is the discipline that ",
    "keeps the whole clear, and clarity is the courtesy we owe the ",
    "willing reader."
)

def build_prompt(topic, salt):
    n = len(VAR)
    parts = []
    t = 0
    while t * 76 < TARGET_TOKENS:     # ~76 tokens per outer loop (measured)
        for k in range(5):
            parts.append(VAR[(t + k + salt) % n])
        parts.append(f" This is the {topic}, cued by note {salt} added later. ")
        t += 1
    return "".join(parts)

def post(prompt, max_tokens, timeout=3600):
    body = json.dumps({
        "prompt": prompt,
        "max_tokens": max_tokens,
        "temperature": 0.0,
        "stream": False,
    }).encode()
    req = urllib.request.Request(
        f"http://{HOST}:{PORT}/v1/completions",
        data=body, headers={"Content-Type": "application/json"})
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        out = json.load(resp)
    dt = time.time() - t0
    return out, dt

# ---- Stream 1: single 16k prompt -------------------------------
p1 = build_prompt("first subject", 1)
print(f"built prompt 1: {len(p1)} chars")
o1, t1 = post(p1, MAX_TOKENS)
u1 = o1.get("usage", {})
print(f"STREAM1 single: prompt_tokens={u1.get('prompt_tokens')} "
      f"completion_tokens={u1.get('completion_tokens')} "
      f"finish={o1.get('choices',[{}])[0].get('finish_reason')} secs={t1:.1f}")

# ---- Streams 2 & 3: concurrent, must queue (parallel=1) ---------
results = {}
def worker(key, prompt):
    try:
        out, dt = post(prompt, MAX_TOKENS, timeout=7200)
        u = out.get("usage", {})
        results[key] = {
            "prompt_tokens": u.get("prompt_tokens"),
            "completion_tokens": u.get("completion_tokens"),
            "finish": out.get("choices", [{}])[0].get("finish_reason"),
            "secs": round(dt, 1),
        }
    except Exception as e:
        results[key] = {"error": str(e)}

p2 = build_prompt("second subject", 2)
p3 = build_prompt("third subject", 3)
print(f"built prompt 2 ({len(p2)} chars) and prompt 3 ({len(p3)} chars); concurrent...")
ta = threading.Thread(target=worker, args=("s2", p2))
tb = threading.Thread(target=worker, args=("s3", p3))
ta.start(); tb.start(); ta.join(); tb.join()
for k in sorted(results):
    print(f"  [{k}] {results[k]}")

# verdict
def suspicious(r):
    if "error" in r:
        return True
    return r.get("finish") == "stop" and r.get("completion_tokens", 0) < MAX_TOKENS

allr = [("stream1", o1.get("choices",[{}])[0].get("finish_reason"),
         u1.get("completion_tokens")),
        ("s2", results.get("s2",{}).get("finish"), results.get("s2",{}).get("completion_tokens")),
        ("s3", results.get("s3",{}).get("finish"), results.get("s3",{}).get("completion_tokens"))]
print("\nSUMMARY")
ok = True
for name, fin, comp in allr:
    sig = fin == "stop" and comp is not None and comp < MAX_TOKENS
    err = comp is None
    if sig or err:
        ok = False
    print(f"  {name}: finish={fin} completion={comp}/{MAX_TOKENS} "
          f"{'SIGNATURE' if sig else ('ERROR' if err else 'ok')}")
print("  RESULT:", "PASS" if ok else "FAIL")