"""Reputation-weighted, budget-balanced audit sampling (ADR-0018 §1, §6-7).

The flat ``sample_rate`` on ``ValidatorProcess`` treats every wallet the same.
This module scores each wallet's audit probability from its tenure
(``completed_job_count``) and reputation -- newcomers and low-reputation
wallets get sampled far more than veterans in good standing, who float down
to a floor -- while a running budget balance keeps the *realized* fleet-wide
average anchored to a configured target even as the wallet population mix
drifts (research §6, §8 layers 2-4).
"""

from __future__ import annotations

import random
from dataclasses import dataclass


@dataclass(frozen=True)
class AuditRateConfig:
    """Tunable audit-rate policy. Defaults match ADR-0018 §1/§6."""

    target_rate: float = 0.05
    newcomer_rate: float = 0.25
    veteran_floor: float = 0.02
    probation_jobs: int = 50
    veteran_jobs: int = 500
    tripwire_multiplier: float = 3.0
    # Bounds the correction the budget balancer can apply in one step so a
    # short burst of skewed traffic can't swing everyone's rate wildly.
    max_budget_scale: float = 20.0

    def __post_init__(self) -> None:
        if not 0.0 <= self.target_rate <= 1.0:
            raise ValueError("target_rate must be between 0 and 1")
        if not 0.0 <= self.veteran_floor <= self.newcomer_rate <= 1.0:
            raise ValueError("veteran_floor must be <= newcomer_rate, both within [0, 1]")
        if self.probation_jobs < 0 or self.veteran_jobs <= self.probation_jobs:
            raise ValueError("veteran_jobs must be greater than probation_jobs")
        if self.tripwire_multiplier < 1.0:
            raise ValueError("tripwire_multiplier must be >= 1.0")


class AdaptiveAuditSampler:
    """Per-wallet audit probability + fleet-wide budget balancer.

    ``wallet_base_rate`` scores tenure and reputation into a rate between
    ``veteran_floor`` and ``newcomer_rate``. ``should_audit`` scales that rate
    by a running budget-balance factor (target_rate / historical mean base
    rate) so the fleet-wide realized audit rate tracks ``target_rate`` even
    when the wallet mix is skewed toward veterans or newcomers, then applies
    the tripwire multiplier to *that single decision only* -- a tripwire flag
    never touches the shared budget-balance history, so it can't punish
    other wallets' audit rate.
    """

    def __init__(self, config: AuditRateConfig | None = None, *, random_seed: int | None = None) -> None:
        self.config = config or AuditRateConfig()
        self._random = random.Random(random_seed)
        self._base_rate_total = 0.0
        self._decisions = 0

    def wallet_base_rate(self, *, completed_job_count: int, reputation: float) -> float:
        cfg = self.config
        if completed_job_count <= cfg.probation_jobs:
            tenure_rate = cfg.newcomer_rate
        elif completed_job_count >= cfg.veteran_jobs:
            tenure_rate = cfg.veteran_floor
        else:
            span = cfg.veteran_jobs - cfg.probation_jobs
            frac = (completed_job_count - cfg.probation_jobs) / span
            tenure_rate = cfg.newcomer_rate + (cfg.veteran_floor - cfg.newcomer_rate) * frac

        reputation_gap = max(0.0, min(1.0, 1.0 - reputation))
        reputation_rate = cfg.veteran_floor + reputation_gap * (cfg.newcomer_rate - cfg.veteran_floor)

        return max(tenure_rate, reputation_rate, cfg.veteran_floor)

    def sample_rate_for(self, *, completed_job_count: int, reputation: float, tripwire: bool = False) -> float:
        """Preview the effective audit probability without recording a decision."""
        base_rate = self.wallet_base_rate(completed_job_count=completed_job_count, reputation=reputation)
        return self._effective_rate(base_rate, tripwire=tripwire)

    def should_audit(self, *, completed_job_count: int, reputation: float, tripwire: bool = False) -> bool:
        base_rate = self.wallet_base_rate(completed_job_count=completed_job_count, reputation=reputation)
        effective_rate = self._effective_rate(base_rate, tripwire=tripwire)
        self._base_rate_total += base_rate
        self._decisions += 1
        return self._random.random() < effective_rate

    def _effective_rate(self, base_rate: float, *, tripwire: bool) -> float:
        rate = base_rate * self._budget_scale()
        if tripwire:
            rate *= self.config.tripwire_multiplier
        return max(self.config.veteran_floor, min(1.0, rate))

    def _budget_scale(self) -> float:
        if self._decisions == 0:
            return 1.0
        mean_base_rate = self._base_rate_total / self._decisions
        if mean_base_rate <= 0:
            return 1.0
        return min(self.config.max_budget_scale, self.config.target_rate / mean_base_rate)


__all__ = ["AuditRateConfig", "AdaptiveAuditSampler"]