//! Cost controls for LLM-jury runs.
//!
//! Three knobs cap what an eval run spends:
//!
//! - [`CostBudget`] tracks cumulative dollars and refuses calls that
//!   would tip the run over the cap.
//! - [`should_sample`] decides whether a juror runs at all for a given
//!   test name, using a deterministic SHA-256 hash so reruns reproduce.
//! - [`ConsensusMethod::would_terminate_early`] tells the jury dispatch
//!   loop when the completed jurors already pin the verdict, so it can
//!   skip the remaining juror calls. The shared loop in
//!   `crate::eval::jury::engine::dispatch_votes` consults it after each
//!   vote when the caller opts in.
//!
//! All three are pure data + math; they are independent of the provider
//! impls so they can be unit-tested without a network.

use sha2::{Digest, Sha256};

use super::provider::Cost;
use crate::report::RunCost;

/// Cumulative cost ledger for an eval run.
///
/// The runner records every completed call via [`CostBudget::record_call`]
/// and consults [`CostBudget::would_exceed`] before dispatching the next.
/// When the cap is `None` the budget is unbounded and `would_exceed`
/// always returns `false`.
#[derive(Debug, Clone, Default)]
pub struct CostBudget {
    /// Maximum cumulative dollars allowed across the whole run. `None`
    /// means unbounded.
    cap_usd: Option<f64>,
    /// Running total in dollars.
    spent_usd: f64,
    /// Number of completed juror calls. Used by reporters to compute
    /// `cost_per_juror_avg`.
    juror_calls: u64,
}

impl CostBudget {
    /// Build an unbounded budget. Records cost but never refuses.
    pub fn unbounded() -> Self {
        Self {
            cap_usd: None,
            spent_usd: 0.0,
            juror_calls: 0,
        }
    }

    /// Build a budget capped at `cap_usd` total spend. Negative caps are
    /// clamped to zero so a config mistake cannot turn the cap into an
    /// infinite-budget surprise.
    pub fn with_cap(cap_usd: f64) -> Self {
        Self {
            cap_usd: Some(cap_usd.max(0.0)),
            spent_usd: 0.0,
            juror_calls: 0,
        }
    }

    /// Whether running `next` would tip the cumulative spend over the
    /// cap. Returns `false` when no cap is set.
    pub fn would_exceed(&self, next: &Cost) -> bool {
        match self.cap_usd {
            Some(cap) => self.spent_usd + next.total() > cap + 1e-9,
            None => false,
        }
    }

    /// Record a completed juror call's cost. The counter advances even
    /// when the cost is zero (free providers, cached responses) so the
    /// per-juror average reflects real call volume.
    pub fn record_call(&mut self, cost: &Cost) {
        self.spent_usd += cost.total();
        self.juror_calls += 1;
    }

    /// Total dollars spent so far.
    pub fn spent(&self) -> f64 {
        self.spent_usd
    }

    /// Number of completed juror calls.
    pub fn juror_calls(&self) -> u64 {
        self.juror_calls
    }

    /// Configured cap, if any.
    pub fn cap(&self) -> Option<f64> {
        self.cap_usd
    }

    /// Average dollars per juror call. Returns `None` when no juror has
    /// been recorded yet.
    pub fn cost_per_juror_avg(&self) -> Option<f64> {
        if self.juror_calls == 0 {
            None
        } else {
            Some(self.spent_usd / self.juror_calls as f64)
        }
    }

    /// Roll up this budget into a [`RunCost`] for the reporter footer.
    ///
    /// `tests_with_cost` is the count of tests that ran at least one
    /// juror call; the runner tracks it because the budget only knows
    /// about juror calls, not parent tests. Returns `None` when no juror
    /// has been recorded (the reporter omits the cost line entirely in
    /// that case).
    pub fn to_run_cost(&self, tests_with_cost: u64) -> Option<RunCost> {
        if self.juror_calls == 0 {
            return None;
        }
        let cost_per_juror_avg = self.spent_usd / self.juror_calls as f64;
        let cost_per_test_avg = if tests_with_cost == 0 {
            self.spent_usd
        } else {
            self.spent_usd / tests_with_cost as f64
        };
        Some(RunCost {
            total_cost_usd: self.spent_usd,
            cost_per_juror_avg,
            cost_per_test_avg,
            juror_calls: self.juror_calls,
            tests_with_cost,
            // This juror-centric builder tracks calls, not parent-test pass/fail,
            // so the completed-task figure is filled by the eval/run report path
            // that knows the verdicts (WOR-1393).
            cost_per_completed_task: None,
        })
    }
}

