"""The versioned safetensors-versus-GGUF performance contract.

The contract is the decision rule the native GGUF track is judged by, written
down *before* the numbers arrive and consumed later by the release gate
(DGR-014).  Its thresholds are ratios against the Transformers/safetensors
reference recipe rather than absolute tokens/sec, because the absolute figure is
a property of whichever machine ran the benchmark and would have to be re-argued
on every host; a ratio is a claim about the runtime.

Three rules give the contract its teeth:

* **Thresholds are locked.**  ``CONTRACT_SCHEMA_VERSION`` and ``locked_at``
  travel with the document.  Moving a threshold after seeing results is a new
  contract version and a human decision, not a tweak.
* **Only like-for-like comparisons count.**  A recipe measured on a different
  device than the reference is marked non-comparable and is granted no benefit,
  so a GPU-versus-CPU mismatch can never be laundered into a speed win.
* **Quantized recipes never claim numerical equivalence.**  Quality is gated on
  the near-lossless quality lane; the performance-fit lane is judged on speed,
  memory and fit alone.

The verdict is one of ``promote``, ``optimize`` or ``stop`` — the three outcomes
the release gate is allowed to reach.
"""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Mapping

from .recipe_benchmark import Lane

# Layout of the contract document understood by this reader.
CONTRACT_SCHEMA_VERSION = 1

VERDICT_PROMOTE = "promote"
VERDICT_OPTIMIZE = "optimize"
VERDICT_STOP = "stop"


class PerformanceContractError(ValueError):
    """Raised when a contract is missing, malformed, or of an unsupported version."""


@dataclass(frozen=True)
class ContractThresholds:
    """The locked decision thresholds.

    Every value is a ratio of a GGUF recipe's metric to the reference recipe's
    metric on the same machine, same device, same plan.

    A *meaningful speed benefit* means the GGUF recipe decodes at least 25%
    faster for a single request without making time-to-first-token materially
    worse, or sustains at least 25% more aggregate throughput under concurrency.
    Either route is a real win for the product: one helps a single user, the
    other helps a loaded node.

    A *meaningful fit benefit* means peak resident memory (RSS plus VRAM) drops
    by at least 25%.  Fit is the product thesis — models larger than one
    consumer node — so it is measured in resident bytes, not in how small the
    file on disk is.  Artifact size has its own reported threshold because a
    smaller download is a real but secondary good.

    25% is chosen to sit well clear of ordinary run-to-run variance on a busy
    developer machine while still being a benefit a user would notice.  A 5%
    edge would not justify owning a native runtime and a patch stack.
    """

    min_decode_speedup: float = 1.25
    max_ttft_ratio: float = 1.25
    min_aggregate_throughput_speedup: float = 1.25
    max_resident_memory_ratio: float = 0.75
    max_artifact_size_ratio: float = 0.60
    min_quality_exact_match_rate: float = 0.90
    min_quality_mean_similarity: float = 0.97
    max_failure_rate: float = 0.0

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


@dataclass(frozen=True)
class PerformanceContract:
    """A locked, versioned contract plus the baseline it was locked against."""

    contract_version: int
    locked_at: str
    locked_by: str
    plan_id: str
    thresholds: ContractThresholds
    baseline: Mapping[str, Any]
    stop_condition: str
    notes: str = ""
    schema_version: int = CONTRACT_SCHEMA_VERSION

    def to_dict(self) -> dict:
        return {
            "schema_version": self.schema_version,
            "contract_version": self.contract_version,
            "locked_at": self.locked_at,
            "locked_by": self.locked_by,
            "plan_id": self.plan_id,
            "thresholds": self.thresholds.to_dict(),
            "baseline": dict(self.baseline),
            "stop_condition": self.stop_condition,
            "notes": self.notes,
        }


