"""Tests for FailureRateAdaptiveSampler."""
from unittest.mock import MagicMock

import pytest

from llm_guard.adaptive_sampler import FailureRateAdaptiveSampler


def _make_result(needs_alert: bool = False, tier: str = "HIGH"):
    """Mock ChainTrustResult."""
    r = MagicMock()
    r.needs_alert = needs_alert
    r.confidence_tier = tier
    return r


def _ok():
    return _make_result(needs_alert=False, tier="HIGH")


def _fail():
    return _make_result(needs_alert=True, tier="LOW")


class TestColdStart:
    def test_always_scores_during_cold_start(self):
        s = FailureRateAdaptiveSampler(window=5, seed=42)
        for _ in range(5):
            assert s.should_score() is True

    def test_mode_is_cold_start_initially(self):
        s = FailureRateAdaptiveSampler(window=5, seed=42)
        assert s.status()["mode"] == "cold_start"

    def test_transitions_to_healthy_after_window_low_failure_rate(self):
        s = FailureRateAdaptiveSampler(window=5, anomaly_threshold=0.25, seed=42)
        for _ in range(5):
            s.should_score()
            s.observe(_ok())  # 0% failure rate
        assert s.status()["mode"] == "healthy"

    def test_transitions_to_burst_after_window_high_failure_rate(self):
        s = FailureRateAdaptiveSampler(window=5, anomaly_threshold=0.25, seed=42)
        for _ in range(5):
            s.should_score()
            s.observe(_fail())  # 100% failure rate
        assert s.status()["mode"] == "burst"


class TestHealthyMode:
    def test_sample_rate_roughly_15_percent(self):
        s = FailureRateAdaptiveSampler(
            window=10, healthy_sample_rate=0.15, anomaly_threshold=0.50, seed=42
        )
        # Warm up through cold start
        for _ in range(10):
            s.should_score()
            s.observe(_ok())
        assert s.status()["mode"] == "healthy"
        # Now sample N=1000 calls
        n = 1000
        scored = sum(1 for _ in range(n) if s.should_score())
        # Expect ~15%, allow ±5pp
        assert 100 <= scored <= 200, f"Expected 100-200 scored, got {scored}"

    def test_failure_triggers_burst(self):
        s = FailureRateAdaptiveSampler(
            window=4, healthy_sample_rate=1.0, anomaly_threshold=0.25, seed=42
        )
        # Warm up: 4 ok observations → healthy
        for _ in range(4):
            s.should_score()
            s.observe(_ok())
        assert s.status()["mode"] == "healthy"
        # Push failure rate above threshold: 2 failures in window of 4 = 50%
        s.observe(_fail())
        s.observe(_fail())
        assert s.status()["mode"] == "burst"


class TestBurstMode:
    def test_burst_always_scores(self):
        s = FailureRateAdaptiveSampler(
            window=3, healthy_sample_rate=0.0, anomaly_threshold=0.25,
            burst_window=5, seed=42
        )
        # Enter burst: 3 failures
        for _ in range(3):
            s.should_score()
            s.observe(_fail())
        assert s.status()["mode"] == "burst"
        for _ in range(5):
            assert s.should_score() is True

    def test_burst_counter_decrements_on_should_score(self):
        s = FailureRateAdaptiveSampler(
            window=3, anomaly_threshold=0.25, burst_window=5, seed=42
        )
        for _ in range(3):
            s.should_score()
            s.observe(_fail())
        assert s.status()["calls_in_burst"] == 5
        s.should_score()
        assert s.status()["calls_in_burst"] == 4

    def test_burst_exits_to_healthy_after_burst_window(self):
        s = FailureRateAdaptiveSampler(
            window=4, anomaly_threshold=0.25, burst_window=4, seed=42
        )
        # Enter burst
        for _ in range(4):
            s.should_score()
            s.observe(_fail())
        assert s.status()["mode"] == "burst"
        # During burst, add good observations to lower failure rate
        for _ in range(4):
            s.should_score()
            s.observe(_ok())  # failure rate drops as window fills with ok
        # After 4 should_score() calls in burst, counter = 0 → re-evaluate
        # Window should now have recent ok observations → healthy
        assert s.status()["mode"] == "healthy"

    def test_burst_persists_if_failure_rate_still_high(self):
        s = FailureRateAdaptiveSampler(
            window=4, anomaly_threshold=0.25, burst_window=3, seed=42
        )
        for _ in range(4):
            s.should_score()
            s.observe(_fail())
        assert s.status()["mode"] == "burst"
        # All burst calls still produce failures → stays in burst
        for _ in range(3):
            s.should_score()
            s.observe(_fail())
        assert s.status()["mode"] == "burst"
        assert s.status()["calls_in_burst"] == 3  # reset


class TestObserveEdgeCases:
    def test_observe_without_prior_should_score_does_not_raise(self):
        s = FailureRateAdaptiveSampler(window=5, seed=42)
        s.observe(_ok())  # no prior should_score()
        s.observe(_fail())

    def test_failure_defined_as_needs_alert_or_low_tier(self):
        s = FailureRateAdaptiveSampler(window=4, anomaly_threshold=0.25, seed=42)
        for _ in range(4):
            s.should_score()
        # needs_alert=False, tier="LOW" → still a failure
        s.observe(_make_result(needs_alert=False, tier="LOW"))
        assert s.status()["failure_rate"] == pytest.approx(0.25)


class TestMultiPipeline:
    def test_pipelines_are_independent(self):
        s = FailureRateAdaptiveSampler(window=3, anomaly_threshold=0.25, seed=42)
        # Push pipeline A into burst
        for _ in range(3):
            s.should_score(pipeline_id="A")
            s.observe(_fail(), pipeline_id="A")
        # Pipeline B should still be in cold_start
        assert s.status("A")["mode"] == "burst"
        assert s.status("B")["mode"] == "cold_start"

    def test_reset_clears_pipeline_state(self):
        s = FailureRateAdaptiveSampler(window=3, anomaly_threshold=0.25, seed=42)
        for _ in range(3):
            s.should_score(pipeline_id="X")
            s.observe(_fail(), pipeline_id="X")
        assert s.status("X")["mode"] == "burst"
        s.reset("X")
        assert s.status("X")["mode"] == "cold_start"
        assert s.status("X")["n_observations"] == 0


class TestStatus:
    def test_status_has_all_required_keys(self):
        s = FailureRateAdaptiveSampler()
        keys = s.status().keys()
        for k in ("pipeline_id", "failure_rate", "mode", "calls_in_burst",
                  "n_observations", "sample_rate"):
            assert k in keys

    def test_sample_rate_is_1_in_cold_start(self):
        s = FailureRateAdaptiveSampler(healthy_sample_rate=0.15)
        assert s.status()["sample_rate"] == 1.0

    def test_sample_rate_is_healthy_rate_in_healthy_mode(self):
        s = FailureRateAdaptiveSampler(
            window=3, healthy_sample_rate=0.15, anomaly_threshold=0.50, seed=42
        )
        for _ in range(3):
            s.should_score()
            s.observe(_ok())
        assert s.status()["mode"] == "healthy"
        assert s.status()["sample_rate"] == pytest.approx(0.15)
