#!/usr/bin/env python3
"""Benchmark Ollama en série, sans réutilisation du cache de prompt.

Python 3 standard, nvidia-smi et une API Ollama déjà disponible. Ne modifie ni
les modèles, ni les pilotes, ni les paramètres GPU. Arrêter les autres clients
Ollama avant le lancement. Le modèle cible est déchargé entre les essais.
Exécuter sur l'hôte du GPU : --url doit désigner l'Ollama de cet hôte, puisque
la télémétrie nvidia-smi est locale. Un label « unlocked » ne prouve pas que le
déverrouillage a réussi : vérifier séparément le journal noyau de ce démarrage.
"""

import argparse
import csv
import hashlib
import json
import os
from pathlib import Path
import platform
import statistics
import subprocess
import sys
import threading
import time
from datetime import datetime, timezone
from urllib.error import HTTPError
from urllib.request import Request, urlopen


GPU_FIELDS = "name,driver_version,memory.total,memory.used,utilization.gpu,power.draw,power.limit,temperature.gpu,clocks.sm,clocks.mem"
PROFILE_LINES = {"short": 2, "medium": 18, "long": 72}


def now():
    return datetime.now(timezone.utc).isoformat()


def api(url, path, payload=None, timeout=600):
    data = None if payload is None else json.dumps(payload).encode()
    request = Request(url + path, data=data, headers={"Content-Type": "application/json"})
    try:
        with urlopen(request, timeout=timeout) as response:
            return json.load(response)
    except HTTPError as error:
        raise RuntimeError(f"Ollama HTTP {error.code}: {error.read().decode()[:500]}") from error


def command(args):
    return subprocess.check_output(args, text=True, timeout=8).strip()


def gpu_snapshot():
    values = next(csv.reader([command(["nvidia-smi", "--id=0", "--query-gpu=" + GPU_FIELDS, "--format=csv,noheader,nounits"])]))
    result = dict(zip(GPU_FIELDS.split(","), (value.strip() for value in values)))
    for key in GPU_FIELDS.split(",")[2:]:
        try:
            result[key] = float(result[key])
        except ValueError:
            result[key] = None
    result["at"] = now()
    processes = command(["nvidia-smi", "--id=0", "--query-compute-apps=process_name", "--format=csv,noheader"])
    result["compute_processes"] = processes.splitlines() if processes else []
    return result


def assert_idle():
    deadline = time.monotonic() + 15
    calm_samples = 0
    while True:
        sample = gpu_snapshot()
        foreign = [p for p in sample["compute_processes"] if "ollama" not in p and "llama-server" not in p]
        if foreign:
            raise RuntimeError("Charge GPU tierce détectée : arrêter les autres clients/charges avant le benchmark.")
        calm_samples = calm_samples + 1 if (sample["utilization.gpu"] or 0) <= 10 else 0
        if calm_samples >= 3:
            return
        if time.monotonic() >= deadline:
            raise RuntimeError("GPU occupé : arrêter les autres clients/charges avant le benchmark.")
        time.sleep(1)


def prompt_for(profile, repeat):
    # Le corpus est entièrement synthétique. Aucun journal de la machine n'est lu.
    lines = []
    for number in range(PROFILE_LINES[profile]):
        lines.append(
            f"Service api-{number:04d} : Debian, conteneur rootless, trois réplicas, "
            f"port interne 8080, base PostgreSQL, sauvegarde nocturne vérifiée, "
            f"latence p95 {40 + number % 29} ms, erreurs HTTP {number % 3} pour mille, "
            "métriques Prometheus. La sonde readiness vérifie la connexion à la base ; "
            "la sonde liveness vérifie seulement le processus. Le dernier déploiement "
            "a augmenté le nombre de connexions sans modifier la taille du pool."
        )
    return (
        f"Dossier synthétique {profile}, répétition {repeat:02d}.\n"
        "Analyse les services ci-dessous. Rédige en français une analyse technique "
        "détaillée d'au moins 600 mots : diagnostic, priorités, vérifications en lecture "
        "seule et corrections argumentées. Ne te contente pas d'un résumé.\n\n"
        + "\n".join(lines)
        + "\n\nDonne maintenant ton analyse opérationnelle détaillée."
    )


def save(path, data):
    temporary = path.with_name(path.name + ".tmp")
    temporary.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n")
    os.replace(temporary, path)