STOP_CONDITION = (
    "Stop the native llama.cpp/GGUF track when, on the same machine and device "
    "as the Transformers/safetensors reference and under this plan, no "
    "performance-fit GGUF recipe delivers either a meaningful speed benefit "
    "(>=25% higher single-request decode tokens/sec without a >25% worse TTFT, "
    "or >=25% higher aggregate throughput under concurrency) or a meaningful fit "
    "benefit (>=25% lower peak resident memory), or when the near-lossless "
    "quality lane fails, which indicates a broken runtime rather than a "
    "quantization trade-off."
)


def _recipe_entries(report: Mapping[str, Any]) -> dict[str, Mapping[str, Any]]:
    return {entry["recipe"]["id"]: entry for entry in report["recipes"]}


def _cell(entry: Mapping[str, Any], concurrency: int) -> Mapping[str, Any] | None:
    return entry["concurrency"].get(str(concurrency))


def _resident_bytes(cell: Mapping[str, Any]) -> int:
    return int(cell["peak_rss_bytes"]) + int(cell["peak_vram_bytes"])


def _ratio(value: float, reference: float) -> float:
    """Ratio guarded against a zero reference, which means "not measured"."""
    if reference <= 0:
        return 0.0
    return round(value / reference, 4)


@dataclass(frozen=True)
class RecipeEvaluation:
    """How one GGUF recipe fared against the reference under the contract."""

    recipe_id: str
    lane: str
    comparable: bool
    incomparable_reason: str
    speed_benefit: bool
    fit_benefit: bool
    quality_pass: bool | None
    failures: int
    measurements: Mapping[str, Any]
    reasons: tuple[str, ...]

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


@dataclass(frozen=True)
class ContractEvaluation:
    """The release-gate answer: a verdict plus every reason behind it."""

    contract_version: int
    plan_id: str
    verdict: str
    quality_lane_pass: bool
    speed_benefit: bool
    fit_benefit: bool
    stop_condition_met: bool
    recipes: tuple[RecipeEvaluation, ...]
    rationale: tuple[str, ...]

    def to_dict(self) -> dict:
        return {
            "contract_version": self.contract_version,
            "plan_id": self.plan_id,
            "verdict": self.verdict,
            "quality_lane_pass": self.quality_lane_pass,
            "speed_benefit": self.speed_benefit,
            "fit_benefit": self.fit_benefit,
            "stop_condition_met": self.stop_condition_met,
            "recipes": [recipe.to_dict() for recipe in self.recipes],
            "rationale": list(self.rationale),
        }


