"""Reproduction script for the Jev measurements taken on 2026-09-20 from Vietnam.

Every payload below is verbatim what was sent. Run with a Vercel AI Gateway key:

    AI_GATEWAY_API_KEY=vck_... python3 repro.py [latency|headtohead|overhead|promptsize]

Total cost of a full run is under $0.002 at list prices. The free tier rate-limits
around 30 requests per minute (advertised) but was observed to reject earlier than
that, so a full run may need to be retried.
"""

import json
import os
import sys
import time
import urllib.error
import urllib.request

KEY = os.environ.get("AI_GATEWAY_API_KEY", "")
JEV_URL = "https://ai-gateway.vercel.sh/typesafe/v1/systemone"
CHAT_URL = "https://ai-gateway.vercel.sh/v1/chat/completions"

JEV_MODEL = "typesafe-ai/jev"
CHAT_MODELS = ["google/gemini-2.5-flash-lite", "openai/gpt-4.1-nano"]

# ---------------------------------------------------------------- shared state

HEADTOHEAD_STATE = {
    "task": "Add a --json flag to the sweep script and update USAGE.md",
    "files_touched": 2,
    "diff_lines": 47,
}

# ------------------------------------------------- experiment 1+2: head-to-head

HEADTOHEAD_JEV_BODY = {
    "model": JEV_MODEL,
    "state": HEADTOHEAD_STATE,
    "questions": {
        "worker": {
            "type": "choice",
            "instructions": "Which worker class should handle this task",
            "criteria": {
                "cheap_loop": "A small mechanical edit, no design judgement",
                "cli_agent": "Needs repo context and multi-file edits",
                "frontier": "Ambiguous spec or cross-module design work",
            },
        },
        "risk": {
            "type": "score",
            "instructions": "How risky is landing this change without review",
            "criteria": [
                "Trivial, safe to auto-merge",
                "Ordinary, a glance is enough",
                "Needs a careful read",
                "Do not land without a human",
            ],
        },
        "needs_docs": {
            "type": "noul",
            "instructions": "This change requires a USAGE.md update",
        },
    },
}

HEADTOHEAD_SYSTEM_PROMPT = (
    "You route coding tasks. Reply with ONLY a JSON object, no prose, no markdown fence. "
    'Schema: {"worker":"cheap_loop|cli_agent|frontier","risk":0|1|2|3,"needs_docs":true|false}. '
    "worker: cheap_loop=small mechanical edit no design judgement; "
    "cli_agent=needs repo context and multi-file edits; "
    "frontier=ambiguous spec or cross-module design work. "
    "risk of landing without review: 0=trivial safe to auto-merge, 1=ordinary a glance is enough, "
    "2=needs a careful read, 3=do not land without a human. "
    "needs_docs: does this require a USAGE.md update."
)


def headtohead_chat_body(model):
    return {
        "model": model,
        "max_tokens": 80,
        "temperature": 0,
        "messages": [
            {"role": "system", "content": HEADTOHEAD_SYSTEM_PROMPT},
            {"role": "user", "content": json.dumps(HEADTOHEAD_STATE)},
        ],
    }


# ------------------------------------------------- experiment 3: token overhead
OVERHEAD_SENTENCE = (
    "The deployment pipeline reported a failure in the staging environment after the "
    "configuration was updated. Engineers reviewed the logs and found no clear cause. "
)
OVERHEAD_SIZES = {"tiny": 1, "small": 8, "medium": 40, "large": 160}
OVERHEAD_JEV_QUESTIONS = {
    "q": {"type": "noul", "instructions": "The text mentions a failure"}
}
OVERHEAD_CHAT_PREFIX = "The text mentions a failure. Answer yes or no.\n\n"

# ------------------------------------------------------ experiment 4: latency

LATENCY_JEV_BODY = {
    "model": JEV_MODEL,
    "state": "Deploy step failed: connection refused to 127.0.0.1:41700",
    "questions": {
        "is_network": {
            "type": "noul",
            "instructions": "This is a network or connectivity failure",
        }
    },
}

# ----------------------------------------------------------------------- calls


def post(url, body, timeout=45):
    req = urllib.request.Request(
        url,
        data=json.dumps(body).encode(),
        headers={"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"},
    )
    t0 = time.perf_counter()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return {"ms": (time.perf_counter() - t0) * 1000, "code": r.status,
                    "body": json.loads(r.read())}
    except urllib.error.HTTPError as e:
        return {"ms": (time.perf_counter() - t0) * 1000, "code": e.code,
                "body": e.read()[:300].decode("utf8", "replace")}
    except Exception as e:
        return {"ms": (time.perf_counter() - t0) * 1000, "code": "EXC",
                "body": f"{type(e).__name__}: {e}"}


