"""Tests for ABRollout."""

from __future__ import annotations

import concurrent.futures

import pytest

from llmhosts.training.ab_rollout import ABRollout


class TestABRollout:
    def test_should_use_new_weights_deterministic(self) -> None:
        """should_use_new_weights must return the same result for the same hash."""
        rollout = ABRollout(split_pct=0.10)
        request_id = "test-request-uuid-12345"
        results = [rollout.should_use_new_weights(request_id) for _ in range(50)]
        assert len(set(results)) == 1, "Result should be identical for the same request_hash"

    def test_split_is_approximately_correct(self) -> None:
        """10% split should route ~10% of traffic to treatment (within 2pp tolerance)."""
        rollout = ABRollout(split_pct=0.10)
        n = 10_000
        treatment_count = sum(1 for i in range(n) if rollout.should_use_new_weights(f"req-{i:06d}"))
        treatment_pct = treatment_count / n
        # Should be within 2% of target
        assert abs(treatment_pct - 0.10) < 0.02, f"Expected ~10% treatment traffic, got {treatment_pct:.2%}"

    def test_zero_split_always_returns_false(self) -> None:
        """0% split should never route to treatment."""
        rollout = ABRollout(split_pct=0.0)
        for i in range(1000):
            assert rollout.should_use_new_weights(f"req-{i}") is False

    def test_full_split_always_returns_true(self) -> None:
        """100% split should always route to treatment."""
        rollout = ABRollout(split_pct=1.0)
        for i in range(1000):
            assert rollout.should_use_new_weights(f"req-{i}") is True

    def test_invalid_split_raises(self) -> None:
        """split_pct outside [0, 1] should raise ValueError."""
        with pytest.raises(ValueError, match="split_pct"):
            ABRollout(split_pct=1.5)
        with pytest.raises(ValueError, match="split_pct"):
            ABRollout(split_pct=-0.1)

    def test_record_result_accumulates_control(self) -> None:
        """record_result should count control successes correctly."""
        rollout = ABRollout(split_pct=0.10)
        for i in range(10):
            rollout.record_result("control", latency_ms=float(i + 1), success=(i % 2 == 0))
        metrics = rollout.get_metrics()
        assert metrics["control"]["requests"] == 10
        # 5 successes out of 10
        assert metrics["control"]["success_rate"] == pytest.approx(0.5)

    def test_record_result_accumulates_treatment(self) -> None:
        """record_result should count treatment successes correctly."""
        rollout = ABRollout(split_pct=0.10)
        for _ in range(5):
            rollout.record_result("treatment", latency_ms=15.0, success=True)
        metrics = rollout.get_metrics()
        assert metrics["treatment"]["requests"] == 5
        assert metrics["treatment"]["success_rate"] == pytest.approx(1.0)

    def test_get_metrics_initial_state(self) -> None:
        """get_metrics() on a fresh rollout should return zero counts."""
        rollout = ABRollout(split_pct=0.10)
        metrics = rollout.get_metrics()
        assert metrics["control"]["requests"] == 0
        assert metrics["treatment"]["requests"] == 0
        assert metrics["control"]["success_rate"] == 0.0
        assert metrics["treatment"]["success_rate"] == 0.0

    def test_get_metrics_structure(self) -> None:
        """get_metrics must return both 'control' and 'treatment' keys."""
        rollout = ABRollout(split_pct=0.10)
        metrics = rollout.get_metrics()
        assert "control" in metrics
        assert "treatment" in metrics
        for variant in ("control", "treatment"):
            assert "requests" in metrics[variant]
            assert "success_rate" in metrics[variant]
            assert "p50_latency_ms" in metrics[variant]

    def test_p50_latency_computed_correctly(self) -> None:
        """p50 latency should be the median of recorded latencies."""
        rollout = ABRollout(split_pct=0.10)
        latencies = [10.0, 20.0, 30.0, 40.0, 50.0]
        for lat in latencies:
            rollout.record_result("control", latency_ms=lat, success=True)
        metrics = rollout.get_metrics()
        # Median of [10, 20, 30, 40, 50] = 30
        assert metrics["control"]["p50_latency_ms"] == pytest.approx(30.0)

    def test_unknown_variant_is_ignored(self) -> None:
        """record_result with unknown variant should not crash or accumulate."""
        rollout = ABRollout(split_pct=0.10)
        rollout.record_result("unknown_variant", latency_ms=5.0, success=True)
        metrics = rollout.get_metrics()
        assert metrics["control"]["requests"] == 0
        assert metrics["treatment"]["requests"] == 0

    def test_thread_safety(self) -> None:
        """100 concurrent calls to record_result should not corrupt state."""
        rollout = ABRollout(split_pct=0.10)
        n_threads = 100

        def worker(i: int) -> None:
            variant = "control" if i % 2 == 0 else "treatment"
            rollout.record_result(variant, latency_ms=float(i), success=True)

        with concurrent.futures.ThreadPoolExecutor(max_workers=20) as pool:
            futures = [pool.submit(worker, i) for i in range(n_threads)]
            for fut in concurrent.futures.as_completed(futures):
                fut.result()  # re-raise any exceptions

        metrics = rollout.get_metrics()
        total = metrics["control"]["requests"] + metrics["treatment"]["requests"]
        assert total == n_threads, f"Expected {n_threads} total requests, got {total}"

    def test_hash_uses_md5_not_identity(self) -> None:
        """The split should not be based on simple integer modulo of the string."""
        # Verify the hash is applied (different string lengths can map to same bucket)
        rollout = ABRollout(split_pct=0.50)
        # Two requests that differ only in length — their routing may differ
        r1 = rollout.should_use_new_weights("a")
        r2 = rollout.should_use_new_weights("aa")
        # Just check no exception is raised and results are boolean
        assert isinstance(r1, bool)
        assert isinstance(r2, bool)

    def test_split_pct_property(self) -> None:
        """split_pct property should return the initialised value."""
        rollout = ABRollout(split_pct=0.25)
        assert rollout.split_pct == pytest.approx(0.25)
