#!/usr/bin/env python3
from __future__ import annotations

import concurrent.futures
import csv
import json
import math
import os
import re
import statistics
import subprocess
import threading
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path

OUTPUT_DIR = Path(
    os.environ.get(
        "OUTPUT_DIR",
        "/Users/nation/Documents/New project/artifacts/spark-open-model-exhaustive-20260727T190000Z",
    )
)
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

SPARK_HOST = os.environ.get("SPARK_HOST", "192.168.1.121").strip()
SPARK_USER = os.environ.get("SPARK_USER", "perknation").strip()
SPARK_SSH_KEY = os.environ.get(
    "SPARK_SSH_KEY",
    str(Path.home() / "Library/Application Support/NVIDIA/Sync/config/nvsync.key"),
).strip()
SAMPLE_INTERVAL_S = float(os.environ.get("SAMPLE_INTERVAL_S", "2.0"))
WARMUP_TIMEOUT_S = int(os.environ.get("WARMUP_TIMEOUT_S", "1800"))
REQUEST_TIMEOUT_S = int(os.environ.get("REQUEST_TIMEOUT_S", "900"))
RUN_DUAL_PROBE = os.environ.get("RUN_DUAL_PROBE", "1").strip().lower() not in {"0", "false", "no"}

LONG_CONTEXT_TEXT = """
Yeti Claw is a public inference surface that fronts several local model lanes. The operator cares about three things:
reliable response delivery, honest capacity guidance, and predictable model behavior under real user concurrency.

Spark serves the heaviest text and image workloads. The fastest public text lane is the TensorRT-hosted Nemotron path.
Two additional open-weight text models sit behind Ollama on the same Spark box: GLM and Kimi. They are useful for comparison,
but the operator needs to understand whether those lanes stay practical once more than one user is active at the same time.

Users do not see kernel traces, swap pressure, or runner load delays. They only see whether the chat box answers quickly,
whether the answer appears at all, and whether the system feels confident and stable. That means the benchmark needs to
measure both front-end consequences and host-side environmental conditions.

The public Yeti Claw site also doubles as an editorial archive. Mission Control articles are expected to publish not just
winners, but tradeoffs, failure cases, and operator recommendations. A model that is technically impressive but operationally
fragile should not be presented the same way as a model that is modest but reliable.

For this benchmark, the committee wants data that can support product positioning. The report needs to say which model is
best as the public default, which model belongs in a secondary comparison slot, and which one should be treated as a specialty
lane with lower concurrency expectations. The report also needs a clean explanation of what happens if two heavy Ollama-backed
models are asked to cold-load at the same time.

The benchmark should not rely on a single toy prompt. One prompt should test short operational chat behavior, another should
test structured JSON output for integration safety, and another should test longer context digestion without crossing the model
context window. The operator also wants the benchmark to be reproducible and the artifacts easy to publish into a Mission
Control article with charts, tables, and a downloadable PDF.
""".strip()


@dataclass(frozen=True)
class ModelConfig:
    key: str
    label: str
    api_kind: str
    endpoint: str
    model: str
    family: str


@dataclass(frozen=True)
class Scenario:
    key: str
    label: str
    messages: list[dict[str, str]]
    max_tokens: int
    levels: list[int]
    min_requests: int
    request_multiplier: int
    num_ctx: int = 4096
    timeout_s: int = REQUEST_TIMEOUT_S


MODELS = {
    "nemotron_super": ModelConfig(
        key="nemotron_super",
        label="Nemotron Super",
        api_kind="openai",
        endpoint=f"http://{SPARK_HOST}:8001/v1/chat/completions",
        model="nvidia/nemotron-3-super",
        family="Nemotron",
    ),
    "glm45air_q6": ModelConfig(
        key="glm45air_q6",
        label="GLM-4.5-Air Q6_K",
        api_kind="ollama",
        endpoint=f"http://{SPARK_HOST}:11434/api/chat",
        model="glm45air-q6",
        family="GLM",
    ),
    "kimi_dev72_q40": ModelConfig(
        key="kimi_dev72_q40",
        label="Kimi Dev 72B Q4_0",
        api_kind="ollama",
        endpoint=f"http://{SPARK_HOST}:11434/api/chat",
        model="hf.co/bartowski/moonshotai_Kimi-Dev-72B-GGUF:Q4_0",
        family="Kimi",
    ),
}

