"""Optimistic fraud validator for completed inference requests."""

import json
import math
import random
import threading
import time
import urllib.request
from typing import Any

from .audit import (
    ToplocAuditConfig,
    ToplocProofClaim,
    ToplocVerificationResult,
    verify_activation_proofs,
    verify_activation_proofs_detailed,
)
from .sampling import AdaptiveAuditSampler, AuditRateConfig
from .tripwire import detect_output_tripwire

__version__ = "0.1.0"


class ValidatorProcess:
    """Separate validator loop that samples completed requests and submits slashes."""

    def __init__(
        self,
        *,
        contracts: Any,
        reference_node_url: str,
        sample_rate: float = 0.05,
        tolerance: float = 1e-6,
        slash_amount: int = 100,
        strike_threshold: int = 3,
        random_seed: int | None = None,
        webhook_url: str | None = None,
        interval_seconds: float = 1.0,
        billing: Any | None = None,
        toploc_config: ToplocAuditConfig | None = None,
        toploc_backend: Any | None = None,
        audit_sampler: AdaptiveAuditSampler | None = None,
    ) -> None:
        if not 0.0 <= sample_rate <= 1.0:
            raise ValueError("sample_rate must be between 0 and 1")
        if tolerance < 0:
            raise ValueError("tolerance must be non-negative")
        if slash_amount <= 0:
            raise ValueError("slash_amount must be positive")
        if strike_threshold <= 0:
            raise ValueError("strike_threshold must be positive")
        if interval_seconds <= 0:
            raise ValueError("interval_seconds must be positive")

        self._contracts = contracts
        self._billing = billing
        self._reference_node_url = reference_node_url.rstrip("/")
        self._sample_rate = sample_rate
        self._tolerance = tolerance
        self._slash_amount = slash_amount
        self._strike_threshold = strike_threshold
        self._webhook_url = webhook_url
        self._interval_seconds = interval_seconds
        self._toploc_config = toploc_config or ToplocAuditConfig()
        self._toploc_backend = toploc_backend
        self._audit_sampler = audit_sampler
        self._random = random.Random(random_seed)
        self._last_event_index = -1
        self._running = False
        self._thread: threading.Thread | None = None
        self.sampled_count = 0

    def validate_once(self) -> list[Any]:
        """Run one validation cycle and return slash receipts submitted this cycle."""
        receipts: list[Any] = []
        events = self._contracts.validation.list_completed_inferences(
            after_index=self._last_event_index,
        )
        for event in events:
            self._last_event_index = max(self._last_event_index, event.index)
            if not self._should_sample(event):
                continue
            self.sampled_count += 1
            audit_result = self._validate_event(event)
            if audit_result.ok:
                self._record_clean_audit(event)
                continue
            receipts.extend(self._slash_node(
                audit_result.culprit_node,
                event.observed_output,
                audit_result.reference_output,
                reason=audit_result.reason,
            ))
        return receipts

    def _should_sample(self, event: Any) -> bool:
        """ADR-0018 §1/§6-7: flat sample_rate stays the default; when an
        AdaptiveAuditSampler is configured, the decision is reputation- and
        tenure-weighted and budget-balanced against the fleet-wide target
        instead of a uniform coin flip."""
        if self._audit_sampler is None:
            return self._random.random() < self._sample_rate

        tripwire = detect_output_tripwire(_event_value(event, "observed_output") or "")
        wallets = _route_wallets(event)
        if not wallets:
            return self._audit_sampler.should_audit(
                completed_job_count=0, reputation=1.0, tripwire=tripwire,
            )
        # A route is only as trustworthy as its least-trusted hop -- audit
        # against whichever wallet on the route looks riskiest.
        riskiest = min(
            (self._contracts.registry.get_wallet(wallet) for wallet in wallets),
            key=lambda wallet: wallet.reputation,
        )
        return self._audit_sampler.should_audit(
            completed_job_count=riskiest.completed_job_count,
            reputation=riskiest.reputation,
            tripwire=tripwire,
        )

    def start(self) -> None:
        if self._running:
            raise RuntimeError("ValidatorProcess is already running")
        self._running = True
        self._thread = threading.Thread(target=self._run_loop, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._running = False
        if self._thread is not None:
            self._thread.join(timeout=2)
        self._thread = None

    def _run_loop(self) -> None:
        while self._running:
            self.validate_once()
            time.sleep(self._interval_seconds)

    def _run_reference(self, messages: list[dict]) -> str:
        response = _post_json(
            f"{self._reference_node_url}/v1/infer",
            {"messages": messages},
        )
        text = response.get("text")
        if not isinstance(text, str):
            raise ValueError("reference node response did not contain text")
        return text

    def _validate_event(self, event: Any) -> "_AuditResult":
        event = self._event_with_on_demand_commitments(event)
        hop_commitments = _hop_commitments_from_event(event)
        if hop_commitments is not None and self._commitment_expired(event):
            # ADR-0018 §3: the on-demand retention window has passed — nodes
            # are no longer expected to hold the boundary activations needed
            # to verify this commitment, so fall back to the text-only path.
            hop_commitments = None

        if hop_commitments is None:
            reference_output = self._run_reference(event.messages)
            ok = _outputs_match(event.observed_output, reference_output, self._tolerance)
            return _AuditResult(
                ok=ok,
                reference_output=reference_output,
                reason="reference output diverged",
                # Text comparison has no per-hop signal; the last hop is the
                # best-effort guess (text-only fallback), never used when
                # hop-boundary commitments make real bisection possible.
                culprit_node=None if ok else _final_text_node(event.route_nodes),
            )

        if len(hop_commitments) == 1:
            # Single-commitment route (AH-006 whole-route format, or a
            # genuinely one-hop pipeline): reuse the original teacher-forced
            # call so existing single-hop reference integrations keep working.
            only = hop_commitments[0]
            reference_activations_by_hop = [self._run_teacher_forced_prefill(
                model=_event_value(event, "model"),
                messages=_event_value(event, "messages"),
                claimed_token_ids=only.token_ids,
                claim=only.claim,
            )]
        else:
            reference_activations_by_hop = self._run_teacher_forced_prefill_hops(
                model=_event_value(event, "model"),
                messages=_event_value(event, "messages"),
                hop_commitments=hop_commitments,
            )
        culprit_index = _first_divergent_hop(
            hop_commitments,
            reference_activations_by_hop,
            config=self._toploc_config,
            backend=self._toploc_backend,
        )
        ok = culprit_index is None
        return _AuditResult(
            ok=ok,
            reference_output=(
                "TOPLOC activation proof accepted"
                if ok
                else f"TOPLOC activation proof mismatch at hop {culprit_index}"
            ),
            reason="TOPLOC activation proof mismatch",
            culprit_node=None if ok else hop_commitments[culprit_index].node,
        )

    def _event_with_on_demand_commitments(self, event: Any) -> Any:
        """Fetch missing per-hop TOPLOC commitments only after audit sampling.

        Tracker validation events deliberately carry ordinary route metadata,
        not a pre-announced audit flag. When this validator samples an event, it
        asks each hop for its short-lived boundary commitment and splices the
        returned proof into a local event copy for the bisection verifier.
        """
        route_nodes = _event_value(event, "route_nodes") or []
        if not isinstance(route_nodes, list) or not route_nodes:
            return event
        updated_nodes: list[dict] = []
        changed = False
        for node in route_nodes:
            if not isinstance(node, dict):
                updated_nodes.append(node)
                continue
            updated = dict(node)
            if _mapping_value(updated, "toploc_proof") is None:
                commitment = self._fetch_hop_commitment(event, updated)
                if commitment is not None:
                    updated.update(commitment)
                    changed = True
            updated_nodes.append(updated)
        if not changed:
            return event
        if isinstance(event, dict):
            copied = dict(event)
        else:
            copied = dict(vars(event))
        copied["route_nodes"] = updated_nodes
        return copied

    def _fetch_hop_commitment(self, event: Any, node: dict) -> dict[str, Any] | None:
        endpoint = node.get("endpoint")
        if not isinstance(endpoint, str) or not endpoint:
            return None
        try:
            response = _post_json(
                f"{endpoint.rstrip('/')}/v1/audit/toploc/commitment",
                {
                    "session_id": _event_value(event, "session_id"),
                    "model": _event_value(event, "model"),
                    "messages": _event_value(event, "messages") or [],
                    "shard_start": node.get("shard_start"),
                    "shard_end": node.get("shard_end"),
                },
                timeout=2.0,
            )
        except (OSError, ValueError, json.JSONDecodeError):
            return None
        proof = response.get("toploc_proof") or response.get("activation_proof")
        token_ids = response.get("claimed_token_ids") or response.get("output_token_ids")
        if not isinstance(proof, dict):
            return None
        if not isinstance(token_ids, list) or not all(isinstance(token, int) for token in token_ids):
            return None
        return {"toploc_proof": proof, "claimed_token_ids": token_ids}

    def _commitment_expired(self, event: Any) -> bool:
        ts = _event_value(event, "ts")
        if ts is None:
            return False
        return (time.time() - float(ts)) > self._toploc_config.commitment_ttl_seconds

    def _run_teacher_forced_prefill(
        self,
        *,
        model: str,
        messages: list[dict],
        claimed_token_ids: list[int],
        claim: ToplocProofClaim,
    ) -> list[Any]:
        response = _post_json(
            f"{self._reference_node_url}/v1/audit/toploc",
            {
                "model": model,
                "messages": messages,
                "claimed_token_ids": claimed_token_ids,
                "dtype": claim.dtype,
                "quantization": claim.quantization,
                "decode_batching_size": claim.decode_batching_size,
                "topk": claim.topk,
                "skip_prefill": claim.skip_prefill,
            },
        )
        activations = response.get("activations")
        if not isinstance(activations, list):
            raise ValueError("reference node audit response did not contain activations")
        return activations

    def _run_teacher_forced_prefill_hops(
        self,
        *,
        model: str,
        messages: list[dict],
        hop_commitments: list["_HopCommitment"],
    ) -> list[list[Any]]:
        """Teacher-force the claimed tokens once and collect reference
        activations at every hop's boundary layer (ADR-0018 §4 / research §1.2:
        one referee forward pass, compared at each cut-point)."""
        reference_claim = hop_commitments[0].claim
        response = _post_json(
            f"{self._reference_node_url}/v1/audit/toploc",
            {
                "model": model,
                "messages": messages,
                "claimed_token_ids": hop_commitments[-1].token_ids,
                "hop_boundaries": [hop.shard_end for hop in hop_commitments],
                "dtype": reference_claim.dtype,
                "quantization": reference_claim.quantization,
                "decode_batching_size": reference_claim.decode_batching_size,
                "topk": reference_claim.topk,
                "skip_prefill": reference_claim.skip_prefill,
            },
        )
        activations_by_hop = response.get("activations_by_hop")
        if not isinstance(activations_by_hop, list) or len(activations_by_hop) != len(hop_commitments):
            raise ValueError("reference node audit response did not contain per-hop activations")
        return activations_by_hop

    def _record_clean_audit(self, event: Any) -> None:
        """ADR-0018 §6: reputation derives only from tracker-verified audit
        outcomes — a clean audit credits every node on the verified route."""
        for wallet_address in _route_wallets(event):
            if self._contracts.registry.get_wallet(wallet_address).banned:
                continue
            self._contracts.registry.record_audit_outcome(wallet_address, passed=True)

    def _slash_node(
        self,
        node: dict | None,
        observed_output: str,
        reference_output: str,
        *,
        reason: str = "reference output diverged",
    ) -> list[Any]:
        receipts: list[Any] = []
        wallet_address = node.get("wallet_address") if node else None
        if not wallet_address:
            return receipts
        if self._contracts.registry.get_wallet(wallet_address).banned:
            return receipts
        receipts.append(self._contracts.registry.submit_slash_proof(
            wallet_address=wallet_address,
            slash_amount=self._slash_amount,
            strike_threshold=self._strike_threshold,
            reason=(
                f"{reason} "
                f"(observed={observed_output!r}, reference={reference_output!r})"
            ),
            webhook_url=self._webhook_url,
        ))
        # ADR-0018 §6: reputation loss is separate from the strike/ban that
        # submit_slash_proof already recorded above — never double-strike.
        self._contracts.registry.record_audit_outcome(wallet_address, passed=False)
        # ADR-0015: the pending balance is the collateral — forfeit it in the
        # same validation cycle as the strike.
        if self._billing is not None:
            forfeit = self._billing.forfeit_pending(wallet_address, reason="fraud-divergence")
            print(
                f"[validator] forfeited pending balance of {wallet_address}: "
                f"{forfeit['amount']:.6f} USDT (fraud-divergence)",
                flush=True,
            )
        return receipts


def _route_wallets(event: Any) -> list[str]:
    """Unique wallet addresses across a route, in hop order."""
    route_nodes = _event_value(event, "route_nodes") or []
    seen: set[str] = set()
    wallets: list[str] = []
    for node in route_nodes:
        wallet_address = node.get("wallet_address") if isinstance(node, dict) else None
        if wallet_address and wallet_address not in seen:
            seen.add(wallet_address)
            wallets.append(wallet_address)
    return wallets


def _final_text_node(route_nodes: list[dict]) -> dict | None:
    """Text-only fallback blame: when the audit has no per-hop fingerprints
    to bisect (free-running text comparison only), guess the last hop.
    Never used once hop-boundary commitments make real bisection possible."""
    if not route_nodes:
        return None
    return max(route_nodes, key=lambda node: int(node.get("shard_end", 0)))


class _AuditResult:
    def __init__(
        self,
        *,
        ok: bool,
        reference_output: str,
        reason: str,
        culprit_node: dict | None = None,
    ) -> None:
        self.ok = ok
        self.reference_output = reference_output
        self.reason = reason
        self.culprit_node = culprit_node


class _HopCommitment:
    """One hop's on-demand TOPLOC commitment plus the route node it blames."""

    def __init__(self, node: dict | None, claim: ToplocProofClaim, token_ids: list[int]) -> None:
        self.node = node
        self.claim = claim
        self.token_ids = token_ids
        self.shard_end = int(node.get("shard_end", 0)) if node else None


def _hop_commitments_from_event(event: Any) -> list[_HopCommitment] | None:
    """Per-hop bisection commitments (ADR-0018 §3/§4): each route node reports
    its own output-boundary fingerprint. Falls back to the AH-006 whole-route
    commitment format (one fingerprint, no bisection) when hops don't carry
    individual commitments."""
    route_nodes = _event_value(event, "route_nodes") or []
    per_hop_nodes = [
        node for node in route_nodes
        if isinstance(node, dict) and _mapping_value(node, "toploc_proof") is not None
    ]
    if per_hop_nodes and len(per_hop_nodes) == len(route_nodes):
        ordered = sorted(route_nodes, key=lambda node: int(node.get("shard_start", 0)))
        default_token_ids = _event_value(event, "claimed_token_ids")
        commitments = []
        for node in ordered:
            token_ids = node.get("claimed_token_ids", default_token_ids)
            if not isinstance(token_ids, list) or not all(isinstance(t, int) for t in token_ids):
                raise ValueError("TOPLOC hop commitments must include claimed_token_ids")
            commitments.append(_HopCommitment(
                node,
                ToplocProofClaim.from_mapping(node["toploc_proof"]),
                token_ids,
            ))
        return commitments

    whole_route = _toploc_audit_from_event(event)
    if whole_route is None:
        return None
    token_ids, claim = whole_route
    return [_HopCommitment(route_nodes[-1] if route_nodes else None, claim, token_ids)]


def _first_divergent_hop(
    hop_commitments: list[_HopCommitment],
    reference_activations_by_hop: list[Any],
    *,
    config: ToplocAuditConfig,
    backend: Any | None,
) -> int | None:
    """First hop whose committed output fingerprint diverges from the
    referee's independently-computed reference activations at that same
    cut-point (research §1.2: no interactive game needed at hop granularity —
    the referee checks every cut-point in one replay)."""
    for index, commitment in enumerate(hop_commitments):
        ok = verify_activation_proofs(
            reference_activations_by_hop[index],
            commitment.claim,
            config=config,
            backend=backend,
        )
        if not ok:
            return index
    return None


def _toploc_audit_from_event(event: Any) -> tuple[list[int], ToplocProofClaim] | None:
    audit = _event_mapping(event, "audit")
    claim_data = (
        _event_mapping(event, "toploc_proof")
        or _event_mapping(event, "activation_proof")
        or _mapping_value(audit, "toploc_proof")
        or _mapping_value(audit, "activation_proof")
        or _mapping_value(audit, "toploc")
    )
    if claim_data is None:
        return None
    token_ids = (
        _event_value(event, "claimed_token_ids")
        or _event_value(event, "output_token_ids")
        or _mapping_value(audit, "claimed_token_ids")
        or _mapping_value(audit, "output_token_ids")
    )
    if not isinstance(token_ids, list) or not all(isinstance(token, int) for token in token_ids):
        raise ValueError("TOPLOC audit events must include claimed_token_ids")
    return token_ids, ToplocProofClaim.from_mapping(claim_data)


def _event_value(event: Any, name: str) -> Any:
    if hasattr(event, name):
        return getattr(event, name)
    if isinstance(event, dict):
        return event.get(name)
    return None


def _event_mapping(event: Any, name: str) -> dict[str, Any] | None:
    value = _event_value(event, name)
    return value if isinstance(value, dict) else None


def _mapping_value(mapping: dict[str, Any] | None, name: str) -> Any:
    if mapping is None:
        return None
    return mapping.get(name)


def _outputs_match(observed: str, reference: str, tolerance: float) -> bool:
    observed_float = _parse_float(observed)
    reference_float = _parse_float(reference)
    if observed_float is not None and reference_float is not None:
        return math.isclose(observed_float, reference_float, rel_tol=tolerance, abs_tol=tolerance)
    return observed == reference


def _parse_float(value: str) -> float | None:
    try:
        return float(value)
    except ValueError:
        return None


def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
    data = json.dumps(payload).encode()
    req = urllib.request.Request(
        url,
        data=data,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=timeout) as response:
        return json.loads(response.read())


__all__ = [
    "ToplocAuditConfig",
    "ToplocProofClaim",
    "ToplocVerificationResult",
    "verify_activation_proofs_detailed",
    "ValidatorProcess",
    "AdaptiveAuditSampler",
    "AuditRateConfig",
    "detect_output_tripwire",
]
