"""TOPLOC activation proof helpers for validator-side audits."""

from __future__ import annotations

from dataclasses import dataclass
from importlib import import_module
from typing import Any, Literal


ProofEncoding = Literal["base64", "bytes"]


@dataclass(frozen=True)
class ToplocAuditConfig:
    """Canonical audit parameters for one model preset."""

    dtype: str = "bfloat16"
    quantization: str = "bfloat16"
    decode_batching_size: int = 32
    topk: int = 8
    skip_prefill: bool = True
    encoding: ProofEncoding = "base64"
    # ADR-0018 §3: nodes retain boundary activations only briefly; a commitment
    # older than this can no longer be verified against a live node and must
    # fall back to the text-only audit path.
    commitment_ttl_seconds: float = 30.0


@dataclass(frozen=True)
class ToplocProofClaim:
    """Prover-provided TOPLOC proof and the parameters it was built with."""

    proofs: Any
    dtype: str
    quantization: str
    decode_batching_size: int
    topk: int
    skip_prefill: bool = True
    encoding: ProofEncoding = "base64"

    @classmethod
    def from_mapping(cls, value: dict[str, Any]) -> "ToplocProofClaim":
        return cls(
            proofs=value["proofs"],
            dtype=str(value.get("dtype", "bfloat16")),
            quantization=str(value.get("quantization", "bfloat16")),
            decode_batching_size=int(value.get("decode_batching_size", 32)),
            topk=int(value.get("topk", 8)),
            skip_prefill=bool(value.get("skip_prefill", True)),
            encoding=_proof_encoding(value.get("encoding", "base64")),
        )

    def as_mapping(self) -> dict[str, Any]:
        return {
            "proofs": self.proofs,
            "dtype": self.dtype,
            "quantization": self.quantization,
            "decode_batching_size": self.decode_batching_size,
            "topk": self.topk,
            "skip_prefill": self.skip_prefill,
            "encoding": self.encoding,
        }


def build_activation_proofs(
    activations: list[Any],
    *,
    config: ToplocAuditConfig | None = None,
    backend: Any | None = None,
) -> ToplocProofClaim:
    """Build a TOPLOC proof claim from captured activation tensors."""
    cfg = config or ToplocAuditConfig()
    module = backend or _load_toploc()
    function_name = f"build_proofs_{cfg.encoding}"
    build = getattr(module, function_name)
    proofs = _call_toploc(
        build,
        activations,
        decode_batching_size=cfg.decode_batching_size,
        topk=cfg.topk,
        skip_prefill=cfg.skip_prefill,
    )
    return ToplocProofClaim(
        proofs=proofs,
        dtype=cfg.dtype,
        quantization=cfg.quantization,
        decode_batching_size=cfg.decode_batching_size,
        topk=cfg.topk,
        skip_prefill=cfg.skip_prefill,
        encoding=cfg.encoding,
    )


@dataclass(frozen=True)
class ToplocVerificationResult:
    """Verification outcome plus the raw TOPLOC divergence metric.

    The `toploc` library's `verify_proofs_*` returns a bool for simple
    prover/verifier config mismatches, but for a real activation comparison
    it returns one `VerificationResult(exp_intersections, mant_err_mean,
    mant_err_median)` per chunk (README §"What it actually is"). Historically
    only `bool(result)` was kept, which is always true for a non-empty list
    of results regardless of how divergent they are (AH-021 gap #1). This
    dataclass surfaces the raw per-chunk metrics (aggregated: worst-case
    `exp_intersections`, mean `mant_err_mean`/`mant_err_median`) so a
    calibration corpus can be built before any threshold is trusted.
    """

    passed: bool
    exp_intersections: float | None = None
    mant_err_mean: float | None = None
    mant_err_median: float | None = None
    chunk_count: int = 0


def verify_activation_proofs(
    reference_activations: list[Any],
    claim: ToplocProofClaim,
    *,
    config: ToplocAuditConfig | None = None,
    backend: Any | None = None,
) -> bool:
    """Verify prover TOPLOC proofs against reference teacher-forced activations."""
    return verify_activation_proofs_detailed(
        reference_activations, claim, config=config, backend=backend,
    ).passed