/// Decide whether a test name should be sampled for jury evaluation.
///
/// Sampling is deterministic across reruns: the same test name plus the
/// same rate always returns the same answer. `rate <= 0.0` skips every
/// test; `rate >= 1.0` samples every test.
///
/// The mapping hashes the test name with SHA-256, takes the leading 8
/// bytes as a `u64`, and compares against `rate * u64::MAX`. The
/// distribution is close to uniform for any non-degenerate test-name
/// corpus.
pub fn should_sample(test_name: &str, rate: f64) -> bool {
    if rate <= 0.0 {
        return false;
    }
    if rate >= 1.0 {
        return true;
    }
    let mut hasher = Sha256::new();
    hasher.update(test_name.as_bytes());
    let digest = hasher.finalize();
    let mut head = [0u8; 8];
    head.copy_from_slice(&digest[..8]);
    let bucket = u64::from_be_bytes(head);
    let threshold = (rate * (u64::MAX as f64)) as u64;
    bucket < threshold
}

/// Consensus method for the jury (ADR-0008).
///
/// This is the single definition shared by the cost-control early-termination
/// logic here and the verdict-aggregation engine in
/// [`crate::eval::jury::engine`], which re-exports it. Unifying the two former
/// copies means a method added in one place is automatically visible
/// to the other.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConsensusMethod {
    /// More passes than fails. Early-terminate when one side has a
    /// guaranteed minority of the total jurors.
    Majority,
    /// Every juror must pass. Early-terminate on the first fail.
    Unanimous,
    /// At least `k` jurors must pass.
    KOfN {
        /// Threshold count.
        k: u32,
    },
}

/// Running tally fed to [`ConsensusMethod::would_terminate_early`].
///
/// `total` is the size of the jury (including jurors that have not yet
/// returned). `passes` and `fails` count completed jurors.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ConsensusTally {
    /// Total juror count, including jurors still in flight.
    pub total: u32,
    /// Jurors that have already voted pass.
    pub passes: u32,
    /// Jurors that have already voted fail.
    pub fails: u32,
}

impl ConsensusTally {
    /// How many jurors have yet to return.
    pub fn unresolved(&self) -> u32 {
        self.total.saturating_sub(self.passes + self.fails)
    }
}

impl ConsensusMethod {
    /// Whether the jury can short-circuit. Returns the early outcome
    /// (`Some(true)` for pass, `Some(false)` for fail) when the
    /// remaining juror calls cannot change the verdict.
    ///
    /// Aligns with the early-termination table in ADR-0008.
    pub fn would_terminate_early(&self, tally: ConsensusTally) -> Option<bool> {
        match self {
            ConsensusMethod::Majority => terminate_majority(tally),
            ConsensusMethod::Unanimous => terminate_unanimous(tally),
            ConsensusMethod::KOfN { k } => terminate_k_of_n(tally, *k),
        }
    }

    /// Stable label the jury reporter surfaces as the consensus method.
    ///
    /// Lives on the canonical type so the engine and any other caller
    /// render the same string instead of keeping a parallel mapping.
    pub fn label(self) -> &'static str {
        match self {
            ConsensusMethod::Majority => "majority",
            ConsensusMethod::Unanimous => "unanimous",
            ConsensusMethod::KOfN { .. } => "k_of_n",
        }
    }
}

fn terminate_majority(tally: ConsensusTally) -> Option<bool> {
    let half = tally.total / 2;
    if tally.passes > half {
        return Some(true);
    }
    // ceil(n / 2) fails guarantees passes cannot exceed half.
    let ceil_half = tally.total.div_ceil(2);
    if tally.fails >= ceil_half {
        return Some(false);
    }
    None
}

fn terminate_unanimous(tally: ConsensusTally) -> Option<bool> {
    if tally.fails > 0 {
        return Some(false);
    }
    if tally.passes == tally.total {
        return Some(true);
    }
    None
}

fn terminate_k_of_n(tally: ConsensusTally, k: u32) -> Option<bool> {
    if tally.passes >= k {
        return Some(true);
    }
    if tally.unresolved() + tally.passes < k {
        return Some(false);
    }
    None
}

#[cfg(test)]
mod tests {
    use super::*;

    fn dollars(d: f64) -> Cost {
        Cost {
            input_dollars: d,
            output_dollars: 0.0,
            cache_read_dollars: None,
            cache_write_dollars: None,
        }
    }

    #[test]
    fn unbounded_budget_never_exceeds() {
        let b = CostBudget::unbounded();
        assert!(!b.would_exceed(&dollars(1_000_000.0)));
    }

    #[test]
    fn capped_budget_blocks_calls_that_would_overshoot() {
        let mut b = CostBudget::with_cap(1.00);
        b.record_call(&dollars(0.40));
        assert!(!b.would_exceed(&dollars(0.50)));
        assert!(b.would_exceed(&dollars(0.65)));
    }