SCENARIOS = [
    Scenario(
        key="short_ops",
        label="Short Ops Reply",
        messages=[
            {
                "role": "system",
                "content": "You are a concise operations assistant. Reply in plain text with exactly three short bullet points.",
            },
            {
                "role": "user",
                "content": (
                    "Give three practical tactics for keeping a public AI chat lane responsive when several users arrive at once."
                ),
            },
        ],
        max_tokens=128,
        levels=[1, 2, 4],
        min_requests=6,
        request_multiplier=3,
    ),
    Scenario(
        key="json_guardrails",
        label="Structured JSON",
        messages=[
            {
                "role": "system",
                "content": (
                    "Return valid JSON only. No prose outside the JSON object. "
                    "The object must contain keys risk, recommendation, and operator_note."
                ),
            },
            {
                "role": "user",
                "content": (
                    "Assess a public model lane that starts swapping during peak traffic. "
                    "Provide a low, medium, or high risk rating, a one-sentence recommendation, "
                    "and one short operator note."
                ),
            },
        ],
        max_tokens=96,
        levels=[1, 2, 4],
        min_requests=4,
        request_multiplier=2,
    ),
    Scenario(
        key="long_context",
        label="Long Context Synthesis",
        messages=[
            {
                "role": "system",
                "content": "Reply in four short bullet points and keep the answer strictly focused on operator decision-making.",
            },
            {
                "role": "user",
                "content": (
                    "Read the following internal benchmark brief and summarize the operator takeaways.\n\n"
                    f"{LONG_CONTEXT_TEXT}\n\n"
                    "Now provide four bullet points: best public default, best secondary lane, biggest risk, and deployment guidance."
                ),
            },
        ],
        max_tokens=144,
        levels=[1, 2],
        min_requests=4,
        request_multiplier=2,
    ),
]


def ssh_target() -> str:
    return f"{SPARK_USER}@{SPARK_HOST}"


def ssh_command(*remote_parts: str) -> list[str]:
    return [
        "ssh",
        "-i",
        SPARK_SSH_KEY,
        "-o",
        "BatchMode=yes",
        "-o",
        "ConnectTimeout=8",
        "-o",
        "StrictHostKeyChecking=no",
        ssh_target(),
        *remote_parts,
    ]


def ssh_run(remote: str, timeout: int = 60) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        ssh_command(remote),
        capture_output=True,
        text=True,
        timeout=timeout,
        check=False,
    )


def parse_loadavg(text: str) -> tuple[float | None, float | None, float | None]:
    match = re.search(r"load average:?\s*([0-9.]+)[, ]+([0-9.]+)[, ]+([0-9.]+)", text, re.I)
    if not match:
        return (None, None, None)
    return tuple(float(match.group(i)) for i in range(1, 4))


def fmt_num(value: float | int | None, digits: int = 2) -> str:
    if value is None:
        return "n/a"
    return f"{value:.{digits}f}"


