#!/usr/bin/env python3
"""
REQUIREMENT 4 -- "Submit P50 / P70 / P100 latency numbers for your pipeline,
measured across a reasonable number of test queries — not a single best-case run."

P50 / P70 / P100 exactly as the brief names them. P100 is the max, so the worst
case is visible rather than averaged away.

WHAT COUNTS TOWARD THE 200 ms
-----------------------------
The brief scopes the target as "chunking + vector DB retrieval + everything
through to final output". Speech-to-text is a mandated third-party API
(requirement 1: Sarvam or ElevenLabs), so its round trip is reported SEPARATELY
rather than folded into a number we do not control. Both are printed. Nothing is
hidden inside a single figure.

THE STAGE THAT ACTUALLY DECIDES CPU-VS-GPU
------------------------------------------
`retrieve` is not just the dot product. It is a full bge-m3 forward pass to embed
the query, and THEN a dot product over the index. On a GPU the forward pass is a
few ms; on a laptop CPU it can be 50-150 ms, which is most of the budget. The
dot product is trivial either way. So this script times query embedding on its
own -- that single number tells you whether a laptop can serve inside 200 ms
before you spend a day moving.

  python scripts/bench_latency.py --n 300 --langs hi,bn,kn
  python scripts/bench_latency.py --embed-only     # the CPU feasibility check
"""
from __future__ import annotations

import argparse
import json
import statistics as st
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from src.schema_utils import default_root  # noqa: E402

QUERIES = {
    "hi": ["कैंटालूप को पकने में कितना समय लगता है", "एक चील कितनी तेजी से उड़ती है",
           "कॉर्पोरेशन क्या है", "स्टबहब का टोल फ्री नंबर क्या है",
           "मुंबई किस राज्य की राजधानी है"],
    "bn": ["তামিল ভাষা কোথায় বলা হয়", "একটি ঈগল কত দ্রুত উড়তে পারে",
           "কর্পোরেশন কী", "ক্যান্টালুপ পাকতে কত সময় লাগে"],
    "kn": ["ಕಾರ್ಪೊರೇಶನ್ ಎಂದರೇನು", "ಹದ್ದು ಎಷ್ಟು ವೇಗವಾಗಿ ಹಾರುತ್ತದೆ",
           "ಕ್ಯಾಂಟಲೋಪ್ ಹಣ್ಣಾಗಲು ಎಷ್ಟು ಸಮಯ ಬೇಕು"],
    "ta": ["தமிழ் மொழி எங்கு பேசப்படுகிறது", "ஒரு கழுகு எவ்வளவு வேகமாக பறக்கும்"],
}


def pct(v: list[float], q: float) -> float:
    """P100 must be the true max, so index by ceil rather than interpolating."""
    if not v:
        return 0.0
    s = sorted(v)
    if q >= 1.0:
        return s[-1]
    return s[min(len(s) - 1, int(q * len(s)))]