def _evaluate_recipe(
    entry: Mapping[str, Any],
    reference: Mapping[str, Any],
    drift_by_recipe: Mapping[str, Mapping[str, Any]],
    thresholds: ContractThresholds,
    concurrency_levels: list[int],
) -> RecipeEvaluation:
    recipe = entry["recipe"]
    lane = recipe["lane"]
    reasons: list[str] = []

    if not entry["available"]:
        return RecipeEvaluation(
            recipe_id=recipe["id"], lane=lane, comparable=False,
            incomparable_reason=entry["unavailable_reason"] or "recipe was not measured",
            speed_benefit=False, fit_benefit=False, quality_pass=None, failures=0,
            measurements={}, reasons=("recipe unavailable; no benefit granted",),
        )

    if recipe["device"] != reference["recipe"]["device"]:
        return RecipeEvaluation(
            recipe_id=recipe["id"], lane=lane, comparable=False,
            incomparable_reason=(
                f"recipe ran on device {recipe['device']!r} but the reference ran on "
                f"{reference['recipe']['device']!r}; a cross-device ratio is not a runtime result"
            ),
            speed_benefit=False, fit_benefit=False, quality_pass=None, failures=0,
            measurements={}, reasons=("cross-device comparison; no benefit granted",),
        )

    single = _cell(entry, 1)
    reference_single = _cell(reference, 1)
    measurements: dict[str, Any] = {}
    speed_benefit = False

    if single and reference_single:
        decode_speedup = _ratio(
            single["decode_tokens_per_sec"], reference_single["decode_tokens_per_sec"]
        )
        ttft_ratio = _ratio(single["ttft_p50_ms"], reference_single["ttft_p50_ms"])
        measurements["decode_speedup"] = decode_speedup
        measurements["ttft_ratio"] = ttft_ratio
        single_request_win = (
            decode_speedup >= thresholds.min_decode_speedup
            and 0 < ttft_ratio <= thresholds.max_ttft_ratio
        )
        if single_request_win:
            speed_benefit = True
            reasons.append(
                f"single-request decode {decode_speedup:.2f}x reference "
                f"(>= {thresholds.min_decode_speedup:.2f}x) at TTFT ratio {ttft_ratio:.2f}"
            )
        else:
            reasons.append(
                f"no single-request speed win: decode {decode_speedup:.2f}x, TTFT {ttft_ratio:.2f}x"
            )

    concurrent = [level for level in concurrency_levels if level > 1]
    if concurrent:
        top = max(concurrent)
        cell, reference_cell = _cell(entry, top), _cell(reference, top)
        if cell and reference_cell:
            aggregate_speedup = _ratio(
                cell["aggregate_decode_tokens_per_sec"],
                reference_cell["aggregate_decode_tokens_per_sec"],
            )
            measurements["aggregate_throughput_speedup"] = aggregate_speedup
            measurements["aggregate_concurrency"] = top
            if aggregate_speedup >= thresholds.min_aggregate_throughput_speedup:
                speed_benefit = True
                reasons.append(
                    f"aggregate throughput at concurrency {top} is {aggregate_speedup:.2f}x reference "
                    f"(>= {thresholds.min_aggregate_throughput_speedup:.2f}x)"
                )
            else:
                reasons.append(
                    f"no concurrency speed win: aggregate throughput at {top} is "
                    f"{aggregate_speedup:.2f}x reference"
                )

    fit_benefit = False
    if single and reference_single:
        resident_ratio = _ratio(_resident_bytes(single), _resident_bytes(reference_single))
        artifact_ratio = _ratio(
            entry["load"]["artifact_bytes"], reference["load"]["artifact_bytes"]
        )
        measurements["resident_memory_ratio"] = resident_ratio
        measurements["artifact_size_ratio"] = artifact_ratio
        if 0 < resident_ratio <= thresholds.max_resident_memory_ratio:
            fit_benefit = True
            reasons.append(
                f"peak resident memory is {resident_ratio:.2f}x reference "
                f"(<= {thresholds.max_resident_memory_ratio:.2f}x)"
            )
        else:
            reasons.append(f"no fit win: peak resident memory is {resident_ratio:.2f}x reference")
        measurements["artifact_size_win"] = (
            0 < artifact_ratio <= thresholds.max_artifact_size_ratio
        )

    failures = sum(cell["failures"] for cell in entry["concurrency"].values())
    requests = sum(cell["requests"] for cell in entry["concurrency"].values())
    failure_rate = _ratio(failures, requests) if requests else 0.0
    measurements["failure_rate"] = failure_rate
    if failure_rate > thresholds.max_failure_rate:
        reasons.append(f"failure rate {failure_rate:.2%} exceeds the contract limit")
        speed_benefit = False
        fit_benefit = False

    # Quality is a claim only the near-lossless lane is allowed to make.  A
    # quantized recipe's drift is recorded elsewhere and deliberately not read
    # here: Q4 disagreeing with bf16 is the trade-off, not a failure.
    quality_pass: bool | None = None
    if lane == Lane.QUALITY.value:
        drift = drift_by_recipe.get(recipe["id"])
        if drift is None:
            quality_pass = False
            reasons.append("quality-lane recipe has no drift measurement against the reference")
        else:
            quality_pass = (
                drift["exact_match_rate"] >= thresholds.min_quality_exact_match_rate
                and drift["mean_similarity"] >= thresholds.min_quality_mean_similarity
            )
            measurements["exact_match_rate"] = drift["exact_match_rate"]
            measurements["mean_similarity"] = drift["mean_similarity"]
            reasons.append(
                f"quality lane exact-match {drift['exact_match_rate']:.2f} / similarity "
                f"{drift['mean_similarity']:.3f} versus the reference "
                f"({'pass' if quality_pass else 'fail'})"
            )

    return RecipeEvaluation(
        recipe_id=recipe["id"], lane=lane, comparable=True, incomparable_reason="",
        speed_benefit=speed_benefit, fit_benefit=fit_benefit, quality_pass=quality_pass,
        failures=failures, measurements=measurements, reasons=tuple(reasons),
    )


