"""Scoped wall-clock micro-benchmark (external-review item 3).

Not a production deployment study. Measures, on the real infra:
  (a) draft-side overhead: embed(query) + nearest-neighbor retrieval against a
      realistically-sized personal store (~48 entries) — timed on this CPU node;
  (b) target-side decode speed of the served gpt-oss-120b: per-token decode
      time from paired short/long generations (sequential, unbatched);
then compares mean retrieval overhead against the decode time represented by
the accepted-token gap (personal 19.96 vs static 15.70 MAT, headline run).
Writes results/wallclock_micro.json.
"""
from __future__ import annotations

import json
import time
from pathlib import Path

import numpy as np
import requests

from .data import load_bfcl
from .memory import Embedder, _best_match, Entry

ROOT = Path(__file__).resolve().parent.parent
RESULTS = ROOT / "results"
URL = "http://localhost:30000/v1/chat/completions"
N_RETRIEVAL = 200
N_DECODE_PAIRS = 30


def main():
    tasks = load_bfcl()[:N_RETRIEVAL]
    emb = Embedder()
    # warm the embedder, then build a 48-entry store
    store = [Entry(emb.embed(t.query), "x") for t in tasks[:48]]
    for t in tasks[:20]:
        emb.embed(t.query + " warm")
    lat = []
    for t in tasks:
        q = t.query + " ?"                      # bust the embed cache
        t0 = time.perf_counter()
        e = emb.embed(q)
        _best_match(e, store)
        lat.append((time.perf_counter() - t0) * 1000)
    lat.sort()
    retr = {"n": len(lat), "mean_ms": round(float(np.mean(lat)), 2),
            "median_ms": round(lat[len(lat) // 2], 2),
            "p95_ms": round(lat[int(0.95 * len(lat))], 2)}

    def timed_gen(prompt, max_tok):
        t0 = time.perf_counter()
        r = requests.post(URL, json={
            "model": "gpt-oss-120b", "temperature": 0.0,
            "max_tokens": max_tok,
            "messages": [{"role": "user", "content": prompt}]}, timeout=120)
        dt = time.perf_counter() - t0
        u = r.json().get("usage", {})
        return dt, u.get("completion_tokens", max_tok)

    per_tok = []
    for t in tasks[:N_DECODE_PAIRS]:
        p = "Write a long detailed paragraph about: " + t.query[:200]
        d1, n1 = timed_gen(p, 8)
        d2, n2 = timed_gen(p, 168)
        if n2 > n1:
            per_tok.append((d2 - d1) / (n2 - n1) * 1000)
    per_tok.sort()
    dec = {"n_pairs": len(per_tok),
           "per_token_ms_median": round(per_tok[len(per_tok) // 2], 2),
           "per_token_ms_mean": round(float(np.mean(per_tok)), 2)}

    ptok = dec["per_token_ms_median"]
    out = {
        "retrieval_overhead": retr,
        "decode": dec,
        "comparison": {
            "personal_MAT_19.96_decode_ms": round(19.96 * ptok, 1),
            "static_MAT_15.70_decode_ms": round(15.70 * ptok, 1),
            "gap_4.26_tokens_decode_ms": round(4.26 * ptok, 1),
            "retrieval_overhead_mean_ms": retr["mean_ms"],
            "overhead_over_gap_pct": round(
                100 * retr["mean_ms"] / (4.26 * ptok), 1),
        },
        "caveats": ("Sequential unbatched decode on the live server (batching "
                    "changes per-token time); retrieval timed on the CPU node,"
                    " single-threaded; no end-to-end speculative decoder is "
                    "integrated — this bounds the overhead-vs-savings ratio, "
                    "not deployed speedup."),
    }
    (RESULTS / "wallclock_micro.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()