def metadata(args):
    tags = api(args.url, "/api/tags")["models"]
    selected = next((m for m in tags if m["name"] == args.model), None)
    if selected is None:
        raise RuntimeError(f"Modèle installé introuvable : {args.model}. Aucun téléchargement effectué.")
    show = api(args.url, "/api/show", {"model": args.model})
    module = Path("/sys/module/gpu_unlock")
    unlock = {"module_present": module.exists()}
    for name in ("targets", "resume"):
        path = module / "parameters" / name
        if path.exists():
            unlock[name] = path.read_text().strip()
    unlock["boot_options"] = [s for s in Path("/proc/cmdline").read_text().split() if s.startswith("gpu_unlock.")]
    return {
        "at": now(), "ollama": api(args.url, "/api/version"), "kernel": platform.release(),
        "script_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "model": selected, "show": {k: show.get(k) for k in ("details", "parameters", "capabilities", "model_info")},
        "template_sha256": hashlib.sha256(show.get("template", "").encode()).hexdigest(),
        "gpu": gpu_snapshot(), "gpu_unlock": unlock,
        "options": {"num_ctx": args.context, "num_predict": args.predict, "num_batch": 512, "temperature": 0, "seed": 42},
        "thinking": False, "profile_lines": {p: PROFILE_LINES[p] for p in args.profiles},
        "cache_policy": "Unload target, preload empty prompt, then one measured request per trial.",
    }


def measure(args, profile, repeat, options):
    result = {"profile": profile, "repeat": repeat, "at": now(), "status": "running"}
    try:
        _measure(args, profile, repeat, options, result)
        result["status"] = "complete"
        return result
    except (Exception, KeyboardInterrupt) as error:
        result["status"] = "interrupted" if isinstance(error, KeyboardInterrupt) else "failed"
        result["error"] = str(error) or "Interruption clavier"
        error.bench_run = result
        raise


def _measure(args, profile, repeat, options, result):
    assert_idle()
    # Pas de cache de prompt inter-essais, même sur les versions sans compteur de cache.
    api(args.url, "/api/generate", {"model": args.model, "keep_alive": 0, "stream": False})
    preload = api(args.url, "/api/generate", {"model": args.model, "keep_alive": "5m", "stream": False,
                                                   "options": options})
    prompt = prompt_for(profile, repeat)
    result.update({"prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
                   "prompt_bytes": len(prompt.encode()), "preload": preload, "gpu_samples": [], "ps_samples": []})
    stop = threading.Event()
    collection_errors = []

    def collect():
        while not stop.is_set():
            try:
                result["gpu_samples"].append(gpu_snapshot())
                loaded = api(args.url, "/api/ps", timeout=8).get("models", [])
                result["ps_samples"].append({"at": now(), "models": loaded})
            except Exception as error:
                collection_errors.append(str(error))
            stop.wait(1)

    thread = threading.Thread(target=collect, daemon=True)
    thread.start()
    payload = {"model": args.model, "prompt": prompt, "stream": True, "think": False,
               "keep_alive": "5m", "options": options}
    request = Request(args.url + "/api/generate", data=json.dumps(payload).encode(),
                      headers={"Content-Type": "application/json"})
    started = time.perf_counter()
    first_token = None
    final = None
    output = hashlib.sha256()
    try:
        with urlopen(request, timeout=600) as response:
            for line in response:
                if not line.strip():
                    continue
                chunk = json.loads(line)
                if chunk.get("error"):
                    raise RuntimeError(chunk["error"])
                text = chunk.get("thinking", "") + chunk.get("response", "")
                if text and first_token is None:
                    first_token = time.perf_counter() - started
                output.update(text.encode())
                if chunk.get("done"):
                    final = chunk
    finally:
        result["wall_seconds"] = time.perf_counter() - started
        stop.set()
        thread.join(timeout=18)
    if final is None:
        raise RuntimeError("Flux incomplet : aucune métrique finale Ollama.")
    result.update({"metrics": {k: v for k, v in final.items() if k not in ("response", "thinking", "context")},
                   "ttft_seconds": first_token, "output_sha256": output.hexdigest(), "collection_errors": collection_errors})
    count, duration = final.get("prompt_eval_count", 0), final.get("prompt_eval_duration", 0)
    cached = final.get("prompt_eval_cached_count", 0)
    if not count or not duration or not final.get("eval_count") or not final.get("eval_duration"):
        raise RuntimeError("Compteurs/durées manquants ou nuls : essai non exploitable.")
    if cached:
        raise RuntimeError(f"Cache de prompt réutilisé ({cached} tokens) malgré la purge.")
    if count + args.predict >= args.context:
        raise RuntimeError("Prompt trop long pour le contexte et la génération demandés.")
    target_ps = [m for sample in result["ps_samples"] for m in sample["models"] if m.get("name") == args.model]
    if any(m.get("name") != args.model for sample in result["ps_samples"] for m in sample["models"]):
        raise RuntimeError("Un autre modèle est chargé pendant l'essai : concurrence possible.")
    if not target_ps or any(m.get("size_vram", 0) < m.get("size", 0) for m in target_ps):
        raise RuntimeError("Résidence GPU complète non vérifiée via /api/ps.")
    if getattr(args, "expected_digest", None) and any(m.get("digest") != args.expected_digest for m in target_ps):
        raise RuntimeError("L'empreinte du modèle chargé diffère de l'inventaire initial.")
    if any(m.get("context_length", args.context) != args.context for m in target_ps):
        raise RuntimeError("Le contexte chargé diffère du contexte demandé.")
    if collection_errors or not result["gpu_samples"]:
        raise RuntimeError("Télémétrie incomplète : " + "; ".join(collection_errors))
    for sample in result["gpu_samples"]:
        if any("ollama" not in p and "llama-server" not in p for p in sample["compute_processes"]):
            raise RuntimeError("Une charge GPU tierce a été détectée pendant l'essai.")
        expected = getattr(args, "expected_power_limit", None)
        if expected is not None and sample.get("power.limit") != expected:
            raise RuntimeError("La limite de puissance GPU a changé depuis l'inventaire.")
    result["prompt_tokens_per_second"] = (count - cached) * 1e9 / duration
    result["generation_tokens_per_second"] = final["eval_count"] * 1e9 / final["eval_duration"]


