#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10"
# dependencies = ["httpx>=0.27"]
# ///
"""Benchmark Ollama models for FabledScribe chat / curator workloads.

Measures time-to-first-token, total wall time, and tokens/sec for each
(model, scenario) pair. Designed to be runnable against both local and
remote Ollama servers so results from your two CPU servers are directly
comparable.

Default forces CPU-only inference (num_gpu=0). Pass --num-gpu 99 for full
GPU offload, or any integer for partial.

Scenarios:
    chat     — small input, small output, no thinking. Mirrors the chat-only
               journal companion's expected load (short user message →
               curious follow-up).
    curator  — longer transcript input, structured-output extraction.
               Thinking defaults ON (qwen3 family) but is overridable via
               --think for cross-family comparisons.

Think modes:
    auto     — chat=off, curator=on (default; matches scenario intent).
    off      — force think disabled for all runs (fair cross-family
               comparison; non-qwen3 models silently ignore the field
               anyway, so this aligns behaviour explicitly).
    on       — force think enabled for all runs (measures what think
               contributes vs `off` on the same model).

Prerequisites:
    - `uv` installed (https://docs.astral.sh/uv/). The script's shebang
      uses `uv run --script`, which auto-creates an ephemeral venv with
      httpx — no project install required.
    - Ollama running and reachable at --server (default http://localhost:11434).
    - The models named in --models must already be pulled
      (`ollama pull qwen3:30b-a3b` etc).

Running:
    # Preferred — uv handles the venv:
    uv run scripts/bench_ollama.py --models qwen3:30b-a3b --scenario curator
    # Or via the shebang (the file is executable):
    ./scripts/bench_ollama.py --models qwen3:30b-a3b --scenario curator
    # `python scripts/bench_ollama.py …` works ONLY if httpx is on the
    # interpreter's path — usually it isn't outside the project venv.

Usage examples:
    # Curator candidate on CPU, 3 runs each (default), one model:
    uv run scripts/bench_ollama.py --models qwen3:30b-a3b --scenario curator

    # Chat candidate on GPU, against a remote server:
    uv run scripts/bench_ollama.py \\
        --server http://dock02:11434 \\
        --models llama3.2:3b \\
        --scenario chat --num-gpu 99

    # Compare two qwen curator candidates on CPU, 5 runs each:
    uv run scripts/bench_ollama.py \\
        --models qwen3:30b-a3b,qwen3:32b \\
        --scenario curator --runs 5 --out bench-qwen-curator.md

    # Cross-family curator comparison with think forced OFF (apples-to-apples):
    uv run scripts/bench_ollama.py \\
        --models qwen3:30b-a3b,gemma2:27b,mistral-small:22b \\
        --scenario curator --think off --runs 3 \\
        --out bench-curator-no-think.md

The first request to a (model, num_gpu) pair triggers a model load and is
excluded from the timing as a warm-up. Subsequent runs reflect warm-cache
performance — closest to how the running app will see the model in practice.
"""

from __future__ import annotations

import argparse
import json
import statistics
import sys
import time
from dataclasses import dataclass, field

import httpx


SYSTEM_PROMPT_CHAT = (
    "You are the user's journal companion. Be warm and curious. Ask one "
    "short follow-up question. Don't presume to help; let the user lead. "
    "Match the user's length."
)

USER_MESSAGE_CHAT = (
    "Got a lot done at Famous Supply today. Finished the network restage "
    "on the Bedford swarm and there was one weird issue with a Windows "
    "node. Tired but it feels good to be making progress."
)

SYSTEM_PROMPT_CURATOR = (
    "You are reading a journal conversation. Extract the user's notable "
    "moments — events, decisions, encounters, observations, plans. "
    "Output a JSON array. Each entry has fields: content (1-2 sentence "
    "summary in the user's voice), tags (list of strings), people (list "
    "of names mentioned), projects (list of project names mentioned). "
    "Skip meta-conversational chatter."
)

USER_TRANSCRIPT_CURATOR = """\
[12:04] User: morning. starting on the Famous Supply network restage today.
[12:04] Assistant: how big a job is that?
[12:05] User: bedford location swarm has 3 nodes. one of them is a Windows
        host that's been flaky for a while. probably going to take all
        afternoon.
[12:05] Assistant: anything you want to flag before you start in case
        something goes sideways?
[12:06] User: yeah — the Windows node sometimes loses its bridge network
        on reboot. if that happens i'll have to manually rejoin it. last
        time it took an hour to figure out.
[12:06] Assistant: noted. how are you feeling going in?
[12:07] User: decent. didn't sleep great but i had coffee with Sarah this
        morning which was nice. she mentioned she's hiring on her team if
        i ever want to make a jump.
[12:07] Assistant: good to know. anything to capture about the work itself
        before you dive in, or you good?
[12:08] User: nope, just wanted to externalize before starting.
[14:42] User: ok the Windows node DID lose its bridge. fixing it now.
[14:42] Assistant: any guess what triggers it?
[14:43] User: i think it's a stale dhcp lease that gets re-issued before
        the bridge comes up. might write that down as a thing to fix
        properly later.
[15:55] User: done. all three nodes restaged, the windows one too. wrote
        a note about the dhcp issue for future me.
[15:55] Assistant: how do you feel about how the day went?
[15:56] User: tired but good. ate lunch standing up which was dumb.
        watching a show with Victoria tonight to decompress.
"""


