"""AH-007: on-demand hop-boundary commitments + first-divergent-hop bisection.

`_final_text_node` (blame `max(shard_end)`) is wrong for multi-hop pipelines:
a corrupt early hop still produces a route whose *last* hop reported perfectly
honest activations. These tests build a two-hop route, corrupt one hop's
committed fingerprint, and assert the referee blames that hop specifically —
not always the last one.
"""

from __future__ import annotations

from types import SimpleNamespace

from meshnet_validator import ToplocAuditConfig, ValidatorProcess
from meshnet_validator.audit import build_activation_proofs


class FakeToploc:
    """Fingerprint = the activations themselves; verify is exact equality,
    matching the encode/verify contract exercised in test_toploc_audit.py."""

    def build_proofs_base64(self, activations, *, decode_batching_size, topk, 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, topk, 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()


def _hop_route_event(*, index, hop0_claim_activations, hop1_claim_activations, config, backend, ts=None):
    """Two-hop route where each hop carries its own TOPLOC commitment built
    from the given (possibly corrupted) activations."""
    hop0_claim = build_activation_proofs(hop0_claim_activations, config=config, backend=backend)
    hop1_claim = build_activation_proofs(hop1_claim_activations, config=config, backend=backend)
    route_nodes = [
        {
            "wallet_address": "wallet-hop0",
            "shard_start": 0,
            "shard_end": 15,
            "toploc_proof": hop0_claim.as_mapping(),
        },
        {
            "wallet_address": "wallet-hop1",
            "shard_start": 16,
            "shard_end": 31,
            "toploc_proof": hop1_claim.as_mapping(),
        },
    ]
    event = SimpleNamespace(
        index=index,
        session_id=f"session-hop-{index}",
        model="Qwen2.5-0.5B-Instruct",
        messages=[{"role": "user", "content": "hello"}],
        observed_output="two-hop pipeline output",
        route_nodes=route_nodes,
        claimed_token_ids=[101, 202, 303],
    )
    if ts is not None:
        event.ts = ts
    return event


class HopReferenceValidator(ValidatorProcess):
    """Stands in for the reference node: returns canned ground-truth
    activations at each hop cut-point instead of making an HTTP call."""

    def __init__(self, *args, reference_activations_by_hop, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.reference_activations_by_hop = reference_activations_by_hop
        self.hops_calls: list[dict] = []
        self.text_reference_calls: list[list[dict]] = []

    def _run_teacher_forced_prefill_hops(self, **kwargs):
        self.hops_calls.append(kwargs)
        return self.reference_activations_by_hop

    def _run_reference(self, messages: list[dict]) -> str:
        self.text_reference_calls.append(messages)
        return "two-hop pipeline output"


def test_bisection_blames_first_divergent_hop_not_last_hop():
    """Red: corrupt hop-0 only. The old `_final_text_node` bug blames
    max(shard_end) == hop-1 (wallet-hop1), which is innocent here."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]
    corrupted_hop0_claim_activations = [[9.0, 9.0], [9.0, 9.0]]  # diverges from reference_hop0

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=corrupted_hop0_claim_activations,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert contracts.registry.slashes[0]["wallet_address"] == "wallet-hop0"
    assert "hop 0" in contracts.registry.slashes[0]["reason"]


def test_bisection_blames_the_actual_faulty_hop_when_it_is_the_second_hop():
    """Integration test: multi-hop pipeline, fault injected at a known
    (non-first) hop — proves blame follows the real culprit, not a fixed index."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]
    corrupted_hop1_claim_activations = [[0.0, 0.0], [0.0, 0.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=corrupted_hop1_claim_activations,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    receipts = validator.validate_once()

    assert len(receipts) == 1
    assert contracts.registry.slashes[0]["wallet_address"] == "wallet-hop1"


def test_honest_two_hop_route_is_not_slashed():
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    assert validator.validate_once() == []
    assert contracts.registry.slashes == []


def test_expired_commitment_window_falls_back_to_text_only_audit():
    """ADR-0018 §3: nodes only retain boundary activations briefly. Once the
    on-demand TTL has passed, bisection can't be verified — the validator must
    fall back to the text-comparison path instead of erroring out."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16, commitment_ttl_seconds=1.0)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
        ts=0.0,  # far outside any TTL relative to real wall-clock time
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=1.0,
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    assert validator.validate_once() == []
    assert validator.hops_calls == []
    assert len(validator.text_reference_calls) == 1


def test_hop_commitments_are_not_requested_unless_the_event_is_audit_selected():
    """On-demand: the (expensive) per-hop commitment retrieval only happens
    for events the tracker RNG actually selects for audit — sample_rate is
    the selection gate, and a miss must not touch the reference node at all."""
    config = ToplocAuditConfig(topk=2, decode_batching_size=16)
    backend = FakeToploc()
    reference_hop0 = [[1.0, 2.0], [3.0, 4.0]]
    reference_hop1 = [[5.0, 6.0], [7.0, 8.0]]

    event = _hop_route_event(
        index=0,
        hop0_claim_activations=reference_hop0,
        hop1_claim_activations=reference_hop1,
        config=config,
        backend=backend,
    )
    contracts = FakeContracts([event])
    validator = HopReferenceValidator(
        contracts=contracts,
        reference_node_url="http://reference.invalid",
        sample_rate=0.0,  # never selects any event for audit
        random_seed=7,
        toploc_config=config,
        toploc_backend=backend,
        reference_activations_by_hop=[reference_hop0, reference_hop1],
    )

    assert validator.validate_once() == []
    assert validator.hops_calls == []
    assert validator.text_reference_calls == []
    assert validator.sampled_count == 0