#!/usr/bin/env python
"""
ollama_bench.py – benchmark locale di modelli Ollama (latency & throughput).

Esempio:
    python ollama_bench.py -m llama3 -p "Ciao, spiegami la relatività" -n 5 --gpu
"""
from __future__ import annotations

from bench_utils import emit, sysinfo

import argparse, time, statistics, sys
from pathlib import Path
import requests
import subprocess
from tabulate import tabulate
import datetime as dt

HOST = "http://localhost:11434"

def load_prompt(path_or_text: str) -> str:
    p = Path(path_or_text)
    return p.read_text() if p.exists() else path_or_text

def call_ollama(model: str, prompt: str, use_gpu: bool, seed: int):
    """Restituisce: wall_time, total_ns, eval_ns, eval_tokens"""
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": False,
        "options": {
            "seed": seed,
            "temperature": 0.0,
            "num_gpu": -1 if use_gpu else 0,
        },
    }

    start = time.perf_counter()
    r = requests.post(f"{HOST}/api/generate", json=payload, timeout=1800)
    wall = time.perf_counter() - start
    r.raise_for_status()
    data = r.json()

    tot_ns  = data.get("total_duration", 0)
    eval_ns = data.get("eval_duration", 0)
    tokens  = data.get("eval_count", 0)
    txt     = data.get("response", "")
    return wall, tot_ns, eval_ns, tokens, len(txt)

def get_ollama_version(host: str = "http://localhost:11434") -> str:
    """
    Prova a leggere la versione del server Ollama via API; se fallisce,
    ripiega su `ollama --version`. Ritorna 'unknown' se non disponibile.
    """
    try:
        r = requests.get(f"{host}/api/version", timeout=5)
        if r.ok:
            j = r.json()
            v = j.get("version") or j.get("data") or j.get("ollama")  # compat
            if isinstance(v, str) and v.strip():
                return v.strip()
    except Exception:
        pass

    try:
        cp = subprocess.run(
            ["ollama", "--version"],
            capture_output=True, text=True, timeout=5
        )
        if cp.returncode == 0:
            return cp.stdout.strip() or cp.stderr.strip() or "unknown"
    except Exception:
        pass

    return "unknown"

def ns2s(ns: int) -> float:
    return ns / 1e9

# ────────────────────────────── CLI ───────────────────────────────────────────
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-m", "--model", default="llama3", help="nome modello Ollama")
    ap.add_argument("-p", "--prompt", required=True,
                    help="testo o file da usare come prompt")
    ap.add_argument("-n", "--repeats", type=int, default=3)
    ap.add_argument("--gpu", action="store_true",
                    help="lascia usare la GPU (default: CPU only)")
    ap.add_argument("--seed", type=int, default=42, help="seed riproducibile")
    ap.add_argument("--json", action="store_true",
                help="stampa il risultato in JSON al termine")
    ap.add_argument("--append-csv",
                help="percorso di un CSV su cui accodare la riga di output")    
    ap.add_argument("--run-id", required=False,
                help="identificativo unico della sessione benchmark")
    args = ap.parse_args()

    run_id = args.run_id or "dev"

    prompt = load_prompt(args.prompt)
    mode   = "GPU" if args.gpu else "CPU"

    w_times, e_times, t_times, sizes = [], [], [], []
    print(f"• Benchmarking {args.model} – {mode} – seed {args.seed}")
    for i in range(1, args.repeats + 1):
        print(f"  Run {i}/{args.repeats} … ", end="", flush=True)
        try:
            wall, tot_ns, eval_ns, tok, out_len = call_ollama(
                args.model, prompt, args.gpu, args.seed
            )
        except requests.exceptions.HTTPError as e:
            status = getattr(e.response, "status_code", None)
            body = ""
            try: body = e.response.text
            except Exception: pass
            if status in (400, 404) and any(k in (body or "").lower() for k in ["model", "not found", "no such model"]):
                print(f"⚠️  Model '{args.model}' not installed. Skipping this run.")
            else:
                print(f"⚠️  HTTP error from Ollama ({status}): {e}. Skipping this run.")
            continue
        except requests.exceptions.ConnectionError:
            print("⚠️  Ollama not reachable at http://localhost:11434. Skipping this run.")
            continue
        except Exception as e:
            print(f"⚠️  Ollama run failed (GPU={args.gpu}): {e}. Skipping this run.")
            continue

        w_times.append(wall)
        e_times.append(ns2s(eval_ns))
        t_times.append(tok / ns2s(eval_ns) if eval_ns else 0)
        sizes.append(out_len)
        print(f"{wall:,.2f}s  (tokens/s {t_times[-1]:.1f})")

    # riepilogo
    def min_med_max(arr):
        return min(arr), statistics.median(arr), max(arr)

    if not w_times:
        print("⚠️  No successful Ollama runs; not writing CSV for this configuration.")
        return 0
    wall_min, wall_med, wall_max = min_med_max(w_times)
    tok_min,  tok_med,  tok_max  = min_med_max(t_times)

    rows = [
        ["Wall latency [s]", f"{wall_min:,.2f}", f"{wall_med:,.2f}", f"{wall_max:,.2f}"],
        ["Tokens/s",          f"{tok_min:,.1f}",  f"{tok_med:,.1f}",  f"{tok_max:,.1f}"],
    ]
    print("\n" + tabulate(rows,
        headers=["Metric", "min", "average", "max"],
        tablefmt="github"))
    
    result = {"run_id": run_id}
    result.update(sysinfo())
    result.update({           # costruisci il dict con tutte le metriche
        "bench":"ollama",
        "version": get_ollama_version(),
        "model": args.model,
        "gpu_used": args.gpu,
        "wall_min_s":   f"{wall_min:.2f}",
        "wall_med_s":   f"{wall_med:.2f}",
        "wall_max_s":   f"{wall_max:.2f}",
        "tok_min_s":    f"{tok_min:.2f}",
        "tok_med_s":    f"{tok_med:.2f}",
        "tok_max_s":    f"{tok_max:.2f}",
        "seed": args.seed,
        "timestamp": dt.datetime.now().isoformat(timespec="seconds"),
        "reference": False,
    })
    result.update(sysinfo())
    emit(result, args)

if __name__ == "__main__":
    try:
        main()
    except requests.exceptions.ConnectionError:
        print("⚠️  Ollama apparently not running on http://localhost:11434. Skipping.")
    except Exception as e:
        print(f"⚠️  Unexpected error: {e}.")