@dataclass
class RunResult:
    ttft_ms: float
    total_ms: float
    prompt_tokens: int
    output_tokens: int
    tokens_per_sec: float


@dataclass
class ScenarioResult:
    model: str
    scenario: str
    runs: list[RunResult] = field(default_factory=list)
    error: str | None = None

    def summary(self) -> dict:
        if self.error or not self.runs:
            return {
                "model": self.model,
                "scenario": self.scenario,
                "error": self.error or "no successful runs",
            }
        return {
            "model": self.model,
            "scenario": self.scenario,
            "runs": len(self.runs),
            "ttft_ms_p50": statistics.median(r.ttft_ms for r in self.runs),
            "total_ms_p50": statistics.median(r.total_ms for r in self.runs),
            "tokens_per_sec_p50": statistics.median(
                r.tokens_per_sec for r in self.runs
            ),
            "output_tokens_mean": statistics.mean(
                r.output_tokens for r in self.runs
            ),
            "prompt_tokens": self.runs[0].prompt_tokens,
        }


def _resolve_think(scenario: str, think_mode: str) -> bool:
    """Map (scenario, --think mode) → boolean think flag."""
    if think_mode == "on":
        return True
    if think_mode == "off":
        return False
    # auto: scenario-driven
    return scenario == "curator"


def build_request(
    scenario: str,
    model: str,
    num_gpu: int,
    keep_alive: str,
    think_mode: str = "auto",
) -> dict:
    if scenario == "chat":
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT_CHAT},
            {"role": "user", "content": USER_MESSAGE_CHAT},
        ]
    elif scenario == "curator":
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT_CURATOR},
            {"role": "user", "content": USER_TRANSCRIPT_CURATOR},
        ]
    else:
        raise ValueError(f"unknown scenario: {scenario}")
    think = _resolve_think(scenario, think_mode)
    return {
        "model": model,
        "messages": messages,
        "stream": True,
        "think": think,
        "keep_alive": keep_alive,
        "options": {
            "num_gpu": num_gpu,
            "temperature": 0.3,
            "num_ctx": 8192,
        },
    }


def run_once(server: str, payload: dict) -> RunResult:
    """Stream one chat request and time it.

    Uses Ollama-reported `eval_count` and `eval_duration` for tokens/sec
    (authoritative; doesn't include client-side stream overhead). TTFT is
    wall-clock from request send to first content chunk.
    """
    url = f"{server.rstrip('/')}/api/chat"
    t_start = time.monotonic()
    ttft = None
    prompt_tokens = 0
    output_tokens = 0
    eval_duration_ns = 0

    with httpx.stream("POST", url, json=payload, timeout=600.0) as resp:
        resp.raise_for_status()
        for line in resp.iter_lines():
            if not line:
                continue
            chunk = json.loads(line)
            if ttft is None and chunk.get("message", {}).get("content"):
                ttft = time.monotonic() - t_start
            if chunk.get("done"):
                prompt_tokens = chunk.get("prompt_eval_count", 0)
                output_tokens = chunk.get("eval_count", 0)
                eval_duration_ns = chunk.get("eval_duration", 0)

    total = time.monotonic() - t_start
    tps = (
        output_tokens / (eval_duration_ns / 1e9)
        if eval_duration_ns
        else 0.0
    )
    return RunResult(
        ttft_ms=(ttft if ttft is not None else total) * 1000,
        total_ms=total * 1000,
        prompt_tokens=prompt_tokens,
        output_tokens=output_tokens,
        tokens_per_sec=tps,
    )