def sample_host(model_key: str, scenario_key: str = "host") -> dict:
    sample = {"timestamp": time.time(), "model_key": model_key, "scenario_key": scenario_key}
    remote = (
        "printf 'UPTIME\\n'; uptime; "
        "printf 'FREE\\n'; free -m | sed -n '1,3p'; "
        "printf 'SWAP\\n'; free -m | sed -n '4p'; "
        "printf 'TOP\\n'; top -bn1 | sed -n '1,5p'; "
        "printf 'OLLAMA\\n'; ollama ps 2>/dev/null || true; "
        "printf 'GPU\\n'; "
        "nvidia-smi --query-gpu=temperature.gpu,utilization.gpu,power.draw "
        "--format=csv,noheader,nounits 2>/dev/null || true"
    )
    try:
        res = ssh_run(remote, timeout=25)
        if res.returncode != 0:
            raise RuntimeError(res.stderr.strip() or f"ssh rc={res.returncode}")

        sections = {"UPTIME": "", "FREE": "", "SWAP": "", "TOP": "", "OLLAMA": "", "GPU": ""}
        current = None
        for line in res.stdout.splitlines():
            tag = line.strip()
            if tag in sections:
                current = tag
                continue
            if current:
                sections[current] += line + "\n"

        l1, l5, l15 = parse_loadavg(sections["UPTIME"])
        sample.update({"load1": l1, "load5": l5, "load15": l15})

        free_match = re.search(r"Mem:\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(\d+)", sections["FREE"])
        if free_match:
            sample["mem_total_mb"] = int(free_match.group(1))
            sample["mem_used_mb"] = int(free_match.group(2))
            sample["mem_free_mb"] = int(free_match.group(3))
            sample["mem_available_mb"] = int(free_match.group(6))

        swap_match = re.search(r"Swap:\s+(\d+)\s+(\d+)\s+(\d+)", sections["SWAP"])
        if swap_match:
            sample["swap_total_mb"] = int(swap_match.group(1))
            sample["swap_used_mb"] = int(swap_match.group(2))
            sample["swap_free_mb"] = int(swap_match.group(3))

        for line in sections["TOP"].splitlines():
            if line.startswith("%Cpu(s):"):
                idle_match = re.search(r"([0-9.]+)\s+id", line)
                if idle_match:
                    idle = float(idle_match.group(1))
                    sample["cpu_idle_pct"] = idle
                    sample["cpu_busy_pct"] = max(0.0, 100.0 - idle)

        sample["ollama_ps"] = sections["OLLAMA"].strip()
        gpu_line = sections["GPU"].strip().splitlines()[0] if sections["GPU"].strip() else ""
        gpu_parts = [part.strip() for part in gpu_line.split(",")] if gpu_line else []
        if len(gpu_parts) >= 3:
            try:
                sample["gpu_temp_c"] = float(gpu_parts[0])
            except ValueError:
                pass
            try:
                sample["gpu_util_pct"] = float(gpu_parts[1])
            except ValueError:
                pass
            power_token = gpu_parts[2].split()[0]
            try:
                sample["gpu_power_w"] = float(power_token)
            except ValueError:
                pass
    except Exception as exc:
        sample["sample_error"] = str(exc)
    return sample


class Sampler(threading.Thread):
    def __init__(self, model_key: str, scenario_key: str, interval_s: float, sink: list[dict]):
        super().__init__(daemon=True)
        self.model_key = model_key
        self.scenario_key = scenario_key
        self.interval_s = interval_s
        self.sink = sink
        self.stop_event = threading.Event()

    def run(self) -> None:
        while not self.stop_event.is_set():
            self.sink.append(sample_host(self.model_key, self.scenario_key))
            self.stop_event.wait(self.interval_s)

    def stop(self) -> None:
        self.stop_event.set()


def build_request_body(config: ModelConfig, scenario: Scenario) -> dict:
    if config.api_kind == "openai":
        return {
            "model": config.model,
            "messages": scenario.messages,
            "temperature": 0,
            "max_tokens": scenario.max_tokens,
            "chat_template_kwargs": {
                "enable_thinking": False,
                "force_nonempty_content": True,
            },
        }
    return {
        "model": config.model,
        "messages": scenario.messages,
        "stream": False,
        "options": {
            "temperature": 0,
            "num_predict": scenario.max_tokens,
            "num_ctx": scenario.num_ctx,
        },
    }


def extract_content(config: ModelConfig, payload: dict) -> str:
    if config.api_kind == "openai":
        choices = payload.get("choices")
        if isinstance(choices, list) and choices:
            message = choices[0].get("message")
            if isinstance(message, dict):
                return str(message.get("content") or "").strip()
        return ""
    message = payload.get("message")
    if isinstance(message, dict):
        return str(message.get("content") or "").strip()
    return str(payload.get("response") or "").strip()