def evaluate_contract(
    contract: PerformanceContract,
    report: Mapping[str, Any],
) -> ContractEvaluation:
    """Judge a benchmark report against the locked contract.

    Only performance-fit recipes can earn a speed or fit benefit; the quality
    lane decides only whether the GGUF runtime is numerically sane.  A GGUF
    runtime that fails the quality lane is broken, and no amount of speed
    redeems it, so the verdict is ``stop`` regardless of the other numbers.
    """
    entries = _recipe_entries(report)
    reference_id = report["reference_recipe_id"]
    reference = entries.get(reference_id)
    if reference is None:
        raise PerformanceContractError(
            f"report names reference recipe {reference_id!r}, which it does not contain"
        )

    drift_by_recipe = {entry["recipe_id"]: entry for entry in report["drift"]}
    concurrency_levels = list(report["plan"]["concurrency_levels"])

    evaluations = tuple(
        _evaluate_recipe(entry, reference, drift_by_recipe, contract.thresholds, concurrency_levels)
        for recipe_id, entry in entries.items()
        if recipe_id != reference_id
    )

    quality_lane = [
        evaluation for evaluation in evaluations
        if evaluation.lane == Lane.QUALITY.value and evaluation.comparable
    ]
    quality_lane_pass = bool(quality_lane) and all(
        evaluation.quality_pass for evaluation in quality_lane
    )
    performance_lane = [
        evaluation for evaluation in evaluations
        if evaluation.lane == Lane.PERFORMANCE_FIT.value
    ]
    speed_benefit = any(evaluation.speed_benefit for evaluation in performance_lane)
    fit_benefit = any(evaluation.fit_benefit for evaluation in performance_lane)

    rationale: list[str] = []
    if not quality_lane:
        rationale.append(
            "no comparable near-lossless GGUF recipe was measured, so the runtime's "
            "numerical correctness is unproven"
        )
    elif not quality_lane_pass:
        rationale.append(
            "the near-lossless quality lane failed: the GGUF runtime disagrees with the "
            "safetensors reference beyond what near-lossless weights can explain"
        )
    else:
        rationale.append("the near-lossless quality lane passed against the safetensors reference")

    rationale.append(
        "a meaningful speed benefit was measured" if speed_benefit
        else "no performance-fit recipe delivered a meaningful speed benefit"
    )
    rationale.append(
        "a meaningful fit benefit was measured" if fit_benefit
        else "no performance-fit recipe delivered a meaningful fit benefit"
    )

    stop_condition_met = not quality_lane_pass or not (speed_benefit or fit_benefit)
    if stop_condition_met:
        verdict = VERDICT_STOP
    elif speed_benefit and fit_benefit:
        verdict = VERDICT_PROMOTE
    else:
        verdict = VERDICT_OPTIMIZE
        rationale.append(
            "exactly one of speed or fit cleared the contract: the benefit is real but partial, "
            "so the measured bottleneck needs a bounded optimization task before promotion"
        )

    return ContractEvaluation(
        contract_version=contract.contract_version,
        plan_id=contract.plan_id,
        verdict=verdict,
        quality_lane_pass=quality_lane_pass,
        speed_benefit=speed_benefit,
        fit_benefit=fit_benefit,
        stop_condition_met=stop_condition_met,
        recipes=evaluations,
        rationale=tuple(rationale),
    )