def benchmark(
    *,
    server: str,
    models: list[str],
    scenarios: list[str],
    runs: int,
    num_gpu: int,
    keep_alive: str,
    think_mode: str,
) -> list[ScenarioResult]:
    results: list[ScenarioResult] = []
    for model in models:
        for scenario in scenarios:
            sr = ScenarioResult(model=model, scenario=scenario)
            payload = build_request(
                scenario, model, num_gpu, keep_alive, think_mode
            )
            # Warm-up run loads the model into RAM/VRAM with the requested
            # num_gpu setting. Excluded from the measured runs because it
            # otherwise dominates TTFT with model-load time.
            print(
                f"[{model} :: {scenario}] warm-up (loading model)...",
                flush=True,
            )
            try:
                run_once(server, payload)
            except httpx.HTTPError as e:
                sr.error = f"warm-up failed: {e}"
                print(f"  {sr.error}", file=sys.stderr)
                results.append(sr)
                continue
            except Exception as e:
                sr.error = f"warm-up exception: {e}"
                print(f"  {sr.error}", file=sys.stderr)
                results.append(sr)
                continue
            for i in range(runs):
                try:
                    r = run_once(server, payload)
                except Exception as e:
                    print(f"  run {i+1} failed: {e}", file=sys.stderr)
                    continue
                sr.runs.append(r)
                print(
                    f"  run {i+1}/{runs}: ttft={r.ttft_ms:.0f}ms "
                    f"total={r.total_ms:.0f}ms tps={r.tokens_per_sec:.1f} "
                    f"out_tokens={r.output_tokens}",
                    flush=True,
                )
            results.append(sr)
    return results


def format_markdown(
    results: list[ScenarioResult],
    *,
    server: str,
    num_gpu: int,
    think_mode: str,
) -> str:
    mode = "CPU only" if num_gpu == 0 else (
        f"GPU offload ({num_gpu} layers)" if num_gpu > 0 else "Ollama default"
    )
    think_label = {
        "auto": "auto (chat=off, curator=on)",
        "on": "forced ON",
        "off": "forced OFF",
    }.get(think_mode, think_mode)
    lines = [
        "# Ollama benchmark",
        "",
        f"- Server: `{server}`",
        f"- Hardware mode: {mode} (`num_gpu={num_gpu}`)",
        f"- Think: {think_label}",
        "",
        "| Model | Scenario | Runs | Prompt tok | TTFT p50 (ms) "
        "| Total p50 (ms) | tok/s p50 | Output tok (mean) |",
        "|---|---|---|---|---|---|---|---|",
    ]
    for sr in results:
        s = sr.summary()
        if "error" in s:
            lines.append(
                f"| {s['model']} | {s['scenario']} | — | — | — | — | — "
                f"| error: {s['error']} |"
            )
            continue
        lines.append(
            f"| {s['model']} | {s['scenario']} | {s['runs']} "
            f"| {s['prompt_tokens']} "
            f"| {s['ttft_ms_p50']:.0f} | {s['total_ms_p50']:.0f} "
            f"| {s['tokens_per_sec_p50']:.1f} "
            f"| {s['output_tokens_mean']:.0f} |"
        )
    return "\n".join(lines) + "\n"


def main():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--server",
        default="http://localhost:11434",
        help="Ollama server URL (default %(default)s)",
    )
    parser.add_argument(
        "--models",
        required=True,
        help="Comma-separated model tags (e.g. qwen2.5:32b,qwen3:14b)",
    )
    parser.add_argument(
        "--scenario",
        choices=["chat", "curator", "both"],
        default="both",
    )
    parser.add_argument(
        "--runs",
        type=int,
        default=3,
        help="Runs per (model,scenario), excluding warm-up (default %(default)s)",
    )
    parser.add_argument(
        "--num-gpu",
        type=int,
        default=0,
        help="0 = CPU only (default), 99 = full offload, -1 = Ollama default",
    )
    parser.add_argument(
        "--keep-alive",
        default="10m",
        help="Ollama keep_alive (default %(default)s)",
    )
    parser.add_argument(
        "--think",
        choices=["auto", "on", "off"],
        default="auto",
        help=(
            "Think-mode control. auto = chat off / curator on (default). "
            "off = force disabled for fair cross-family comparison "
            "(non-qwen3 models silently ignore think anyway). "
            "on = force enabled to measure what think contributes."
        ),
    )
    parser.add_argument(
        "--out",
        help="Write markdown table to this file (also prints to stdout)",
    )
    args = parser.parse_args()

    models = [m.strip() for m in args.models.split(",") if m.strip()]
    scenarios = (
        ["chat", "curator"] if args.scenario == "both" else [args.scenario]
    )

    results = benchmark(
        server=args.server,
        models=models,
        scenarios=scenarios,
        runs=args.runs,
        num_gpu=args.num_gpu,
        keep_alive=args.keep_alive,
        think_mode=args.think,
    )

    md = format_markdown(
        results,
        server=args.server,
        num_gpu=args.num_gpu,
        think_mode=args.think,
    )
    print("\n" + md)
    if args.out:
        with open(args.out, "w") as f:
            f.write(md)
        print(f"Wrote results to {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()