"""AH-006: validator TOPLOC audit primitive."""

from __future__ import annotations

from collections import namedtuple
from types import SimpleNamespace

from meshnet_validator import ToplocAuditConfig, ValidatorProcess
from meshnet_validator.audit import (
    build_activation_proofs,
    verify_activation_proofs,
    verify_activation_proofs_detailed,
)


class FakeToploc:
    def __init__(self) -> None:
        self.build_calls: list[dict] = []
        self.verify_calls: list[dict] = []

    def build_proofs_base64(
        self,
        activations,
        *,
        decode_batching_size: int,
        topk: int,
        skip_prefill: bool,
    ):
        self.build_calls.append({
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        })
        return {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        }

    def verify_proofs_base64(
        self,
        activations,
        proofs,
        *,
        decode_batching_size: int,
        topk: int,
        skip_prefill: bool,
    ):
        self.verify_calls.append({
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        })
        return proofs == {
            "activation_fingerprint": tuple(tuple(row) for row in activations),
            "decode_batching_size": decode_batching_size,
            "topk": topk,
            "skip_prefill": skip_prefill,
        }


class FakeValidationLog:
    def __init__(self, events) -> None:
        self._events = events

    def list_completed_inferences(self, *, after_index: int = -1):
        return [event for event in self._events if event.index > after_index]


class FakeRegistry:
    def __init__(self) -> None:
        self.slashes: list[dict] = []
        self.audit_outcomes: list[dict] = []

    def get_wallet(self, wallet_address: str):
        return SimpleNamespace(banned=False)

    def submit_slash_proof(self, **kwargs):
        self.slashes.append(kwargs)
        return kwargs

    def record_audit_outcome(self, wallet_address: str, *, passed: bool):
        self.audit_outcomes.append({"wallet_address": wallet_address, "passed": passed})


class FakeContracts:
    def __init__(self, events) -> None:
        self.validation = FakeValidationLog(events)
        self.registry = FakeRegistry()


class TeacherForcedValidator(ValidatorProcess):
    def __init__(self, *args, reference_activations, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.reference_activations = reference_activations
        self.teacher_forced_calls: list[dict] = []

    def _run_reference(self, messages: list[dict]) -> str:
        raise AssertionError("TOPLOC audits must not free-generate reference text")

    def _run_teacher_forced_prefill(self, **kwargs):
        self.teacher_forced_calls.append(kwargs)
        return self.reference_activations


def test_stub_activation_tensors_round_trip_through_toploc_proofs():
    "Stub activation tensors round trip through toploc proofs\n\nTags: audit, calibration"
    fake_toploc = FakeToploc()
    activations = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)

    claim = build_activation_proofs(activations, config=config, backend=fake_toploc)

    assert verify_activation_proofs(activations, claim, config=config, backend=fake_toploc) is True
    assert fake_toploc.build_calls == [{
        "decode_batching_size": 16,
        "topk": 2,
        "skip_prefill": True,
    }]
    assert fake_toploc.verify_calls == [{
        "decode_batching_size": 16,
        "topk": 2,
        "skip_prefill": True,
    }]


def test_validator_teacher_forces_claimed_tokens_for_toploc_audit():
    "Validator teacher forces claimed tokens for toploc audit\n\nTags: audit, calibration"
    fake_toploc = FakeToploc()
    activations = [[0.25, 0.5], [0.75, 1.0]]
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    claim = build_activation_proofs(activations, config=config, backend=fake_toploc)
    event = SimpleNamespace(
        index=0,
        session_id="session-toploc-ok",
        model="Qwen2.5-0.5B-Instruct",
        messages=[{"role": "user", "content": "hello"}],
        observed_output="honest but hardware-divergent text",
        route_nodes=[{"wallet_address": "wallet-good", "shard_end": 31}],
        claimed_token_ids=[101, 202, 303],
        toploc_proof=claim.as_mapping(),
    )
    contracts = FakeContracts([event])
    validator = TeacherForcedValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=fake_toploc,
        reference_activations=activations,
    )

    assert validator.validate_once() == []
    assert contracts.registry.slashes == []
    assert validator.teacher_forced_calls == [{
        "model": "Qwen2.5-0.5B-Instruct",
        "messages": [{"role": "user", "content": "hello"}],
        "claimed_token_ids": [101, 202, 303],
        "claim": claim,
    }]


