"""Prompt-cache lab for the "AI served someone else's pages" post.

E1  cold vs warm: same long prefix, time to first token + cached_tokens
E2  over-capacity: prompt bigger than the context window -> clean refusal?
E3  neighbour test: warm resource A, ask resource B the same prompt

Keys are fetched from `az` at runtime and never written to disk.
Run:  python3 -I cache_lab.py [e1|e2|e3|all]
"""
import json, os, random, statistics, string, subprocess, sys, time, urllib.request, urllib.error

API = "2025-01-01-preview"
HERE = os.path.dirname(os.path.abspath(__file__))
RES = {
    "A": ("mearytripathi-resource", "rg-mearytripathi"),
    "B": ("aiprodready8372", "rg-prodready-test2"),
}
_keys = {}


def key(res):
    if res not in _keys and os.environ.get(f"KEY_{res}"):
        _keys[res] = os.environ[f"KEY_{res}"]
    if res not in _keys:
        n, g = RES[res]
        _keys[res] = subprocess.check_output(
            ["az", "cognitiveservices", "account", "keys", "list", "-n", n, "-g", g, "--query", "key1", "-o", "tsv"],
            text=True).strip()
    return _keys[res]


def url(res, dep):
    return f"https://{RES[res][0]}.cognitiveservices.azure.com/openai/deployments/{dep}/chat/completions?api-version={API}"


def call(res, dep, prompt, max_tokens=8, timeout=120, tries=6):
    """call_once, retrying 429s after the retry-after the service asks for."""
    for attempt in range(tries):
        out = call_once(res, dep, prompt, max_tokens, timeout)
        if out["status"] != 429 or "context" in (out["error"] or ""):
            break
        time.sleep(out.pop("retry_after", None) or 5)
    out.pop("retry_after", None)
    out["retries"] = attempt
    return out


def call_once(res, dep, prompt, max_tokens=8, timeout=120):
    """Streamed call. Returns ttft_ms, total_ms, usage, status, error."""
    body = json.dumps({
        "messages": [{"role": "system", "content": prompt},
                     {"role": "user", "content": "Reply with one word: ok"}],
        "max_tokens": max_tokens, "temperature": 0,
        "stream": True, "stream_options": {"include_usage": True},
    }).encode()
    req = urllib.request.Request(url(res, dep), data=body, method="POST",
                                 headers={"api-key": key(res), "Content-Type": "application/json"})
    t0 = time.perf_counter()
    out = {"res": res, "dep": dep, "req_bytes": len(body), "status": None, "ttft_ms": None,
           "total_ms": None, "prompt_tokens": None, "cached_tokens": None, "error": None}
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            out["status"] = r.status
            for raw in r:
                line = raw.decode().strip()
                if not line.startswith("data:") or line == "data: [DONE]":
                    continue
                chunk = json.loads(line[5:])
                ch = chunk.get("choices") or []
                if out["ttft_ms"] is None and ch and (ch[0].get("delta") or {}).get("content"):
                    out["ttft_ms"] = round((time.perf_counter() - t0) * 1000)
                if chunk.get("usage"):
                    u = chunk["usage"]
                    out["prompt_tokens"] = u.get("prompt_tokens")
                    out["cached_tokens"] = (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0)
    except urllib.error.HTTPError as e:
        out["status"] = e.code
        out["error"] = e.read().decode()[:600]
        ra = e.headers.get("retry-after")
        out["retry_after"] = float(ra) if ra and ra.replace(".", "").isdigit() else None
    except Exception as e:  # timeouts, resets
        out["error"] = f"{type(e).__name__}: {e}"
    out["total_ms"] = round((time.perf_counter() - t0) * 1000)
    return out


def long_prefix(nonce, words=24000):
    """~30K-token 'policy manual' that starts with a nonce, so each nonce is a fresh cache entry."""
    rnd = random.Random(42)  # same body every time; only the nonce changes
    vocab = ["refund", "invoice", "customer", "ticket", "priority", "escalate", "warehouse", "delivery",
             "GST", "Bengaluru", "Pune", "order", "within", "hours", "manager", "approve", "policy",
             "section", "clause", "exception", "weekend", "festival", "Diwali", "return", "damaged"]
    body = " ".join(rnd.choice(vocab) for _ in range(words))
    return f"Session {nonce}. Internal support manual follows.\n{body}"


