#!/usr/bin/env python3
"""probe-frank-archon-kv.py — verify the seqs x max-model-len KV contract on a
frank-archon (radiance) service. Run it ON dankserver: it talks to
127.0.0.1:7770 and builds its prompts from the repo checkout there.
No scratch container, no mask: the service under test is the service.

    ssh dankserver 'python3 /home/dank/serverAppConfigs/scripts/probe-frank-archon-kv.py'

Four probes, in increasing cost:
  A  decode at depth — one ~51k-token chat completion, 512 tokens, greedy.
     Reports the engine's own `timings` (prefill, decode, draft acceptance).
     The bench law applies: >=128 generated tokens, greedy, real text.
  B  the capacity contract — as many CONCURRENT chats of ~197k prompt tokens as
     the deployed `FRANK_ARCHON_MAX_SEQS` (read from the repo env file on this
     host), each carrying a unique marker at its head. All must answer their own
     marker. This is what the endpoint owes: N full-height conversations
     resident at once, exact retrieval in every one of them.
  C  prefix-cache restore — probe A's cold prompt again, so the disk tier has
     to give the blocks back (cache_n > 0, TTFT collapses).
  D  concurrency decode — `seqs` CONCURRENT 512-token generations at ~25 % of
     the window. This is the cost a slot count imposes on DECODE, which no
     prefill number shows: the fused-step schedule is unfair, so every stream is
     printed plus the aggregate. Measured 2026-10-13 at 3 x 200,000 fp8:
     single-stream 99.0 t/s (43.5k) and 93.3 (197k) — decode is nearly
     window-free — against 3-stream 113.9/43.3/10.9 t/s = 18.4 t/s aggregate.
     Concurrency, not the slab, is what decode pays for slots. The 2026-10-13
     shape hunt missed this for a whole day because leg B generates 24 tokens.
PROBE LAW (learned the hard way, twice): radiance routes the thinking block to
`reasoning_content`, so a chat request whose budget is spent inside
`<think>` returns EMPTY `content` with `finish_reason: stop`. A probe that
reads `content` alone scores 0/2 on a healthy server. Every request here sets
chat_template_kwargs.enable_thinking=false and reads both fields.

A/B method for a shape change (no scratch container): write a variant env file
on the host, point a transient systemd drop-in at it
(`Environment=FRANK_ARCHON_ENV_FILE=...` — drop-ins are read AFTER the unit's
own Environment= lines, so they win), restart, probe, delete the drop-in,
daemon-reload, restart. `systemctl mask` is not available: the unit file in
/etc/systemd/system is a relink symlink into the checkout.
"""
import json
import glob
import time
import urllib.request
import concurrent.futures

BASE = "http://127.0.0.1:7770"
MODEL = "frank-archon"
REPO = "/home/dank/serverAppConfigs"
GLOBS = [f"{REPO}/scripts/*.sh", f"{REPO}/serverAppConfigs-docs/**/*.md"]


def post(payload, timeout=1800):
    req = urllib.request.Request(
        BASE + "/v1/chat/completions",
        data=json.dumps(payload).encode(),
        headers={"Content-Type": "application/json"},
    )
    t0 = time.time()
    with urllib.request.urlopen(req, timeout=timeout) as r:
        body = json.loads(r.read().decode())
        hdr = dict(r.headers.items())
    return body, time.time() - t0, hdr


def get(path, timeout=60):
    with urllib.request.urlopen(BASE + path, timeout=timeout) as r:
        return json.loads(r.read().decode())


def corpus(min_chars):
    files = []
    for g in GLOBS:
        files.extend(sorted(glob.glob(g, recursive=True)))
    out, total = [], 0
    for f in files:
        try:
            t = open(f, encoding="utf-8", errors="replace").read()
        except OSError:
            continue
        out.append(t)
        total += len(t)
        if total > min_chars:
            break
    return "".join(text for text in out)


