"""Serving-layer benchmark: a scripted conversation through the real pipeline.

Measures what the serving rework is judged on (workspace serving plan,
2026-09-11): per-turn wall time, time-to-first-token in --stream mode, token
usage per turn, and VRAM peak while the conversation runs. Run it against a
backend to produce a baseline, run it again after a serving change, and
compare the JSON artifacts side by side.

The conversation replays full history each turn, exactly like a real client.
Planted facts are prefixed BENCH-MARKER so anything the biographer learns
from a bench run is recognizable in memory stores later.

Usage:
    .venv/bin/python scripts/bench_serving.py                      # dev server
    .venv/bin/python scripts/bench_serving.py --base-url http://localhost:8000
    .venv/bin/python scripts/bench_serving.py --stream             # TTFT mode
    .venv/bin/python scripts/bench_serving.py --interleave-url http://localhost:11434/v1 --interleave-every 20

Exit codes: 0 ok, 69 backend unreachable.
"""

from __future__ import annotations

import argparse
import json
import subprocess
import sys
import threading
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import httpx

TURN_TIMEOUT_S = 300.0

# Labels name the pipeline path each turn is expected to exercise.
CONVERSATION: list[tuple[str, str]] = [
    (
        "plant-fact",
        "For the record: BENCH-MARKER my favorite tea is lapsang souchong. "
        "Please remember that.",
    ),
    ("tool-arithmetic", "What is 61 plus 12?"),
    ("small-talk", "Good evening Tatlock, how are the household systems?"),
    ("memory-recall", "What did I tell you my favorite tea was?"),
    (
        "librarian-routed",
        "What do you know about the desklock device in the living room?",
    ),
    ("tool-arithmetic-2", "Multiply 17 by 23, then add 100."),
    ("synthesis-light", "Suggest a two-line dinner idea using rice."),
    (
        "context-summary",
        "Summarize what we discussed in this conversation in one sentence.",
    ),
]


class VramPoller:
    """Samples nvidia-smi memory.used until stopped; keeps peak and samples."""

    def __init__(self, interval_s: float = 0.5):
        self.interval_s = interval_s
        self.samples_mib: list[int] = []
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)

    def _read_mib(self) -> int | None:
        try:
            out = subprocess.run(
                [
                    "nvidia-smi",
                    "--query-gpu=memory.used",
                    "--format=csv,noheader,nounits",
                ],
                capture_output=True,
                text=True,
                timeout=5,
            )
            return int(out.stdout.strip().splitlines()[0])
        except (OSError, ValueError, IndexError, subprocess.TimeoutExpired):
            return None

    def _run(self) -> None:
        while not self._stop.is_set():
            mib = self._read_mib()
            if mib is not None:
                self.samples_mib.append(mib)
            self._stop.wait(self.interval_s)

    def __enter__(self) -> VramPoller:
        self._thread.start()
        return self

    def __exit__(self, *exc: object) -> None:
        self._stop.set()
        self._thread.join(timeout=5)

    @property
    def peak_mib(self) -> int | None:
        return max(self.samples_mib) if self.samples_mib else None


class Interleaver:
    """Fires a plain completion at a second backend on a fixed cadence.

    Simulates an Open-WebUI-style tenant competing for the same card while
    the main conversation runs — the slot-eviction case the serving plan
    calls out. Each request's wall time is recorded so contention shows on
    both sides of the measurement.
    """

    def __init__(self, url: str, model: str, every_s: float):
        self.url = url.rstrip("/")
        self.model = model
        self.every_s = every_s
        self.results: list[dict[str, Any]] = []
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, daemon=True)

    def _run(self) -> None:
        n = 0
        while not self._stop.is_set():
            n += 1
            started = time.monotonic()
            try:
                r = httpx.post(
                    f"{self.url}/chat/completions",
                    json={
                        "model": self.model,
                        "messages": [
                            {
                                "role": "user",
                                "content": f"Interleave probe {n}: name one European capital.",
                            }
                        ],
                        "max_tokens": 32,
                    },
                    timeout=TURN_TIMEOUT_S,
                )
                self.results.append(
                    {
                        "n": n,
                        "wall_s": round(time.monotonic() - started, 3),
                        "status": r.status_code,
                    }
                )
            except httpx.HTTPError as e:
                self.results.append(
                    {
                        "n": n,
                        "wall_s": round(time.monotonic() - started, 3),
                        "error": type(e).__name__,
                    }
                )
            self._stop.wait(self.every_s)

    def __enter__(self) -> Interleaver:
        self._thread.start()
        return self

    def __exit__(self, *exc: object) -> None:
        self._stop.set()
        self._thread.join(timeout=10)


def run_turn(
    client: httpx.Client, base_url: str, model: str, messages: list[dict[str, str]]
) -> dict[str, Any]:
    """One non-streamed turn: wall time and usage."""
    started = time.monotonic()
    r = client.post(
        f"{base_url}/v1/chat/completions",
        json={"model": model, "messages": messages},
        timeout=TURN_TIMEOUT_S,
    )
    wall = time.monotonic() - started
    r.raise_for_status()
    body = r.json()
    reply = body["choices"][0]["message"]["content"] or ""
    return {
        "wall_s": round(wall, 3),
        "usage": body.get("usage"),
        "reply_chars": len(reply),
        "reply": reply,
    }