def nonce():
    return "".join(random.choices(string.ascii_lowercase + string.digits, k=12))


def e1(trials=12):
    rows = []
    for i in range(trials):
        p = long_prefix(nonce())
        for label in ("cold", "warm1", "warm2", "warm3"):
            r = call("A", "gpt-41-mini", p)
            r.update(trial=i, label=label, hit=bool(r["cached_tokens"]))
            rows.append(r)
            print("E1", i, label, r["status"], r["ttft_ms"], r["prompt_tokens"], r["cached_tokens"], (r["error"] or "")[:80], flush=True)
            time.sleep(4)
    return rows


def e2():
    """gpt-4.1-nano context is 1,047,576 tokens. Ask for ~1.2M. Rejected requests are not billed."""
    rows = []
    for words in (1_200_000,):
        p = " ".join(["apple"] * words)
        r = call("A", "gpt-41-nano", p, timeout=300)
        r.update(words=words)
        rows.append(r)
        print("E2", words, r["status"], r["total_ms"], (r["error"] or "")[:300], flush=True)
    return rows


def e3(trials=10, b_dep="gpt-41-mini"):
    """Per trial: control (B's first sight of a prefix nobody sent), then warm A until it hits,
    then B's first request for A's prefix. B-first hitting the cache means A's cache served B."""
    rows = []

    def log(i, label, r):
        r.update(trial=i, label=label)
        rows.append(r)
        print("E3", i, label, r["status"], r["ttft_ms"], r["prompt_tokens"], r["cached_tokens"], r["retries"], flush=True)
        time.sleep(4)

    for i in range(trials):
        log(i, "control-B-cold", call("B", b_dep, long_prefix(nonce())))
        p = long_prefix(nonce())
        for k in range(4):
            r = call("A", "gpt-41-mini", p)
            log(i, f"warm-A-{k + 1}", r)
            if r["cached_tokens"]:
                break
        log(i, "B-first", call("B", b_dep, p))
        log(i, "B-second", call("B", b_dep, p))
    return rows


def summarise(rows, by="label"):
    groups = {}
    for r in rows:
        if r["status"] == 200:
            groups.setdefault(r[by], []).append(r)
    s = {}
    for g, rs in groups.items():
        s[g] = {
            "n": len(rs),
            "ttft_median_ms": statistics.median(x["ttft_ms"] for x in rs if x["ttft_ms"]),
            "prompt_tokens_median": statistics.median(x["prompt_tokens"] for x in rs),
            "cached_share_median": round(statistics.median(x["cached_tokens"] / x["prompt_tokens"] for x in rs), 3),
            "calls_with_cache_hit": sum(1 for x in rs if x["cached_tokens"]),
        }
    return s


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    path = os.path.join(HERE, "results.json")
    results = json.load(open(path)) if os.path.exists(path) else {}
    results["run_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
    if which in ("e1", "all"):
        rows = e1(); results["e1"] = {"rows": rows, "summary": summarise(rows), "by_hit": summarise(rows, "hit")}
    if which in ("e2", "all"):
        results["e2"] = {"rows": e2()}
    if which in ("e3", "all"):
        rows = e3(); results["e3"] = {"rows": rows, "summary": summarise(rows)}
    if which == "e3b":  # retest: B's deployment has a different name than A's
        rows = e3(trials=16, b_dep="neighbour-b")
        results["e3b"] = {"rows": rows, "summary": summarise(rows), "a_dep": "gpt-41-mini", "b_dep": "neighbour-b"}
    json.dump(results, open(path, "w"), indent=1)
    for k in ("e1", "e3", "e3b"):
        if k in results:
            print(k, json.dumps(results[k]["summary"], indent=1))
    if "e1" in results:
        print("e1 by_hit", json.dumps(results["e1"]["by_hit"], indent=1))