def verify_activation_proofs_detailed(
    reference_activations: list[Any],
    claim: ToplocProofClaim,
    *,
    config: ToplocAuditConfig | None = None,
    backend: Any | None = None,
) -> ToplocVerificationResult:
    """Verify prover TOPLOC proofs and surface the raw divergence metric.

    Same pass/fail contract as `verify_activation_proofs` (kept as a thin
    wrapper for existing call sites); this is the entry point for anything
    that needs the underlying distance value, e.g. the AH-021 honest-noise
    calibration corpus.
    """
    cfg = config or ToplocAuditConfig(
        dtype=claim.dtype,
        quantization=claim.quantization,
        decode_batching_size=claim.decode_batching_size,
        topk=claim.topk,
        skip_prefill=claim.skip_prefill,
        encoding=claim.encoding,
    )
    if claim.dtype != cfg.dtype or claim.quantization != cfg.quantization:
        return ToplocVerificationResult(passed=False)
    if claim.decode_batching_size != cfg.decode_batching_size or claim.topk != cfg.topk:
        return ToplocVerificationResult(passed=False)
    if claim.skip_prefill != cfg.skip_prefill or claim.encoding != cfg.encoding:
        return ToplocVerificationResult(passed=False)

    module = backend or _load_toploc()
    function_name = f"verify_proofs_{claim.encoding}"
    verify = getattr(module, function_name)
    raw = _call_toploc(
        verify,
        reference_activations,
        claim.proofs,
        decode_batching_size=claim.decode_batching_size,
        topk=claim.topk,
        skip_prefill=claim.skip_prefill,
    )
    divergence = _extract_divergence(raw)
    return ToplocVerificationResult(passed=bool(raw), **divergence)


def _extract_divergence(raw: Any) -> dict[str, Any]:
    """Aggregate per-chunk TOPLOC `VerificationResult`s, if present.

    `raw` is a plain bool for the simple fake backends used in existing unit
    tests (no per-chunk metric available). The real `toploc` library returns
    a list of per-chunk results; `exp_intersections` is aggregated by min
    (worst honest-noise case across chunks) and the mantissa errors by mean.
    """
    chunks = raw if isinstance(raw, (list, tuple)) else None
    if not chunks:
        return {"exp_intersections": None, "mant_err_mean": None, "mant_err_median": None, "chunk_count": 0}
    exp_vals = [v for v in (_chunk_field(c, "exp_intersections") for c in chunks) if v is not None]
    mean_vals = [v for v in (_chunk_field(c, "mant_err_mean") for c in chunks) if v is not None]
    median_vals = [v for v in (_chunk_field(c, "mant_err_median") for c in chunks) if v is not None]
    return {
        "exp_intersections": min(exp_vals) if exp_vals else None,
        "mant_err_mean": (sum(mean_vals) / len(mean_vals)) if mean_vals else None,
        "mant_err_median": (sum(median_vals) / len(median_vals)) if median_vals else None,
        "chunk_count": len(chunks),
    }


def _chunk_field(chunk: Any, name: str) -> float | None:
    value = chunk.get(name) if isinstance(chunk, dict) else getattr(chunk, name, None)
    return float(value) if isinstance(value, (int, float)) else None


def _load_toploc() -> Any:
    try:
        return import_module("toploc")
    except ModuleNotFoundError as exc:
        raise RuntimeError(
            "toploc is required for activation proof audits; install meshnet-validator with dependencies"
        ) from exc


def _call_toploc(function: Any, activations: list[Any], *args: Any, **kwargs: Any) -> Any:
    try:
        return function(activations, *args, **kwargs)
    except TypeError:
        if kwargs:
            ordered = [
                kwargs["decode_batching_size"],
                kwargs["topk"],
                kwargs["skip_prefill"],
            ]
            return function(activations, *args, *ordered)
        raise


def _proof_encoding(value: object) -> ProofEncoding:
    if value == "bytes":
        return "bytes"
    return "base64"


__all__ = [
    "ToplocAuditConfig",
    "ToplocProofClaim",
    "ToplocVerificationResult",
    "build_activation_proofs",
    "verify_activation_proofs",
    "verify_activation_proofs_detailed",
]