def extract_usage(config: ModelConfig, payload: dict) -> dict[str, float | int | None]:
    usage: dict[str, float | int | None] = {
        "prompt_tokens": None,
        "completion_tokens": None,
        "total_tokens": None,
        "load_duration_s": None,
        "prompt_eval_duration_s": None,
        "eval_duration_s": None,
    }
    if config.api_kind == "openai":
        raw = payload.get("usage") or {}
        usage["prompt_tokens"] = raw.get("prompt_tokens")
        usage["completion_tokens"] = raw.get("completion_tokens")
        usage["total_tokens"] = raw.get("total_tokens")
        return usage

    prompt_tokens = payload.get("prompt_eval_count")
    completion_tokens = payload.get("eval_count")
    total_tokens = None
    if isinstance(prompt_tokens, int) and isinstance(completion_tokens, int):
        total_tokens = prompt_tokens + completion_tokens
    usage["prompt_tokens"] = prompt_tokens
    usage["completion_tokens"] = completion_tokens
    usage["total_tokens"] = total_tokens
    for key in ("load_duration", "prompt_eval_duration", "eval_duration"):
        value = payload.get(key)
        if isinstance(value, (int, float)):
            usage[f"{key}_s"] = round(float(value) / 1_000_000_000, 4)
    return usage


def request_model(
    config: ModelConfig,
    scenario: Scenario,
    request_id: str,
    timeout_s: int | None = None,
) -> dict:
    body = build_request_body(config, scenario)
    req = urllib.request.Request(
        config.endpoint,
        data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json"},
    )
    started_at = time.time()
    start_mono = time.monotonic()
    timeout_s = timeout_s or scenario.timeout_s
    try:
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            payload = json.loads(resp.read().decode("utf-8"))
        latency_s = time.monotonic() - start_mono
        content = extract_content(config, payload)
        usage = extract_usage(config, payload)
        completion_tokens = usage.get("completion_tokens")
        tok_s = None
        if isinstance(completion_tokens, (int, float)) and latency_s > 0:
            tok_s = float(completion_tokens) / latency_s
        row = {
            "request_id": request_id,
            "model_key": config.key,
            "model_label": config.label,
            "model_family": config.family,
            "model": config.model,
            "scenario_key": scenario.key,
            "scenario_label": scenario.label,
            "ok": bool(content),
            "started_at": started_at,
            "latency_s": latency_s,
            "content_preview": content[:240],
            "payload_preview": json.dumps(payload)[:600],
            "tokens_per_s": tok_s,
            **usage,
        }
        return row
    except urllib.error.HTTPError as err:
        detail = err.read().decode("utf-8", "replace") if hasattr(err, "read") else str(err)
        return {
            "request_id": request_id,
            "model_key": config.key,
            "model_label": config.label,
            "model_family": config.family,
            "model": config.model,
            "scenario_key": scenario.key,
            "scenario_label": scenario.label,
            "ok": False,
            "started_at": started_at,
            "latency_s": time.monotonic() - start_mono,
            "error": f"HTTP {err.code}: {detail[:600]}",
        }
    except Exception as exc:
        return {
            "request_id": request_id,
            "model_key": config.key,
            "model_label": config.label,
            "model_family": config.family,
            "model": config.model,
            "scenario_key": scenario.key,
            "scenario_label": scenario.label,
            "ok": False,
            "started_at": started_at,
            "latency_s": time.monotonic() - start_mono,
            "error": str(exc),
        }


def quantile(values: list[float], p: float) -> float | None:
    if not values:
        return None
    if len(values) == 1:
        return values[0]
    idx = (len(values) - 1) * p
    lo = math.floor(idx)
    hi = math.ceil(idx)
    if lo == hi:
        return values[lo]
    return values[lo] + (values[hi] - values[lo]) * (idx - lo)


