"""
benchmark.py — Speed-test LLM models and select the fastest available.

Benchmarks CANDIDATE_MODELS at startup, picks top 5 as ACTIVE_MODELS.
Lumo is included in the active set only when its benchmark succeeds.
"""

import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

from json_storage import update_json
from providers import CANDIDATE_MODELS, MODELS, stream_benchmark

_BENCH_CACHE_FILE = os.path.join(os.path.dirname(__file__), "data", "bench_cache.json")


def load_bench_cache() -> dict:
    """Load persisted benchmark results."""
    try:
        with open(_BENCH_CACHE_FILE) as f:
            return json.load(f)
    except (OSError, ValueError):  # missing or corrupt cache: start empty
        return {}


def save_bench_cache(results: dict) -> None:
    """Save benchmark results, merging with existing cache."""
    ts = int(time.time())

    def merge_results(cache: dict) -> None:
        for model, result in results.items():
            if result.get("latency") is not None:
                cache[model] = {**result, "ts": ts}

    update_json(_BENCH_CACHE_FILE, dict, merge_results)


BENCH_SYSTEM = "Jsi vtipný asistent. Odpovídej česky ve 2 odstavcích."
BENCH_USER = "Co je to Brno a proč je slavné?"


def benchmark_models(candidates: list[str] | None = None, top_n: int = 5,
                     on_progress=None):
    """Benchmark candidate models in parallel, return speed-ranked results.

    Args:
        candidates: Models to test (defaults to CANDIDATE_MODELS).
        top_n: How many to select as active (includes Lumo when successful).
        on_progress: Optional callback(model, status, result) for live updates.
            status: "testing", "ok", "fail"
            result: dict with latency/ttft/tok_sec/error (None for "testing")

    Returns (active_models, ranked_chain, results_dict) where:
      - active_models: up to top_n successful model names
      - ranked_chain: list of (model_name, latency) sorted by TTFT
      - results_dict: model_name -> {latency, ttft, tok_sec, error}
    """
    if candidates is None:
        candidates = list(CANDIDATE_MODELS)

    # Filter to models that exist in registry
    candidates = [n for n in candidates if n in MODELS]
    if not candidates:
        return [], [], {}

    results: dict[str, dict] = {}
    _lock = threading.Lock()

    def bench_one(model_name: str) -> None:
        if on_progress:
            on_progress(model_name, "testing", None)
        total, ttft, tok_sec, err = stream_benchmark(
            model_name, BENCH_SYSTEM, BENCH_USER, max_tokens=80)
        r = {"latency": total, "ttft": ttft, "tok_sec": tok_sec, "error": err}
        with _lock:
            results[model_name] = r
        if on_progress:
            on_progress(model_name, "ok" if not err else "fail", r)

    # Full parallelism: one thread per model, hard 12s total timeout
    # Don't use `with` — it waits for all threads to finish on exit
    _TOTAL_TIMEOUT = 12
    executor = ThreadPoolExecutor(max_workers=len(candidates))
    futures = {executor.submit(bench_one, name): name for name in candidates}
    try:
        for f in as_completed(futures, timeout=_TOTAL_TIMEOUT):
            try:
                f.result()
            except Exception:  # noqa: BLE001, S110 - one crashed model must not stop the rest; it is marked below
                pass
    except TimeoutError:
        pass
    # Mark any still-running models as timed out
    for f, name in futures.items():
        if name not in results:
            r = {"latency": None, "ttft": None, "tok_sec": None, "error": "Timeout (total)"}
            results[name] = r
            if on_progress:
                on_progress(name, "fail", r)
    # Let straggler threads die in background — don't block
    executor.shutdown(wait=False)

    save_bench_cache(results)

    # Rank by TTFT (responsiveness). NVIDIA sorts after every other working
    # provider regardless of speed — its free tier's rate limits are too low
    # to rely on as a primary pick, so it's kept as a last-resort fallback.
    successful = [
        (name, r["latency"])
        for name, r in results.items()
        if r.get("latency") is not None and not r.get("error")
    ]
    successful.sort(key=lambda x: (
        MODELS[x[0]]["provider"] == "nvidia",
        results[x[0]].get("ttft") or 999,
    ))

    # Keep failed candidates at the end for later fallback retries.
    ranked_ids = {name for name, _ in successful}
    ranked_chain = list(successful)
    for name in candidates:
        if name not in ranked_ids:
            ranked_chain.append((name, None))

    # Pick top_n working models, preserving Lumo only when it also worked.
    active = []
    for name, _ in successful:
        if len(active) >= top_n:
            break
        active.append(name)
    if top_n > 0 and "lumo" in ranked_ids and "lumo" not in active:
        if len(active) >= top_n:
            active[-1] = "lumo"
        else:
            active.append("lumo")

    return active, ranked_chain, results