"""Controlled safetensors-versus-GGUF recipe benchmark.

This is a *model recipe* benchmark, unlike
:mod:`meshnet_node.route_session_benchmark`, which is a transport harness.  It
answers one question: on one machine, with one model revision and one fixed
workload, what do the Transformers/safetensors recipe and the whole-model
llama.cpp/GGUF recipes actually cost in speed, memory, fit, and output drift?

Two ideas keep the answer honest.

**Lanes.**  A recipe belongs to exactly one :class:`Lane`.  The quality lane
holds near-lossless recipes (bf16/f16 weights) whose outputs may legitimately be
compared for numerical agreement.  The performance-fit lane holds quantized
recipes (Q8_0, Q4_K_M, ...).  Quantized recipes are judged on speed, memory, and
artifact size only; their drift is *reported* but never read as evidence that
Q4 and bf16 are numerically equivalent, because they are not.  The lane is a
property of the recipe, so nothing downstream can quietly cross the boundary.

**Drivers.**  The measurement core here is pure and runtime-free: it drives a
:class:`RecipeDriver` and computes metrics.  Real runtimes live in
:mod:`meshnet_node.recipe_drivers` and are imported only on demand, which keeps
the default test suite deterministic, GPU-free and model-download-free while the
same code path produces the real evidence.
"""

from __future__ import annotations

import argparse
import json
import statistics
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict, dataclass, field
from difflib import SequenceMatcher
from enum import Enum
from pathlib import Path
from typing import Any, Mapping, Protocol, Sequence

# Layout of the report document produced by :func:`build_report`.
REPORT_SCHEMA_VERSION = 1


class Lane(str, Enum):
    """Why a recipe is being measured at all.

    ``QUALITY`` recipes carry near-lossless weights, so comparing their output
    with the reference recipe is a meaningful correctness signal.
    ``PERFORMANCE_FIT`` recipes carry quantized weights: they exist to be faster
    or to fit, and their drift is descriptive, never a pass/fail equivalence
    claim.
    """

    QUALITY = "quality"
    PERFORMANCE_FIT = "performance-fit"


class BenchmarkError(RuntimeError):
    """Raised when a benchmark cannot be run as specified."""


@dataclass(frozen=True)
class SamplingPolicy:
    """The sampling policy every recipe must be given, identically.

    Greedy by default: sampling noise would otherwise be indistinguishable from
    quantization drift, and the whole point of the quality lane is to tell those
    two apart.
    """

    temperature: float = 0.0
    top_p: float = 1.0
    top_k: int = 1
    seed: int = 1234
    max_output_tokens: int = 64

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass(frozen=True)
class PromptSpec:
    """One fixed prompt, tagged with the context length it is meant to exercise."""

    id: str
    text: str
    context_class: str = "short"

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass(frozen=True)
class BenchmarkPlan:
    """The controlled variables: identical for every recipe in a report.

    A plan is the experiment.  If two recipes were measured under different
    plans, their numbers are not comparable and the report must not pretend they
    are, so the plan is recorded once at the top of the document rather than
    per-recipe.
    """

    plan_id: str
    model_id: str
    model_revision: str
    prompts: tuple[PromptSpec, ...]
    sampling: SamplingPolicy = SamplingPolicy()
    concurrency_levels: tuple[int, ...] = (1, 4)
    repeats: int = 1
    warmup_requests: int = 1

    def __post_init__(self) -> None:
        if not self.prompts:
            raise BenchmarkError("a benchmark plan needs at least one prompt")
        if not self.concurrency_levels or any(level < 1 for level in self.concurrency_levels):
            raise BenchmarkError("concurrency levels must all be >= 1")
        if self.repeats < 1:
            raise BenchmarkError("repeats must be >= 1")
        if 1 not in self.concurrency_levels or 4 not in self.concurrency_levels:
            raise BenchmarkError("a controlled baseline must include concurrency levels 1 and 4")

    def to_dict(self) -> dict:
        return {
            "plan_id": self.plan_id,
            "model_id": self.model_id,
            "model_revision": self.model_revision,
            "prompts": [prompt.to_dict() for prompt in self.prompts],
            "sampling": self.sampling.to_dict(),
            "concurrency_levels": list(self.concurrency_levels),
            "repeats": self.repeats,
            "warmup_requests": self.warmup_requests,
        }