def summarize_step(
    config: ModelConfig,
    scenario: Scenario,
    concurrency: int,
    requests: list[dict],
    samples: list[dict],
    wall_s: float,
) -> dict:
    ok = [r for r in requests if r.get("ok")]
    latencies = sorted(r["latency_s"] for r in ok if isinstance(r.get("latency_s"), (int, float)))
    completion_tokens = [float(r["completion_tokens"]) for r in ok if isinstance(r.get("completion_tokens"), (int, float))]
    prompt_tokens = [float(r["prompt_tokens"]) for r in ok if isinstance(r.get("prompt_tokens"), (int, float))]
    token_s_values = [float(r["tokens_per_s"]) for r in ok if isinstance(r.get("tokens_per_s"), (int, float))]
    summary = {
        "model_key": config.key,
        "model_label": config.label,
        "model_family": config.family,
        "scenario_key": scenario.key,
        "scenario_label": scenario.label,
        "concurrency": concurrency,
        "request_count": len(requests),
        "success_count": len(ok),
        "failure_count": len(requests) - len(ok),
        "throughput_rps": (len(ok) / wall_s) if wall_s > 0 else None,
        "latency_avg_s": statistics.mean(latencies) if latencies else None,
        "latency_p50_s": quantile(latencies, 0.5) if latencies else None,
        "latency_p95_s": quantile(latencies, 0.95) if latencies else None,
        "latency_max_s": max(latencies) if latencies else None,
        "wall_time_s": wall_s,
        "avg_completion_tokens": statistics.mean(completion_tokens) if completion_tokens else None,
        "avg_prompt_tokens": statistics.mean(prompt_tokens) if prompt_tokens else None,
        "avg_tokens_per_s": statistics.mean(token_s_values) if token_s_values else None,
        "aggregate_tokens_per_s": (sum(completion_tokens) / wall_s) if completion_tokens and wall_s > 0 else None,
    }
    for key in (
        "cpu_busy_pct",
        "load1",
        "mem_used_mb",
        "mem_available_mb",
        "swap_used_mb",
        "gpu_temp_c",
        "gpu_util_pct",
        "gpu_power_w",
    ):
        values = [s.get(key) for s in samples if isinstance(s.get(key), (int, float))]
        if not values:
            continue
        summary[f"peak_{key}"] = max(values)
        summary[f"min_{key}"] = min(values)
        summary[f"avg_{key}"] = statistics.mean(values)
    return summary


def stop_loaded_ollama_models(except_model: str | None = None) -> None:
    list_res = ssh_run("ollama ps 2>/dev/null | awk 'NR>1 {print $1}'", timeout=20)
    if list_res.returncode != 0:
        return
    models = [line.strip() for line in list_res.stdout.splitlines() if line.strip()]
    for model in models:
        if except_model and model == except_model:
            continue
        ssh_run(f"ollama stop {json.dumps(model)} >/dev/null 2>&1 || true", timeout=20)


def prepare_model(config: ModelConfig) -> None:
    if config.api_kind == "ollama":
        stop_loaded_ollama_models()


def cleanup_model(config: ModelConfig) -> None:
    if config.api_kind == "ollama":
        ssh_run(f"ollama stop {json.dumps(config.model)} >/dev/null 2>&1 || true", timeout=20)


def warm_model(config: ModelConfig, scenario: Scenario) -> dict:
    return request_model(config, scenario, f"{config.key}-{scenario.key}-warmup", timeout_s=WARMUP_TIMEOUT_S)


