#!/usr/bin/env python3
"""probe-frank-archon-kv.py — verify the 2 x 200,000 KV contract on a LIVE
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'

Three 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).
     Method: >=128 generated tokens (below that, first-token latency
     dominates the decode rate), greedy (sampling noise is unattributable),
     real text (synthetic text inflates draft acceptance).
  B  the capacity contract — TWO concurrent chats of ~197k prompt tokens
     (~394k tokens resident) each carrying a unique marker at its head. Both
     must answer their own marker. This is what a 2 x 200,000 endpoint owes:
     two full-height conversations resident at once, exact retrieval in both.
  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).

PROBE PITFALL: 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 main():
    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)

    facts = {"B1": "cobalt-harbor-13", "B2": "quartz-meadow-42"}
    bodies = {s: chat(wide, 646_000, s, f, 24) for s, f in facts.items()}
    print(f"[B] {len(bodies['B1']['messages'][0]['content']):,} chars per "
          f"request, x2 concurrent")

    # A + C first (cheap), B last (the expensive pair).
    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=2) 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")
    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()