def prompt(text, chars, salt, fact):
    """Real repo text, salted so the prefix cache cannot answer, marker planted
    near the head, question at the tail."""
    head = f"\n[SALT {salt}] confidential marker: {fact}\n"
    body = text + text
    p = (head + body * max(1, chars // len(body)))[:chars]
    return p + "\n\nQuestion: what is the confidential marker in the " \
        "SALT block near the top of this document? Answer with the marker " \
        "alone, no other words.\n"


def chat(text, chars, salt, fact, max_tokens):
    return {"model": MODEL,
            "messages": [{"role": "user",
                          "content": prompt(text, chars, salt, fact)}],
            "max_tokens": max_tokens, "temperature": 0.0,
            "chat_template_kwargs": {"enable_thinking": False}}


def show(tag, j, wall, expect=None):
    tim, u = j.get("timings") or {}, j.get("usage", {})
    ch = (j.get("choices") or [{}])[0]
    msg = ch.get("message") or {}
    txt = (msg.get("content") or msg.get("reasoning_content") or "").strip()
    acc = (tim["draft_n_accepted"] / tim["draft_n"]) if tim.get("draft_n") else None
    line = (f"[{tag}] prompt={u.get('prompt_tokens')} gen={u.get('completion_tokens')} "
            f"wall={wall:.1f}s prefill={tim.get('prompt_per_second', 0):.0f} t/s "
            f"decode={tim.get('predicted_per_second', 0):.1f} t/s "
            f"draft={tim.get('draft_n')}/{tim.get('draft_n_accepted')} "
            f"accept={acc if acc is None else round(acc, 4)} "
            f"cache_n={tim.get('cache_n')} finish={ch.get('finish_reason')}")
    if expect:
        hit = expect.split("-")[0] in txt.lower()
        line += f" expected={expect} answered={txt[:80]!r} HIT={hit}"
    print(line)


def deployed_shape():
    """The contract is the deployed env, not a constant: slot count and window
    come from the checkout on this host, so the probe cannot drift from the
    service it is testing (the 2026-10-13 shape is 4 x 200,000; before that
    2 x 200,000, and the fp8 era 3 x 200,000)."""
    shape = {}
    with open(f"{REPO}/scripts/frank-archon.env") as fh:
        for line in fh:
            line = line.strip()
            if line.startswith("FRANK_ARCHON_") and "=" in line:
                key, _, val = line.partition("=")
                shape[key] = val.strip().strip('"')
    return (int(shape["FRANK_ARCHON_MAX_SEQS"]),
            int(shape["FRANK_ARCHON_MAX_MODEL_LEN"]),
            shape["FRANK_ARCHON_KV_DTYPE"])


def main():
    seqs, maxlen, dtype = deployed_shape()
    print(f"shape: {seqs} x {maxlen:,} tokens, kv {dtype}")
    print("health:", json.dumps(get("/health")))
    print("models:", [m.get("id") for m in get("/v1/models").get("data", [])])
    deep = corpus(160_000)
    wide = corpus(1_100_000)

    # 3.28 chars/token is the measured ratio of this corpus through the
    # froggeric template (646,000 chars -> 196,919 tokens at the 2026-10-11
    # deploy); 98.5 % of the window leaves the response room under the ceiling.
    chars = int(0.985 * maxlen * 3.28)
    # Each marker's FIRST hyphen-separated word is unique, because show()
    # scores the hit on that word — a shared prefix would score every slot HIT
    # on any answer that says the word.
    words = ["cobalt", "quartz", "basalt", "ember", "willow", "cobweb", "flint"]
    facts = {f"B{i+1}": f"{words[i % len(words)]}-flag-{(i + 1) * 137}"
             for i in range(seqs)}
    bodies = {s: chat(wide, chars, s, f, 24) for s, f in facts.items()}
    print(f"[B] {len(bodies['B1']['messages'][0]['content']):,} chars per "
          f"request, x{seqs} concurrent")

    # A + C first (cheap), B last (the expensive set).
    for tag, salt in (("A-cold", "A1-9137"), ):
        j, wall, _ = post(chat(deep, 140_000, salt, "zephyr-lantern-7", 512))
        show(tag, j, wall, "zephyr-lantern-7")
    t0 = time.time()
    with concurrent.futures.ThreadPoolExecutor(max_workers=seqs) as ex:
        futs = {s: ex.submit(post, b) for s, b in bodies.items()}
        for s, fut in futs.items():
            try:
                j, wall, _ = fut.result()
                show(f"B-{s}", j, wall, facts[s])
            except Exception as e:  # a preemption/timeout lands here
                print(f"[B-{s}] FAILED after {time.time() - t0:.1f}s: {e!r}")
    print(f"[B] total {time.time() - t0:.1f}s")
    # D — concurrency decode: seqs streams, 512 tokens each, ~25 % depth.
    dchars = int(0.25 * maxlen * 3.28)
    dbodies = [chat(deep, dchars, f"D{i + 1}-9137", f"flint-marker-{i}", 512)
               for i in range(seqs)]
    print(f"[D] {seqs} concurrent streams x 512 tokens, "
          f"{dchars:,} chars each")
    t0 = time.time()
    with concurrent.futures.ThreadPoolExecutor(max_workers=seqs) as ex:
        futs = [ex.submit(post, b) for b in dbodies]
        for i, fut in enumerate(futs, 1):
            try:
                j, wall, _ = fut.result()
                show(f"D-s{i}", j, wall)
            except Exception as e:  # a preemption/timeout lands here
                print(f"[D-s{i}] FAILED after {time.time() - t0:.1f}s: {e!r}")
    tot = time.time() - t0
    print(f"[D] {seqs} x 512 tokens in {tot:.1f}s = "
          f"{seqs * 512 / tot:.1f} t/s aggregate decode")
    j, wall, _ = post(chat(deep, 140_000, "A1-9137", "zephyr-lantern-7", 512))
    show("C-prefix-hit", j, wall, "zephyr-lantern-7")


if __name__ == "__main__":
    main()