def run_model(config: ModelConfig) -> dict:
    prepare_model(config)
    warmup = warm_model(config, SCENARIOS[0])
    request_rows = [warmup]
    sample_rows: list[dict] = []
    scenario_results = []
    if not warmup.get("ok"):
        cleanup_model(config)
        return {
            "label": config.label,
            "model": config.model,
            "warmup": warmup,
            "scenarios": scenario_results,
            "requests": request_rows,
            "samples": sample_rows,
        }

    for scenario in SCENARIOS:
        steps = []
        for concurrency in scenario.levels:
            total_requests = max(concurrency * scenario.request_multiplier, scenario.min_requests)
            requests_for_step: list[dict] = []
            samples_for_step: list[dict] = []
            sampler = Sampler(config.key, scenario.key, SAMPLE_INTERVAL_S, samples_for_step)
            sampler.start()
            started = time.time()
            with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as pool:
                futures = [
                    pool.submit(request_model, config, scenario, f"{config.key}-{scenario.key}-c{concurrency}-r{i + 1}")
                    for i in range(total_requests)
                ]
                for future in concurrent.futures.as_completed(futures):
                    requests_for_step.append(future.result())
            wall_s = time.time() - started
            sampler.stop()
            sampler.join(timeout=5)
            summary = summarize_step(config, scenario, concurrency, requests_for_step, samples_for_step, wall_s)
            steps.append({"summary": summary, "requests": requests_for_step, "samples": samples_for_step})
            request_rows.extend([{**row, "concurrency": concurrency} for row in requests_for_step])
            sample_rows.extend([{**row, "concurrency": concurrency} for row in samples_for_step])

            if summary["success_count"] == 0 or summary["failure_count"] > summary["success_count"]:
                break
        scenario_results.append(
            {
                "scenario": {
                    "key": scenario.key,
                    "label": scenario.label,
                    "max_tokens": scenario.max_tokens,
                    "levels": scenario.levels,
                    "min_requests": scenario.min_requests,
                    "request_multiplier": scenario.request_multiplier,
                    "num_ctx": scenario.num_ctx,
                },
                "steps": steps,
            }
        )
    cleanup_model(config)
    return {
        "label": config.label,
        "model": config.model,
        "warmup": warmup,
        "scenarios": scenario_results,
        "requests": request_rows,
        "samples": sample_rows,
    }


def run_dual_ollama_probe(config_a: ModelConfig, config_b: ModelConfig) -> dict:
    stop_loaded_ollama_models()
    scenario = SCENARIOS[0]
    samples: list[dict] = []
    sampler = Sampler("dual_ollama", "dual_ollama_cold_load", SAMPLE_INTERVAL_S, samples)
    sampler.start()
    started = time.time()
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
        future_a = pool.submit(request_model, config_a, scenario, f"{config_a.key}-dual-probe")
        future_b = pool.submit(request_model, config_b, scenario, f"{config_b.key}-dual-probe")
        result_a = future_a.result()
        result_b = future_b.result()
    wall_s = time.time() - started
    sampler.stop()
    sampler.join(timeout=5)
    requests = [result_a, result_b]
    summary = summarize_step(
        ModelConfig(
            key="dual_ollama",
            label=f"{config_a.label} + {config_b.label}",
            api_kind="ollama",
            endpoint=config_a.endpoint,
            model=f"{config_a.model} | {config_b.model}",
            family="Dual Ollama Probe",
        ),
        Scenario(
            key="dual_ollama_cold_load",
            label="Dual Ollama Cold Load",
            messages=scenario.messages,
            max_tokens=scenario.max_tokens,
            levels=[2],
            min_requests=2,
            request_multiplier=1,
        ),
        concurrency=2,
        requests=requests,
        samples=samples,
        wall_s=wall_s,
    )
    stop_loaded_ollama_models()
    return {
        "label": "Dual Ollama Cold Load Probe",
        "steps": [{"summary": summary, "requests": requests, "samples": samples}],
        "requests": requests,
        "samples": samples,
    }