def summarize(runs):
    summary = {}
    for profile in dict.fromkeys(r["profile"] for r in runs):
        selected = [r for r in runs if r["profile"] == profile]
        entry = {"runs": len(selected), "prompt_tokens": [r["metrics"]["prompt_eval_count"] for r in selected],
                 "generated_tokens": [r["metrics"]["eval_count"] for r in selected]}
        for name in ("prompt_tokens_per_second", "generation_tokens_per_second", "ttft_seconds", "wall_seconds"):
            values = [r[name] for r in selected if r[name] is not None]
            entry[name] = {"median": statistics.median(values), "min": min(values), "max": max(values)} if values else None
        summary[profile] = entry
    return summary


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="http://127.0.0.1:11434")
    parser.add_argument("--model", default="qwen3.8:27b")
    parser.add_argument("--label", required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--profiles", default="short,medium,long")
    parser.add_argument("--profile-lines", default="short=2,medium=18,long=72", help="Lignes synthétiques par profil, pour calibrer les tailles.")
    parser.add_argument("--context", type=int, default=16384)
    parser.add_argument("--predict", type=int, default=256)
    parser.add_argument("--check", action="store_true", help="Inventaire uniquement, sans génération ni déchargement.")
    args = parser.parse_args()
    args.url = args.url.rstrip("/")
    args.profiles = args.profiles.split(",")
    try:
        overrides = dict(item.split("=") for item in args.profile_lines.split(","))
        if set(overrides) != set(PROFILE_LINES) or any(int(v) < 1 for v in overrides.values()):
            raise ValueError()
        PROFILE_LINES.update({k: int(v) for k, v in overrides.items()})
    except ValueError:
        parser.error("--profile-lines exige short=N,medium=N,long=N, tous strictement positifs.")
    if args.repeats < 1 or args.predict < 1 or args.context <= args.predict or any(p not in PROFILE_LINES for p in args.profiles):
        parser.error("Paramètres de répétition, de contexte, de génération ou de profils invalides.")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    try:
        with args.output.open("x"):
            pass
    except FileExistsError:
        parser.error(f"Le fichier existe déjà : {args.output}")
    document = {"schema": 1, "label": args.label, "started_at": now(), "status": "running", "runs": []}
    save(args.output, document)
    try:
        document["metadata"] = metadata(args)
        args.expected_digest = document["metadata"]["model"]["digest"]
        args.expected_power_limit = document["metadata"]["gpu"]["power.limit"]
        save(args.output, document)
        if not args.check:
            for profile in args.profiles:
                for repeat in range(1, args.repeats + 1):
                    run = measure(args, profile, repeat, document["metadata"]["options"])
                    document["runs"].append(run)
                    document["summary"] = summarize(document["runs"])
                    save(args.output, document)
                    print(f"{profile} {repeat}/{args.repeats}: {run['metrics']['prompt_eval_count']} tokens, "
                          f"prompt {run['prompt_tokens_per_second']:.2f} tok/s, "
                          f"génération {run['generation_tokens_per_second']:.2f} tok/s", flush=True)
        document["status"] = "inventory_only" if args.check else "complete"
        return_code = 0
    except (Exception, KeyboardInterrupt) as error:
        document["status"] = "interrupted" if isinstance(error, KeyboardInterrupt) else "failed"
        document["error"] = str(error) or "Interruption clavier"
        if hasattr(error, "bench_run"):
            document["runs"].append(error.bench_run)
        print(document["error"], file=sys.stderr, flush=True)
        return_code = 1
    document["finished_at"] = now()
    save(args.output, document)
    return return_code


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