def run_turn_streamed(
    client: httpx.Client, base_url: str, model: str, messages: list[dict[str, str]]
) -> dict[str, Any]:
    """One streamed turn: wall time and time-to-first-token."""
    started = time.monotonic()
    ttft: float | None = None
    chunks: list[str] = []
    with client.stream(
        "POST",
        f"{base_url}/v1/chat/completions",
        json={"model": model, "messages": messages, "stream": True},
        timeout=TURN_TIMEOUT_S,
    ) as r:
        r.raise_for_status()
        for line in r.iter_lines():
            if not line.startswith("data: "):
                continue
            payload = line[len("data: ") :]
            if payload.strip() == "[DONE]":
                break
            try:
                delta = json.loads(payload)["choices"][0]["delta"]
            except (json.JSONDecodeError, KeyError, IndexError):
                continue
            content = delta.get("content")
            if content:
                if ttft is None:
                    ttft = time.monotonic() - started
                chunks.append(content)
    reply = "".join(chunks)
    return {
        "wall_s": round(time.monotonic() - started, 3),
        "ttft_s": round(ttft, 3) if ttft is not None else None,
        "reply_chars": len(reply),
        "reply": reply,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--base-url", default="http://localhost:8777")
    parser.add_argument(
        "--model",
        default="Tatlock",
        help="model id as the target serves it (tatlock registry: 'Tatlock'; raw backend: e.g. 'gemma4:e2b')",
    )
    parser.add_argument("--turns", type=int, default=len(CONVERSATION))
    parser.add_argument(
        "--stream",
        action="store_true",
        help="measure time-to-first-token instead of usage",
    )
    parser.add_argument(
        "--interleave-url",
        default=None,
        help="OpenAI-compat base (e.g. http://localhost:11434/v1) for a competing tenant",
    )
    parser.add_argument("--interleave-every", type=float, default=20.0)
    parser.add_argument(
        "--interleave-model",
        default="gemma4:e2b",
        help="model id at the interleave backend (raw Ollama/llama-server tag)",
    )
    parser.add_argument("--label", default=None, help="run label stored in the artifact")
    parser.add_argument("--out-dir", default="build/bench")
    args = parser.parse_args()

    base_url = args.base_url.rstrip("/")
    client = httpx.Client()

    try:
        health = client.get(f"{base_url}/health", timeout=10).json()
    except httpx.HTTPError as e:
        print(f"backend unreachable at {base_url}: {type(e).__name__}", file=sys.stderr)
        return 69

    started_at = datetime.now(UTC).isoformat()
    turns: list[dict[str, Any]] = []
    messages: list[dict[str, str]] = []

    interleaver = (
        Interleaver(args.interleave_url, args.interleave_model, args.interleave_every)
        if args.interleave_url
        else None
    )

    with VramPoller() as vram:
        ctx = interleaver if interleaver else _NullCtx()
        with ctx:
            for label, text in CONVERSATION[: args.turns]:
                messages.append({"role": "user", "content": text})
                runner = run_turn_streamed if args.stream else run_turn
                try:
                    result = runner(client, base_url, args.model, messages)
                except httpx.HTTPStatusError as e:
                    detail = e.response.text[:200]
                    print(
                        f"  {label}: FAILED {e.response.status_code} {detail}",
                        file=sys.stderr,
                    )
                    turns.append(
                        {
                            "label": label,
                            "error": type(e).__name__,
                            "status": e.response.status_code,
                            "detail": detail,
                        }
                    )
                    messages.pop()
                    continue
                except httpx.HTTPError as e:
                    print(f"  {label}: FAILED ({type(e).__name__})", file=sys.stderr)
                    turns.append({"label": label, "error": type(e).__name__})
                    messages.pop()
                    continue
                messages.append({"role": "assistant", "content": result.pop("reply")})
                turns.append({"label": label, **result})
                shown = result.get("ttft_s") if args.stream else result.get("wall_s")
                print(f"  {label}: {shown}s  ({result})")

    artifact = {
        "label": args.label,
        "started_at": started_at,
        "base_url": base_url,
        "model": args.model,
        "mode": "stream" if args.stream else "usage",
        "backend_health": health,
        "interleave": {
            "url": args.interleave_url,
            "every_s": args.interleave_every,
            "results": interleaver.results if interleaver else [],
        },
        "vram": {"peak_mib": vram.peak_mib, "samples": len(vram.samples_mib)},
        "turns": turns,
    }

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    stamp = started_at.replace(":", "").replace("-", "")[:15]
    out_path = out_dir / f"bench-{stamp}-{args.label or 'run'}.json"
    out_path.write_text(json.dumps(artifact, indent=2))

    ok_walls = [t["wall_s"] for t in turns if "wall_s" in t]
    print(f"\nturns: {len(ok_walls)}/{len(turns)} ok")
    if ok_walls:
        print(f"wall  total {sum(ok_walls):.1f}s  mean {sum(ok_walls) / len(ok_walls):.1f}s")
    if args.stream:
        ttfts = [t["ttft_s"] for t in turns if t.get("ttft_s") is not None]
        if ttfts:
            print(f"ttft  mean {sum(ttfts) / len(ttfts):.2f}s  max {max(ttfts):.2f}s")
    print(f"vram  peak {vram.peak_mib} MiB over {len(vram.samples_mib)} samples")
    print(f"artifact: {out_path}")
    return 0


class _NullCtx:
    def __enter__(self) -> _NullCtx:
        return self

    def __exit__(self, *exc: object) -> None:
        return None


if __name__ == "__main__":
    sys.exit(main())