def write_csv(path: Path, rows: list[dict]) -> None:
    if not rows:
        path.write_text("")
        return
    fieldnames = sorted({key for row in rows for key in row.keys()})
    with path.open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def write_summary(results: dict, dual_probe: dict) -> None:
    lines = [
        "# Spark Open-Weight Model Benchmark",
        "",
        f"- Generated at: {time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime())}",
        f"- Spark host: `{SPARK_HOST}`",
        "",
    ]
    for key, result in results.items():
        warmup = result.get("warmup") or {}
        lines.append(f"## {result['label']}")
        lines.append("")
        lines.append(f"- Model: `{result['model']}`")
        if warmup.get("ok"):
            lines.append(
                f"- Warmup latency: {fmt_num(warmup.get('latency_s'))}s | "
                f"prompt tokens {warmup.get('prompt_tokens') or 'n/a'} | "
                f"completion tokens {warmup.get('completion_tokens') or 'n/a'}"
            )
        else:
            lines.append(f"- Warmup failed: {warmup.get('error', 'unknown error')}")
        for scenario_result in result["scenarios"]:
            scenario = scenario_result["scenario"]
            lines.append(f"- Scenario: {scenario['label']}")
            for step in scenario_result["steps"]:
                s = step["summary"]
                lines.append(
                    f"  - c{s['concurrency']}: {s['success_count']}/{s['request_count']} ok, "
                    f"avg {fmt_num(s['latency_avg_s'])}s, p95 {fmt_num(s['latency_p95_s'])}s, "
                    f"throughput {fmt_num(s['throughput_rps'], 3)} rps, agg tok/s {fmt_num(s.get('aggregate_tokens_per_s'))}, "
                    f"min avail RAM {fmt_num(s.get('min_mem_available_mb'), 0)} MB, peak swap {fmt_num(s.get('peak_swap_used_mb'), 0)} MB, "
                    f"peak GPU temp {fmt_num(s.get('peak_gpu_temp_c'), 0)}C"
                )
        lines.append("")

    lines.append("## Cross-model contention probe")
    lines.append("")
    for step in dual_probe.get("steps", []):
        s = step["summary"]
        lines.append(
            f"- Dual Ollama cold load: {s['success_count']}/{s['request_count']} ok, "
            f"avg {fmt_num(s['latency_avg_s'])}s, wall {fmt_num(s['wall_time_s'])}s, "
            f"min avail RAM {fmt_num(s.get('min_mem_available_mb'), 0)} MB, "
            f"peak swap {fmt_num(s.get('peak_swap_used_mb'), 0)} MB, peak GPU temp {fmt_num(s.get('peak_gpu_temp_c'), 0)}C"
        )
    (OUTPUT_DIR / "SUMMARY.md").write_text("\n".join(lines))


def main() -> None:
    started_at = time.time()
    order = os.environ.get("MODEL_ORDER", "nemotron_super,glm45air_q6,kimi_dev72_q40")
    results = {}
    for key in [item.strip() for item in order.split(",") if item.strip()]:
        config = MODELS[key]
        print(json.dumps({"starting_model": key, "label": config.label}, indent=2), flush=True)
        results[key] = run_model(config)

    dual_probe = (
        run_dual_ollama_probe(MODELS["glm45air_q6"], MODELS["kimi_dev72_q40"])
        if RUN_DUAL_PROBE
        else {"label": "Dual Ollama Cold Load Probe", "steps": [], "requests": [], "samples": []}
    )
    payload = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "duration_s": round(time.time() - started_at, 2),
        "spark_host": SPARK_HOST,
        "scenarios": [
            {
                "key": scenario.key,
                "label": scenario.label,
                "max_tokens": scenario.max_tokens,
                "levels": scenario.levels,
                "min_requests": scenario.min_requests,
                "request_multiplier": scenario.request_multiplier,
            }
            for scenario in SCENARIOS
        ],
        "models": results,
        "dual_ollama_probe": dual_probe,
    }
    (OUTPUT_DIR / "results.json").write_text(json.dumps(payload, indent=2))
    write_summary(results, dual_probe)

    request_rows: list[dict] = []
    sample_rows: list[dict] = []
    step_rows: list[dict] = []
    for key, result in results.items():
        request_rows.extend(result["requests"])
        sample_rows.extend(result["samples"])
        for scenario_result in result["scenarios"]:
            for step in scenario_result["steps"]:
                step_rows.append(step["summary"])
    request_rows.extend([{**row, "concurrency": 2} for row in dual_probe["requests"]])
    sample_rows.extend([{**row, "concurrency": 2} for row in dual_probe["samples"]])
    for step in dual_probe["steps"]:
        step_rows.append(step["summary"])

    write_csv(OUTPUT_DIR / "requests.csv", request_rows)
    write_csv(OUTPUT_DIR / "samples.csv", sample_rows)
    write_csv(OUTPUT_DIR / "step_summary.csv", step_rows)
    print(json.dumps({"output_dir": str(OUTPUT_DIR), "duration_s": payload["duration_s"]}, indent=2))


if __name__ == "__main__":
    main()