def parse_contract(data: Any, source: str = "<memory>") -> PerformanceContract:
    """Validate an already-decoded contract document."""
    if not isinstance(data, Mapping):
        raise PerformanceContractError(f"contract root in {source} must be a JSON object")

    schema_version = data.get("schema_version")
    if not isinstance(schema_version, int) or isinstance(schema_version, bool):
        raise PerformanceContractError(f"'schema_version' in {source} must be an integer")
    if schema_version != CONTRACT_SCHEMA_VERSION:
        raise PerformanceContractError(
            f"{source} declares contract schema version {schema_version}, but this node reads "
            f"version {CONTRACT_SCHEMA_VERSION}; upgrade the node or use a supported contract"
        )

    for required in ("contract_version", "locked_at", "locked_by", "plan_id", "stop_condition"):
        if not data.get(required):
            raise PerformanceContractError(f"{source} is missing {required!r}")

    raw_thresholds = data.get("thresholds")
    if not isinstance(raw_thresholds, Mapping):
        raise PerformanceContractError(f"'thresholds' in {source} must be a JSON object")
    known = {field for field in ContractThresholds().to_dict()}
    unknown = set(raw_thresholds) - known
    if unknown:
        raise PerformanceContractError(
            f"{source} carries unknown thresholds {sorted(unknown)}; this node enforces {sorted(known)}"
        )
    thresholds = ContractThresholds(**{
        key: float(value) for key, value in raw_thresholds.items()
    })

    return PerformanceContract(
        contract_version=int(data["contract_version"]),
        locked_at=str(data["locked_at"]),
        locked_by=str(data["locked_by"]),
        plan_id=str(data["plan_id"]),
        thresholds=thresholds,
        baseline=dict(data.get("baseline") or {}),
        stop_condition=str(data["stop_condition"]),
        notes=str(data.get("notes", "")),
        schema_version=schema_version,
    )


def load_contract(path: Path) -> PerformanceContract:
    """Load and validate the contract at ``path``."""
    try:
        raw = path.read_text(encoding="utf-8")
    except OSError as exc:
        raise PerformanceContractError(
            f"cannot read performance contract {path}: {exc.strerror or exc}"
        ) from exc
    try:
        data = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise PerformanceContractError(
            f"{path} is not valid JSON: {exc.msg} at line {exc.lineno} column {exc.colno}"
        ) from exc
    return parse_contract(data, source=str(path))


def baseline_from_report(report: Mapping[str, Any]) -> dict[str, Any]:
    """Distil the reference numbers a later gate needs to compare against."""
    entries = _recipe_entries(report)
    baseline: dict[str, Any] = {
        "evidence_class": report["evidence_class"],
        "model_id": report["plan"]["model_id"],
        "model_revision": report["plan"]["model_revision"],
        "reference_recipe_id": report["reference_recipe_id"],
        "host": report["host"],
        "recipes": {},
    }
    for recipe_id, entry in entries.items():
        if not entry["available"]:
            baseline["recipes"][recipe_id] = {"available": False,
                                              "reason": entry["unavailable_reason"]}
            continue
        baseline["recipes"][recipe_id] = {
            "available": True,
            "lane": entry["recipe"]["lane"],
            "device": entry["recipe"]["device"],
            "artifact_bytes": entry["load"]["artifact_bytes"],
            "concurrency": {
                level: {
                    "ttft_p50_ms": cell["ttft_p50_ms"],
                    "ttft_p95_ms": cell["ttft_p95_ms"],
                    "latency_p50_ms": cell["latency_p50_ms"],
                    "latency_p95_ms": cell["latency_p95_ms"],
                    "prefill_tokens_per_sec": cell["prefill_tokens_per_sec"],
                    "decode_tokens_per_sec": cell["decode_tokens_per_sec"],
                    "aggregate_decode_tokens_per_sec": cell["aggregate_decode_tokens_per_sec"],
                    "peak_rss_bytes": cell["peak_rss_bytes"],
                    "peak_vram_bytes": cell["peak_vram_bytes"],
                    "failures": cell["failures"],
                }
                for level, cell in entry["concurrency"].items()
            },
        }
    return baseline