    #[test]
    fn record_call_advances_spent_and_count() {
        let mut b = CostBudget::with_cap(5.00);
        b.record_call(&dollars(0.10));
        b.record_call(&dollars(0.20));
        assert!((b.spent() - 0.30).abs() < 1e-9);
        assert_eq!(b.juror_calls(), 2);
        assert!((b.cost_per_juror_avg().unwrap() - 0.15).abs() < 1e-9);
    }

    #[test]
    fn negative_caps_clamp_to_zero() {
        let b = CostBudget::with_cap(-1.0);
        assert_eq!(b.cap(), Some(0.0));
    }

    #[test]
    fn sampling_at_one_includes_every_test() {
        assert!(should_sample("foo", 1.0));
        assert!(should_sample("anything goes", 1.0));
    }

    #[test]
    fn sampling_at_zero_skips_every_test() {
        assert!(!should_sample("foo", 0.0));
        assert!(!should_sample("anything goes", 0.0));
    }

    #[test]
    fn sampling_is_deterministic_across_calls() {
        for name in ["test-a", "test-b", "another"].iter() {
            let a = should_sample(name, 0.5);
            let b = should_sample(name, 0.5);
            assert_eq!(a, b, "sampling for {name} flipped between calls");
        }
    }

    #[test]
    fn sampling_at_half_lands_near_50_percent_over_many_names() {
        let n = 1_000;
        let hits: usize = (0..n)
            .filter(|i| should_sample(&format!("test-{i}"), 0.5))
            .count();
        // Allow plenty of slack; we just want to confirm we are not biased
        // to one side. 35-65% is a wide enough band that flakes are rare.
        assert!(
            hits > 350 && hits < 650,
            "expected ~500 of {n} samples, got {hits}"
        );
    }

    #[test]
    fn majority_terminates_when_passes_clear_half() {
        let tally = ConsensusTally {
            total: 5,
            passes: 3,
            fails: 0,
        };
        assert_eq!(
            ConsensusMethod::Majority.would_terminate_early(tally),
            Some(true)
        );
    }

    #[test]
    fn majority_terminates_on_guaranteed_minority() {
        let tally = ConsensusTally {
            total: 5,
            passes: 0,
            fails: 3,
        };
        assert_eq!(
            ConsensusMethod::Majority.would_terminate_early(tally),
            Some(false)
        );
    }

    #[test]
    fn majority_keeps_running_when_outcome_is_open() {
        let tally = ConsensusTally {
            total: 5,
            passes: 1,
            fails: 1,
        };
        assert_eq!(ConsensusMethod::Majority.would_terminate_early(tally), None);
    }

    #[test]
    fn unanimous_terminates_on_first_fail() {
        let tally = ConsensusTally {
            total: 4,
            passes: 2,
            fails: 1,
        };
        assert_eq!(
            ConsensusMethod::Unanimous.would_terminate_early(tally),
            Some(false)
        );
    }

    #[test]
    fn unanimous_terminates_when_every_juror_passed() {
        let tally = ConsensusTally {
            total: 3,
            passes: 3,
            fails: 0,
        };
        assert_eq!(
            ConsensusMethod::Unanimous.would_terminate_early(tally),
            Some(true)
        );
    }

    #[test]
    fn k_of_n_short_circuits_on_threshold_reached() {
        let tally = ConsensusTally {
            total: 5,
            passes: 2,
            fails: 0,
        };
        assert_eq!(
            ConsensusMethod::KOfN { k: 2 }.would_terminate_early(tally),
            Some(true)
        );
    }

    #[test]
    fn to_run_cost_returns_none_when_no_juror_recorded() {
        let b = CostBudget::with_cap(5.0);
        assert!(b.to_run_cost(0).is_none());
    }

    #[test]
    fn to_run_cost_rolls_totals_and_averages() {
        let mut b = CostBudget::unbounded();
        b.record_call(&dollars(0.10));
        b.record_call(&dollars(0.30));
        b.record_call(&dollars(0.20));
        let cost = b.to_run_cost(2).expect("expected RunCost");
        assert!((cost.total_cost_usd - 0.60).abs() < 1e-9);
        assert_eq!(cost.juror_calls, 3);
        assert_eq!(cost.tests_with_cost, 2);
        assert!((cost.cost_per_juror_avg - 0.20).abs() < 1e-9);
        assert!((cost.cost_per_test_avg - 0.30).abs() < 1e-9);
    }

    #[test]
    fn k_of_n_short_circuits_when_unreachable() {
        let tally = ConsensusTally {
            total: 5,
            passes: 1,
            fails: 3,
        };
        assert_eq!(
            ConsensusMethod::KOfN { k: 3 }.would_terminate_early(tally),
            Some(false)
        );
    }
}