def table(name: str, v: list[float], budget: float | None = None) -> dict:
    d = {"stage": name, "n": len(v), "mean": round(st.mean(v), 2),
         "p50": round(pct(v, .50), 2), "p70": round(pct(v, .70), 2),
         "p100": round(pct(v, 1.0), 2)}
    if budget:
        d["within_budget_pct"] = round(100 * sum(1 for x in v if x <= budget) / len(v), 1)
    return d


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=Path, default=None)
    ap.add_argument("--langs", default="hi,bn,kn")
    ap.add_argument("--n", type=int, default=300, help="total requests")
    ap.add_argument("--budget-ms", type=float, default=200.0)
    ap.add_argument("--embed-only", action="store_true",
                    help="time query embedding alone — the CPU feasibility check")
    ap.add_argument("--warmup", type=int, default=5)
    args = ap.parse_args()

    root = args.root.expanduser().resolve() if args.root else default_root()
    langs = [l for l in args.langs.split(",") if l in QUERIES]
    if not langs:
        raise SystemExit(f"no canned queries for {args.langs}; have {sorted(QUERIES)}")

    import torch
    dev = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"==> device: {dev}")
    if dev == "cpu":
        import os
        print(f"    threads: {torch.get_num_threads()}  cpus: {os.cpu_count()}")

    # ---------------------------------------------------------- embed only
    from src.evaluate_retrieval import Embedder
    man_p = root / "index" / "manifest.json"
    man = json.loads(man_p.read_text()) if man_p.exists() else {}
    model = man.get("model")
    if model is None:
        hits = list((root / "hf_cache" / "hub").glob("models--BAAI--bge-m3/snapshots/*"))
        model = str(hits[0]) if hits else "BAAI/bge-m3"
    emb = Embedder(model, dev, 1, man.get("max_len", 192))

    flat = [(lg, q) for lg in langs for q in QUERIES[lg]]
    for _ in range(args.warmup):                       # first pass allocates
        emb.encode([flat[0][1]])

    e_ms = []
    for i in range(args.n):
        _lg, q = flat[i % len(flat)]
        t0 = time.perf_counter()
        emb.encode([q])
        e_ms.append((time.perf_counter() - t0) * 1000)

    print(f"\n{'='*66}\nQUERY EMBEDDING — bge-m3 forward pass, batch of 1\n{'='*66}")
    t = table("embed_query", e_ms, args.budget_ms)
    print(f"  n={t['n']}  mean {t['mean']} ms   P50 {t['p50']}   P70 {t['p70']}   "
          f"P100 {t['p100']}")
    print(f"\n  budget headroom AFTER embedding, per percentile:")
    for name in ("p50", "p70", "p100"):
        head_p = args.budget_ms - t[name]
        flag = "  <-- OVER BUDGET" if head_p < 0 else ""
        print(f"    {name.upper():5s}{t[name]:9.1f} ms   leaves {head_p:8.1f} ms{flag}")

    # Judge on P100, not P70. The brief asks for P100 specifically, so the tail
    # is reported whether or not we like it -- and a percentile that already
    # exceeds the budget before retrieval even starts is a fail, not a pass.
    # (`within_budget_pct` compares EMBEDDING alone against the WHOLE budget,
    # which is the wrong denominator; it is printed for reference only.)
    head = args.budget_ms - t["p100"]
    print(f"\n  reference only: {t['within_budget_pct']}% of embeddings fit in the "
          f"full {args.budget_ms:.0f} ms\n  (wrong denominator — embedding is one stage, "
          f"not the whole pipeline)")
    if head < 0:
        print(f"\n  >> FAILS AT P100. Embedding alone is {-head:.0f} ms over budget before")
        print("     retrieval, the reader or TTS run. Do NOT plan to serve from this")
        print("     machine unquantised. Run scripts/optimize_cpu.py — int8 dynamic")
        print("     quantisation typically gives 2-3x on CPU transformers, which would")
        print("     bring P100 inside. Verify retrieval agreement before trusting it.")
    elif head < 20:
        print("  >> TOO TIGHT. Query embedding alone eats the budget on this machine.")
        print("     Options: keep serving on the GPU pod; or cut max_len (192 -> 128);")
        print("     or use a smaller embedder for QUERIES only, keeping bge-m3 for the")
        print("     index — asymmetric encoders are a real technique, but the pair must")
        print("     then be re-evaluated, not assumed compatible.")
    elif head < 80:
        print("  >> WORKABLE but leaves little headroom. Measure the full path below.")
    else:
        print("  >> COMFORTABLE. This machine can serve inside the budget.")

    if args.embed_only:
        out = root / "results" / "latency_embed.json"
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps({"device": dev, "model": model, **t}, indent=2))
        print(f"\n==> wrote {out}")
        return 0

    # ---------------------------------------------------------- full pipeline
    if not man_p.exists():
        raise SystemExit(f"\nno index at {root/'index'} — run src/index_build.py, "
                         f"or use --embed-only for the feasibility check alone")

    from src.guardrails import Guardrails
    from src.harness import AskRequest, Harness
    from src.reader import LexicalSpanReader
    from src.serve import Index
    from src.voice import VoiceStack

    idx = Index(root, langs, None, 0)
    h = Harness(idx, LexicalSpanReader(prior_weight=1.0, answer_mode="sentence"),
                VoiceStack(notes=["bench: TTS off"]), asr=None,
                guards=Guardrails(), budget_ms=args.budget_ms)

    stages: dict = {}
    pipe, answered = [], 0
    for i in range(args.n):
        lg, q = flat[i % len(flat)]
        if lg not in idx.langs:
            continue
        r = h.run(AskRequest(query=q, lang=lg, want_audio=False))
        answered += int(r.answered)
        pipe.append(r.timing["pipeline_ms"])
        for k, v in r.timing.items():
            if k.endswith("_ms") and k not in ("total_ms", "pipeline_ms", "budget_ms"):
                stages.setdefault(k, []).append(v)

    print(f"\n{'='*66}\nFULL PIPELINE — P50 / P70 / P100  (requirement 4)\n{'='*66}")
    print(f"  {'stage':18s}{'n':>6}{'mean':>9}{'P50':>9}{'P70':>9}{'P100':>9}")
    print("  " + "-" * 60)
    rows = []
    for k in sorted(stages, key=lambda k: -st.mean(stages[k])):
        t = table(k, stages[k])
        rows.append(t)
        print(f"  {k:18s}{t['n']:>6}{t['mean']:>9.2f}{t['p50']:>9.2f}"
              f"{t['p70']:>9.2f}{t['p100']:>9.2f}")
    tot = table("PIPELINE", pipe, args.budget_ms)
    rows.append(tot)
    print("  " + "-" * 60)
    print(f"  {'PIPELINE':18s}{tot['n']:>6}{tot['mean']:>9.2f}{tot['p50']:>9.2f}"
          f"{tot['p70']:>9.2f}{tot['p100']:>9.2f}")
    print(f"\n  budget {args.budget_ms:.0f} ms  ->  {tot['within_budget_pct']}% of requests inside")
    print(f"  answered {answered}/{len(pipe)} "
          f"({100*answered/max(1,len(pipe)):.0f}% — the rest were guardrail abstentions)")
    print("\n  NOTE: speech-to-text is a third-party API (requirement 1) and is NOT")
    print("  included above. Report it as its own line, measured against your")
    print("  provider, so the number you control is not confused with the one you")
    print("  do not.")

    out = root / "results" / "latency.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps({
        "device": dev, "model": model, "budget_ms": args.budget_ms,
        "languages": langs, "n_requests": len(pipe),
        "percentiles": "P50/P70/P100 as specified in the brief",
        "embed_query": table("embed_query", e_ms),
        "stages": rows,
        "asr_note": "Sarvam/ElevenLabs round trip excluded — third-party, timed separately",
    }, indent=2))
    print(f"\n==> wrote {out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())