def test_validator_rejects_swapped_precision_toploc_claim():
    "Validator rejects swapped precision toploc claim\n\nTags: audit, calibration"
    fake_toploc = FakeToploc()
    activations = [[0.25, 0.5], [0.75, 1.0]]
    canonical = ToplocAuditConfig(
        dtype="bfloat16",
        quantization="bfloat16",
        topk=2,
        decode_batching_size=16,
    )
    swapped = ToplocAuditConfig(
        dtype="bfloat16",
        quantization="int8",
        topk=2,
        decode_batching_size=16,
    )
    claim = build_activation_proofs(activations, config=swapped, backend=fake_toploc)
    event = SimpleNamespace(
        index=0,
        session_id="session-toploc-bad",
        model="Qwen2.5-0.5B-Instruct",
        messages=[{"role": "user", "content": "hello"}],
        observed_output="looks plausible",
        route_nodes=[{"wallet_address": "wallet-bad", "shard_end": 31}],
        audit={
            "claimed_token_ids": [101, 202, 303],
            "toploc_proof": claim.as_mapping(),
        },
    )
    contracts = FakeContracts([event])
    validator = TeacherForcedValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=canonical,
        toploc_backend=fake_toploc,
        reference_activations=activations,
    )

    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert contracts.registry.slashes[0]["wallet_address"] == "wallet-bad"
    assert "TOPLOC activation proof mismatch" in contracts.registry.slashes[0]["reason"]


# AH-021: verify_activation_proofs_detailed surfaces the raw divergence
# metric a calibration corpus needs, instead of only a pass/fail bool.

ChunkResult = namedtuple("ChunkResult", ["exp_intersections", "mant_err_mean", "mant_err_median"])


class FakeToplocWithChunkResults:
    """Mimics the real `toploc` library: verify returns per-chunk results,
    not a bool, so `bool(result)` alone (the AH-021 gap #1 bug) is always
    true for any non-empty response regardless of divergence."""

    def build_proofs_base64(self, activations, *, decode_batching_size, topk, skip_prefill):
        return {"activations": activations}

    def verify_proofs_base64(self, activations, proofs, *, decode_batching_size, topk, skip_prefill):
        return [
            ChunkResult(exp_intersections=8, mant_err_mean=0.01, mant_err_median=0.008),
            ChunkResult(exp_intersections=6, mant_err_mean=0.03, mant_err_median=0.02),
        ]


def test_verify_activation_proofs_detailed_aggregates_per_chunk_divergence():
    "Verify activation proofs detailed aggregates per chunk divergence\n\nTags: audit, calibration"
    fake_toploc = FakeToplocWithChunkResults()
    activations = [[1.0, 2.0], [3.0, 4.0]]
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    claim = build_activation_proofs(activations, config=config, backend=fake_toploc)

    result = verify_activation_proofs_detailed(activations, claim, config=config, backend=fake_toploc)

    assert result.passed is True  # non-empty list is truthy, same as legacy behavior
    assert result.chunk_count == 2
    assert result.exp_intersections == 6  # worst-case (min) across chunks
    assert result.mant_err_mean == 0.02  # mean of per-chunk means
    assert result.mant_err_median == 0.014  # mean of per-chunk medians

    # verify_activation_proofs still returns just the bool for existing callers.
    assert verify_activation_proofs(activations, claim, config=config, backend=fake_toploc) is True


def test_verify_activation_proofs_detailed_no_metric_from_plain_bool_backend():
    "Verify activation proofs detailed no metric from plain bool backend\n\nTags: audit, calibration"
    fake_toploc = FakeToploc()
    activations = [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    claim = build_activation_proofs(activations, config=config, backend=fake_toploc)

    result = verify_activation_proofs_detailed(activations, claim, config=config, backend=fake_toploc)

    assert result.passed is True
    assert result.chunk_count == 0
    assert result.exp_intersections is None
    assert result.mant_err_mean is None
    assert result.mant_err_median is None


def test_verify_activation_proofs_detailed_rejects_config_mismatch_without_calling_backend():
    "Verify activation proofs detailed rejects config mismatch without calling backend\n\nTags: audit, calibration"
    fake_toploc = FakeToplocWithChunkResults()
    activations = [[1.0, 2.0]]
    canonical = ToplocAuditConfig(dtype="bfloat16", quantization="bfloat16", topk=2, decode_batching_size=16)
    swapped = ToplocAuditConfig(dtype="bfloat16", quantization="int8", topk=2, decode_batching_size=16)
    claim = build_activation_proofs(activations, config=swapped, backend=fake_toploc)

    result = verify_activation_proofs_detailed(activations, claim, config=canonical, backend=fake_toploc)

    assert result.passed is False
    assert result.chunk_count == 0