def median(xs):
    xs = sorted(xs)
    if not xs:
        return float("nan")
    mid = len(xs) // 2
    return xs[mid] if len(xs) % 2 else (xs[mid - 1] + xs[mid]) / 2


def run_latency(n=25):
    """New connection per call. Reported with curl: median 670 ms, max 818 ms."""
    rs = [post(JEV_URL, LATENCY_JEV_BODY) for _ in range(n)]
    ok = [r["ms"] for r in rs if r["code"] == 200]
    print(f"latency n={len(rs)} ok={len(ok)} median={median(ok):.0f}ms max={max(ok):.0f}ms")


def run_headtohead(rounds=12):
    """Interleaved so time-of-day drift hits every model equally.
    Reported: jev median 667 ms / 462 in tok; flash-lite 820 ms / 185; nano 1110 ms / 185."""
    out = {"jev": [], **{m: [] for m in CHAT_MODELS}}
    for _ in range(rounds):
        out["jev"].append(post(JEV_URL, HEADTOHEAD_JEV_BODY))
        for m in CHAT_MODELS:
            out[m].append(post(CHAT_URL, headtohead_chat_body(m)))
    for name, rs in out.items():
        ok = [r for r in rs if r["code"] == 200]
        ms = [r["ms"] for r in ok]
        u = ok[0]["body"]["usage"]
        tin = u.get("input_tokens", u.get("prompt_tokens"))
        tout = u.get("output_tokens", u.get("completion_tokens"))
        print(f"{name:32s} ok={len(ok)}/{len(rs)} median={median(ms):.0f}ms "
              f"max={max(ms):.0f}ms in={tin} out={tout}")


def run_overhead():
    """Four state sizes spanning 160x. Reported: diff is exactly 259 at every size,
    fit jev_in = 1.00 * chat_in + 259. Only the DIFF column identifies the overhead
    as fixed; the ratio column falls from 7.64 to 1.06 and is misleading alone."""
    print(f"{'size':8s} {'chars':>7s} {'chat_in':>8s} {'chat_out':>9s} {'jev_in':>7s} {'diff':>6s}")
    for name, n in OVERHEAD_SIZES.items():
        txt = OVERHEAD_SENTENCE * n
        j = post(JEV_URL, {"model": JEV_MODEL, "state": txt,
                           "questions": OVERHEAD_JEV_QUESTIONS})
        g = post(CHAT_URL, {"model": CHAT_MODELS[0], "max_tokens": 4, "temperature": 0,
                            "messages": [{"role": "user",
                                          "content": OVERHEAD_CHAT_PREFIX + txt}]})
        ji = j["body"]["usage"]["input_tokens"]
        gi = g["body"]["usage"]["prompt_tokens"]
        go = g["body"]["usage"]["completion_tokens"]
        print(f"{name:8s} {len(txt):>7d} {gi:>8d} {go:>9d} {ji:>7d} {ji - gi:>6d}")


def run_promptsize():
    """The chat side's fixed prompt cost, which sets the real break-even.
    Reported: system prompt alone 148 tokens, full head-to-head request 185."""
    def toks(msgs):
        b = {"model": CHAT_MODELS[0], "max_tokens": 1, "temperature": 0, "messages": msgs}
        return post(CHAT_URL, b)["body"]["usage"]["prompt_tokens"]

    only_sys = toks([{"role": "system", "content": HEADTOHEAD_SYSTEM_PROMPT},
                     {"role": "user", "content": "x"}])
    full = toks([{"role": "system", "content": HEADTOHEAD_SYSTEM_PROMPT},
                 {"role": "user", "content": json.dumps(HEADTOHEAD_STATE)}])
    print(f"chat system prompt alone: {only_sys} tokens; full request: {full} tokens")


if __name__ == "__main__":
    if not KEY:
        sys.exit("set AI_GATEWAY_API_KEY")
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    for name, fn in [("latency", run_latency), ("headtohead", run_headtohead),
                     ("overhead", run_overhead), ("promptsize", run_promptsize)]:
        if which in (name, "all"):
            print(f"\n=== {name} ===")
            fn()