@dataclass(frozen=True)
class RecipeSpec:
    """One runtime recipe under test.

    ``is_reference`` marks the single recipe every other recipe's output drift is
    measured against — the current Transformers/safetensors route, which
    decision Gate 8 keeps as the correctness backend.
    """

    id: str
    runtime: str
    weight_format: str
    weight_quantization: str
    lane: Lane
    device: str
    artifact_path: str = ""
    source_model_id: str = ""
    source_model_revision: str = ""
    artifact_sha256: str = ""
    is_reference: bool = False
    notes: str = ""

    def to_dict(self) -> dict:
        data = asdict(self)
        data["lane"] = self.lane.value
        return data


@dataclass(frozen=True)
class LoadStats:
    """What loading the recipe cost, before any token is generated."""

    artifact_bytes: int
    load_ms: float
    rss_bytes: int = 0
    vram_bytes: int = 0
    backend_detail: str = ""

    def to_dict(self) -> dict:
        return asdict(self)


@dataclass(frozen=True)
class GenerationSample:
    """One completed generation as reported by a driver.

    ``prefill_ms``/``decode_ms`` are the runtime's own split where it exposes one
    (llama.cpp does); drivers that cannot split honestly report ``prefill_ms`` as
    the time to the first token.  ``queue_wait_ms`` separates time spent waiting
    for a runtime slot from time spent computing, so a concurrency-4 TTFT is not
    silently read as a slower prefill.
    """

    text: str
    prompt_tokens: int
    decode_tokens: int
    ttft_ms: float
    prefill_ms: float
    decode_ms: float
    total_ms: float
    queue_wait_ms: float = 0.0


class RecipeDriver(Protocol):
    """The seam every runtime implements; the measurement core knows nothing else."""

    def load(self) -> LoadStats:
        """Load the artifact and return its cost."""

    def generate(self, prompt: str, sampling: SamplingPolicy) -> GenerationSample:
        """Run one complete generation under the given sampling policy."""

    def memory_probe(self) -> tuple[int, int]:
        """Return ``(rss_bytes, vram_bytes)`` observed right now."""

    def close(self) -> None:
        """Release the runtime."""


@dataclass(frozen=True)
class RequestOutcome:
    """One request attempt, successful or not.

    A failure is a first-class result, not an exception that aborts the run: a
    recipe that cannot sustain concurrency 4 has told us something, and the
    report must carry it rather than lose it.
    """

    recipe_id: str
    concurrency: int
    prompt_id: str
    repeat: int
    ok: bool
    latency_ms: float
    ttft_ms: float = 0.0
    prefill_ms: float = 0.0
    decode_ms: float = 0.0
    queue_wait_ms: float = 0.0
    prompt_tokens: int = 0
    decode_tokens: int = 0
    text: str = ""
    error: str = ""

    def to_dict(self) -> dict:
        return asdict(self)


