"""512k acceptance: N consecutive >=500k-token prompts through a live server.

Usage: python3 longctx.py <port> <model> [n_rounds=3] [target_tokens=505000]

Each round builds a DIFFERENT ~target-token prompt (defeats prefix caching so
every round is a fresh full-length prefill — the fragmentation test), asks for
a short completion, and reports usage.prompt_tokens, prefill time, and decode
tok/s. Round 1 additionally re-sends its own prompt (prefix-cache hit) with a
longer completion to measure steady decode tok/s at full context.
"""

import json
import random
import sys
import time
import urllib.request

PORT, MODEL = sys.argv[1], sys.argv[2]
ROUNDS = int(sys.argv[3]) if len(sys.argv) > 3 else 3
TARGET = int(sys.argv[4]) if len(sys.argv) > 4 else 505000

WORDS = (
    "the quick brown fox jumps over a lazy dog while seventeen green wizards "
    "quietly brew potent elixirs behind twelve ancient marble columns near "
    "the harbor as autumn rain drums softly on copper rooftops and distant "
    "bells mark the passing hours for patient scholars reading dusty scrolls"
).split()


def build_prompt(seed, n_tokens):
    rng = random.Random(seed)
    # ~1.3 tokens/word for this vocab; overshoot handled by the caller loop.
    words = []
    target_words = int(n_tokens / 1.35)
    while len(words) < target_words:
        chunk = WORDS[:]
        rng.shuffle(chunk)
        words.extend(chunk)
        words.append(f"marker{rng.randint(1000, 9999)}.")
    return " ".join(words)


def post(path, payload, timeout=3600):
    req = urllib.request.Request(
        f"http://127.0.0.1:{PORT}{path}",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def count_tokens(text):
    r = post("/tokenize", {"model": MODEL, "prompt": text})
    return r["count"]


def chat(prompt_text, max_tokens):
    t0 = time.time()
    r = post(
        "/v1/chat/completions",
        {
            "model": MODEL,
            "messages": [
                {
                    "role": "user",
                    "content": prompt_text
                    + "\n\nThe text above is filler. Reply with the single "
                    "word ACKNOWLEDGED followed by one short sentence.",
                }
            ],
            "max_tokens": max_tokens,
            "temperature": 0,
        },
    )
    dt = time.time() - t0
    return r, dt


# calibrate tokens/char on a small sample
sample = build_prompt(0, 2000)
spt = count_tokens(sample) / len(sample)
print(f"calibration: {spt:.5f} tok/char on {len(sample)} chars")

results = []
for rnd in range(1, ROUNDS + 1):
    # build to target using measured ratio, then trim/verify via /tokenize
    text = build_prompt(rnd, int(TARGET * 1.02))
    need_chars = int(TARGET / spt)
    text = text[:need_chars]
    n_tok = count_tokens(text)
    print(f"[round {rnd}] built prompt: {n_tok} tokens ({len(text)} chars)")
    if n_tok < 500000:
        # extend deterministically until over 500k
        while n_tok < 500500:
            text = text + " " + build_prompt(100 + rnd, 6000)
            n_tok = count_tokens(text)
        print(f"[round {rnd}] extended to {n_tok} tokens")

    r, dt = chat(text, 64)
    u = r["usage"]
    reply = r["choices"][0]["message"].get("content") or r["choices"][0][
        "message"
    ].get("reasoning_content", "")
    ok = u["prompt_tokens"] >= 500000 and r["choices"][0].get("finish_reason") in (
        "stop",
        "length",
    )
    print(
        f"[round {rnd}] prompt_tokens={u['prompt_tokens']} "
        f"completion={u['completion_tokens']} wall={dt:.1f}s "
        f"finish={r['choices'][0].get('finish_reason')} ok={ok}"
    )
    print(f"[round {rnd}] reply: {reply[:200]!r}")
    entry = {
        "round": rnd,
        "prompt_tokens": u["prompt_tokens"],
        "wall_s": dt,
        "ok": ok,
        "reply_head": reply[:200],
    }

    if rnd == 1:
        # decode-rate at full context: resend (prefix cached) with more tokens
        r1, dt1 = chat(text, 1)
        r2, dt2 = chat(text, 257)
        gen = r2["usage"]["completion_tokens"]
        # crude: dt2 includes cached-prefill overhead ~= dt1
        rate = (gen - 1) / max(dt2 - dt1, 1e-6)
        print(
            f"[round 1] decode at {u['prompt_tokens']} ctx: {gen} tok in "
            f"{dt2:.1f}s (1-tok call {dt1:.1f}s) -> ~{rate:.1f} tok/s"
        )
        entry["decode_tok_s_at_ctx"] = round(rate, 1)
    results.append(entry)

json.dump(results, open("/tmp/longctx-results.json", "w"), indent=1)
n_ok = sum(1 for e in results if e["ok"])
print(f"PASS {n_ok}/{ROUNDS}" if n_ok == ROUNDS else f"FAIL {n_ok}/{ROUNDS}")
sys.exit(0 if n_ok == ROUNDS else 1)