def _percentile(values: Sequence[float], percentile: float) -> float:
    """Nearest-rank percentile; 0.0 for an empty sample."""
    ordered = sorted(values)
    if not ordered:
        return 0.0
    rank = max(1, -(-len(ordered) * percentile // 100))
    return round(ordered[int(rank) - 1], 4)


def _mean(values: Sequence[float]) -> float:
    return round(statistics.fmean(values), 4) if values else 0.0


@dataclass(frozen=True)
class ConcurrencyMetrics:
    """Aggregate metrics for one recipe at one concurrency level."""

    concurrency: int
    requests: int
    failures: int
    wall_ms: float
    ttft_p50_ms: float
    ttft_p95_ms: float
    latency_p50_ms: float
    latency_p95_ms: float
    prefill_tokens_per_sec: float
    decode_tokens_per_sec: float
    aggregate_decode_tokens_per_sec: float
    peak_rss_bytes: int
    peak_vram_bytes: int
    failure_reasons: tuple[str, ...] = ()

    def to_dict(self) -> dict:
        data = asdict(self)
        data["failure_reasons"] = list(self.failure_reasons)
        return data


def summarize_concurrency(
    outcomes: Sequence[RequestOutcome],
    *,
    concurrency: int,
    wall_ms: float,
    peak_rss_bytes: int,
    peak_vram_bytes: int,
) -> ConcurrencyMetrics:
    """Aggregate one recipe/concurrency cell.

    Per-request rates are averaged over successful requests; aggregate
    throughput is total decoded tokens over the wall clock of the whole cell,
    which is the only figure that credits a runtime for overlapping work.
    """
    ok = [outcome for outcome in outcomes if outcome.ok]
    failures = [outcome for outcome in outcomes if not outcome.ok]
    decode_tokens = sum(outcome.decode_tokens for outcome in ok)

    prefill_rates = [
        outcome.prompt_tokens / (outcome.prefill_ms / 1000)
        for outcome in ok
        if outcome.prefill_ms > 0 and outcome.prompt_tokens
    ]
    decode_rates = [
        outcome.decode_tokens / (outcome.decode_ms / 1000)
        for outcome in ok
        if outcome.decode_ms > 0 and outcome.decode_tokens
    ]
    return ConcurrencyMetrics(
        concurrency=concurrency,
        requests=len(outcomes),
        failures=len(failures),
        wall_ms=round(wall_ms, 4),
        ttft_p50_ms=_percentile([outcome.ttft_ms for outcome in ok], 50),
        ttft_p95_ms=_percentile([outcome.ttft_ms for outcome in ok], 95),
        latency_p50_ms=_percentile([outcome.latency_ms for outcome in ok], 50),
        latency_p95_ms=_percentile([outcome.latency_ms for outcome in ok], 95),
        prefill_tokens_per_sec=_mean(prefill_rates),
        decode_tokens_per_sec=_mean(decode_rates),
        aggregate_decode_tokens_per_sec=round(decode_tokens / max(1e-6, wall_ms / 1000), 4),
        peak_rss_bytes=peak_rss_bytes,
        peak_vram_bytes=peak_vram_bytes,
        failure_reasons=tuple(sorted({outcome.error for outcome in failures if outcome.error})),
    )


@dataclass
class RecipeMeasurement:
    """Everything measured for one recipe across every concurrency level."""

    recipe: RecipeSpec
    load: LoadStats
    metrics: dict[int, ConcurrencyMetrics] = field(default_factory=dict)
    outcomes: list[RequestOutcome] = field(default_factory=list)
    unavailable_reason: str = ""

    @property
    def available(self) -> bool:
        return not self.unavailable_reason

    def outputs_by_prompt(self) -> dict[str, str]:
        """First successful output per prompt, at the lowest concurrency measured.

        Drift is a property of the recipe, not of load: concurrency must not
        change greedy output, so the least-contended sample is the fair one.
        """
        best: dict[str, tuple[int, str]] = {}
        for outcome in self.outcomes:
            if not outcome.ok:
                continue
            seen = best.get(outcome.prompt_id)
            if seen is None or outcome.concurrency < seen[0]:
                best[outcome.prompt_id] = (outcome.concurrency, outcome.text)
        return {prompt_id: text for prompt_id, (_, text) in best.items()}

    def to_dict(self) -> dict:
        return {
            "recipe": self.recipe.to_dict(),
            "available": self.available,
            "unavailable_reason": self.unavailable_reason,
            "load": self.load.to_dict(),
            "concurrency": {
                str(level): metrics.to_dict() for level, metrics in sorted(self.metrics.items())
            },
            "outcomes": [outcome.to_dict() for outcome in self.outcomes],
        }


@dataclass(frozen=True)
class DriftReport:
    """Output drift of one recipe against the reference recipe.

    ``advisory`` is true for every performance-fit recipe: the number is
    published, but a Q4 recipe disagreeing with bf16 is expected behaviour, not a
    defect, and no gate may read it as one.
    """

    recipe_id: str
    lane: Lane
    reference_id: str
    compared_prompts: int
    exact_match_rate: float
    mean_similarity: float
    advisory: bool
    per_prompt: tuple[dict[str, Any], ...] = ()

    def to_dict(self) -> dict:
        return {
            "recipe_id": self.recipe_id,
            "lane": self.lane.value,
            "reference_id": self.reference_id,
            "compared_prompts": self.compared_prompts,
            "exact_match_rate": self.exact_match_rate,
            "mean_similarity": self.mean_similarity,
            "advisory": self.advisory,
            "per_prompt": list(self.per_prompt),
        }


def _first_divergence(left: str, right: str) -> int:
    """Index of the first differing character, or -1 when the strings agree."""
    for index, (a, b) in enumerate(zip(left, right)):
        if a != b:
            return index
    return -1 if len(left) == len(right) else min(len(left), len(right))


def compute_drift(
    measurement: RecipeMeasurement,
    reference: RecipeMeasurement,
) -> DriftReport:
    """Compare one recipe's greedy outputs with the reference recipe's."""
    reference_outputs = reference.outputs_by_prompt()
    outputs = measurement.outputs_by_prompt()
    shared = sorted(set(outputs) & set(reference_outputs))

    per_prompt: list[dict[str, Any]] = []
    exact = 0
    similarities: list[float] = []
    for prompt_id in shared:
        got, want = outputs[prompt_id], reference_outputs[prompt_id]
        matches = got == want
        exact += matches
        similarity = round(SequenceMatcher(None, want, got).ratio(), 4)
        similarities.append(similarity)
        per_prompt.append({
            "prompt_id": prompt_id,
            "exact_match": matches,
            "similarity": similarity,
            "first_divergence_char": _first_divergence(want, got),
            "reference_text": want,
            "recipe_text": got,
        })

    return DriftReport(
        recipe_id=measurement.recipe.id,
        lane=measurement.recipe.lane,
        reference_id=reference.recipe.id,
        compared_prompts=len(shared),
        exact_match_rate=round(exact / len(shared), 4) if shared else 0.0,
        mean_similarity=_mean(similarities),
        advisory=measurement.recipe.lane is Lane.PERFORMANCE_FIT,
        per_prompt=tuple(per_prompt),
    )


class _PeakMemory:
    """Continuously sample a driver's memory while requests are in flight."""

    def __init__(self, driver: RecipeDriver) -> None:
        self._driver = driver
        self.peak_rss = 0
        self.peak_vram = 0

    def sample(self) -> None:
        try:
            rss, vram = self._driver.memory_probe()
        except Exception:  # a probe must never fail a benchmark run
            return
        self.peak_rss = max(self.peak_rss, rss)
        self.peak_vram = max(self.peak_vram, vram)

    def sample_until(self, stop: threading.Event, interval_s: float = 0.01) -> None:
        """Sample until ``stop`` is set, including one final post-request probe."""
        while not stop.wait(interval_s):
            self.sample()
        self.sample()


def _run_request(
    driver: RecipeDriver,
    recipe: RecipeSpec,
    prompt: PromptSpec,
    sampling: SamplingPolicy,
    concurrency: int,
    repeat: int,
    memory: _PeakMemory,
) -> RequestOutcome:
    started = time.monotonic()
    try:
        sample = driver.generate(prompt.text, sampling)
    except Exception as exc:  # a failed request is data, not a crashed benchmark
        return RequestOutcome(
            recipe_id=recipe.id,
            concurrency=concurrency,
            prompt_id=prompt.id,
            repeat=repeat,
            ok=False,
            latency_ms=round((time.monotonic() - started) * 1000, 4),
            error=f"{type(exc).__name__}: {exc}",
        )
    finally:
        memory.sample()

    return RequestOutcome(
        recipe_id=recipe.id,
        concurrency=concurrency,
        prompt_id=prompt.id,
        repeat=repeat,
        ok=True,
        latency_ms=round(sample.total_ms, 4),
        ttft_ms=round(sample.ttft_ms, 4),
        prefill_ms=round(sample.prefill_ms, 4),
        decode_ms=round(sample.decode_ms, 4),
        queue_wait_ms=round(sample.queue_wait_ms, 4),
        prompt_tokens=sample.prompt_tokens,
        decode_tokens=sample.decode_tokens,
        text=sample.text,
    )


def measure_recipe(
    driver: RecipeDriver,
    recipe: RecipeSpec,
    plan: BenchmarkPlan,
) -> RecipeMeasurement:
    """Load one recipe and run the whole plan against it.

    The driver is closed exactly once, whatever happens, so a recipe that dies at
    concurrency 4 still releases its weights before the next recipe loads.
    """
    load = driver.load()
    measurement = RecipeMeasurement(recipe=recipe, load=load)

    try:
        for _ in range(plan.warmup_requests):
            try:
                driver.generate(plan.prompts[0].text, plan.sampling)
            except Exception:  # a failing warmup is reported by the real requests
                break

        for concurrency in plan.concurrency_levels:
            # Measure each cell independently and sample while requests are in
            # flight; post-request probes alone miss transient KV/workspace peaks.
            memory = _PeakMemory(driver)
            memory.sample()
            stop_sampling = threading.Event()
            sampler = threading.Thread(
                target=memory.sample_until,
                args=(stop_sampling,),
                name=f"recipe-memory-{recipe.id}-c{concurrency}",
                daemon=True,
            )
            requests = [
                (prompt, repeat)
                for repeat in range(plan.repeats)
                for prompt in plan.prompts
                for _ in range(concurrency)
            ]
            started = time.monotonic()
            sampler.start()
            try:
                with ThreadPoolExecutor(max_workers=concurrency) as pool:
                    outcomes = list(pool.map(
                        lambda item: _run_request(
                            driver, recipe, item[0], plan.sampling, concurrency, item[1], memory
                        ),
                        requests,
                    ))
            finally:
                stop_sampling.set()
                sampler.join()
            wall_ms = (time.monotonic() - started) * 1000

            measurement.outcomes.extend(outcomes)
            measurement.metrics[concurrency] = summarize_concurrency(
                outcomes,
                concurrency=concurrency,
                wall_ms=wall_ms,
                peak_rss_bytes=memory.peak_rss,
                peak_vram_bytes=memory.peak_vram,
            )
    finally:
        driver.close()

    return measurement


def build_report(
    plan: BenchmarkPlan,
    measurements: Sequence[RecipeMeasurement],
    *,
    host: dict[str, Any],
    evidence_class: str,
    provenance: Mapping[str, Any] | None = None,
) -> dict:
    """Assemble the machine-readable benchmark document.

    ``evidence_class`` is one of ``synthetic``, ``local-real`` or
    ``multi-machine-real`` and is never inferred: a report that cannot say how it
    was produced cannot be trusted by a release gate.
    """
    if evidence_class not in {"synthetic", "local-real", "multi-machine-real"}:
        raise BenchmarkError(f"unknown evidence class {evidence_class!r}")
    if evidence_class != "synthetic" and not isinstance(provenance, Mapping):
        raise BenchmarkError("non-synthetic reports require canonical signed provenance")

    references = [m for m in measurements if m.recipe.is_reference]
    if len(references) != 1:
        raise BenchmarkError(
            f"exactly one reference recipe is required, got {len(references)}"
        )
    reference = references[0]
    if reference.recipe.lane is not Lane.QUALITY:
        raise BenchmarkError("the reference recipe must sit in the quality lane")

    drift = [
        compute_drift(measurement, reference).to_dict()
        for measurement in measurements
        if measurement is not reference and measurement.available
    ]
    report = {
        "schema_version": REPORT_SCHEMA_VERSION,
        "evidence_class": evidence_class,
        "plan": plan.to_dict(),
        "host": host,
        "reference_recipe_id": reference.recipe.id,
        "recipes": [measurement.to_dict() for measurement in measurements],
        "drift": drift,
    }
    if provenance is not None:
        report["provenance"] = dict(provenance)
    return report


def format_summary(report: dict) -> str:
    """Render the human-readable companion to the JSON artifact."""
    plan = report["plan"]
    lines = [
        f"Recipe benchmark {plan['plan_id']} ({report['evidence_class']})",
        f"model {plan['model_id']}@{plan['model_revision']}",
    ]
    for entry in report["recipes"]:
        recipe = entry["recipe"]
        if not entry["available"]:
            lines.append(f"{recipe['id']:38} UNAVAILABLE: {entry['unavailable_reason']}")
            continue
        artifact_gb = entry["load"]["artifact_bytes"] / 1e9
        for level, metrics in entry["concurrency"].items():
            lines.append(
                f"{recipe['id']:38} [{recipe['lane']:16}] c={level:>2} "
                f"ttft p50/p95 {metrics['ttft_p50_ms']:8.1f}/{metrics['ttft_p95_ms']:8.1f} ms; "
                f"prefill {metrics['prefill_tokens_per_sec']:7.1f} tok/s; "
                f"decode {metrics['decode_tokens_per_sec']:6.1f} tok/s; "
                f"aggregate {metrics['aggregate_decode_tokens_per_sec']:7.1f} tok/s; "
                f"rss {metrics['peak_rss_bytes'] / 1e9:5.2f} GB; "
                f"vram {metrics['peak_vram_bytes'] / 1e9:5.2f} GB; "
                f"artifact {artifact_gb:5.2f} GB; failures {metrics['failures']}"
            )
    for entry in report["drift"]:
        tag = "advisory" if entry["advisory"] else "gated"
        lines.append(
            f"drift {entry['recipe_id']:32} vs {entry['reference_id']:28} "
            f"exact {entry['exact_match_rate']:.2f}; similarity {entry['mean_similarity']:.3f} ({tag})"
        )
    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Run the controlled safetensors-versus-GGUF recipe benchmark"
    )
    parser.add_argument("--config", type=Path, required=True, help="benchmark configuration JSON")
    parser.add_argument(
        "--profile",
        choices=("contract-v1", "gpu-diagnostic"),
        default="contract-v1",
        help="validation and provenance profile (GPU diagnostics are not v1-eligible)",
    )
    parser.add_argument("--json-out", type=Path, help="write the JSON report to this path")
    parser.add_argument("--summary-out", type=Path, help="write the text summary to this path")
    args = parser.parse_args(argv)

    from .recipe_drivers import (  # heavy runtimes: import on demand
        run_configured_benchmark,
        run_configured_gpu_diagnostic,
    )

    runner = (
        run_configured_gpu_diagnostic
        if args.profile == "gpu-diagnostic"
        else run_configured_benchmark
    )
    report = runner(json.loads(args.config.read_text(encoding="utf-8")))
    summary = format_summary(report)
    if args.json_out:
        args.json_out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    if args.summary_out:
        args.summary_out.write_text(summary + "\n", encoding="utf-8")
    print(summary)
    return 0


if __name__ == "__main__":  # pragma: no cover - CLI entry point
    raise SystemExit(main())