"""Tests for Heuristic Lifetime Predictors.

This module tests the heuristic-based predictors:
- LengthBasedPredictor
- HistoricalPredictor
- HybridHeuristicPredictor
"""

from __future__ import annotations

import pytest

from sagellm_control.predictor.heuristic import (
    HistoricalPredictor,
    HybridHeuristicPredictor,
    LengthBasedPredictor,
)
from sagellm_control.predictor.types import RequestInfo


class TestLengthBasedPredictor:
    """Tests for LengthBasedPredictor."""

    @pytest.fixture
    def predictor(self) -> LengthBasedPredictor:
        """Create a length-based predictor with default config."""
        return LengthBasedPredictor()

    @pytest.fixture
    def custom_predictor(self) -> LengthBasedPredictor:
        """Create a length-based predictor with custom config."""
        config = {
            "alpha": 0.0002,
            "beta": 0.1,
            "gamma": 1.0,
            "min_lifetime": 0.5,
            "max_lifetime": 100.0,
        }
        return LengthBasedPredictor(config)

    def test_initialization_default(self, predictor: LengthBasedPredictor) -> None:
        """Test initialization with default config."""
        assert predictor.alpha == 0.0001
        assert predictor.beta == 0.05
        assert predictor.gamma == 0.5
        assert predictor.min_lifetime == 0.1
        assert predictor.max_lifetime == 300.0

    def test_initialization_custom(self, custom_predictor: LengthBasedPredictor) -> None:
        """Test initialization with custom config."""
        assert custom_predictor.alpha == 0.0002
        assert custom_predictor.beta == 0.1
        assert custom_predictor.gamma == 1.0
        assert custom_predictor.min_lifetime == 0.5
        assert custom_predictor.max_lifetime == 100.0

    def test_invalid_alpha(self) -> None:
        """Test that negative alpha raises ValueError."""
        with pytest.raises(ValueError, match="alpha must be >= 0"):
            LengthBasedPredictor({"alpha": -0.1})

    def test_invalid_beta(self) -> None:
        """Test that negative beta raises ValueError."""
        with pytest.raises(ValueError, match="beta must be >= 0"):
            LengthBasedPredictor({"beta": -0.1})

    def test_invalid_gamma(self) -> None:
        """Test that negative gamma raises ValueError."""
        with pytest.raises(ValueError, match="gamma must be >= 0"):
            LengthBasedPredictor({"gamma": -0.1})

    def test_invalid_min_lifetime(self) -> None:
        """Test that non-positive min_lifetime raises ValueError."""
        with pytest.raises(ValueError, match="min_lifetime must be > 0"):
            LengthBasedPredictor({"min_lifetime": 0.0})

    def test_invalid_max_lifetime(self) -> None:
        """Test that max_lifetime <= min_lifetime raises ValueError."""
        with pytest.raises(ValueError, match="max_lifetime.*must be > min_lifetime"):
            LengthBasedPredictor({"min_lifetime": 10.0, "max_lifetime": 5.0})

    def test_predict_basic(self, predictor: LengthBasedPredictor) -> None:
        """Test basic prediction."""
        info = RequestInfo(
            request_id="req1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        lifetime = predictor.predict(info)

        # Expected: 0.0001 * 100 + 0.05 * 50 + 0.5 = 0.01 + 2.5 + 0.5 = 3.01
        assert abs(lifetime - 3.01) < 0.001

    def test_predict_zero_tokens(self, predictor: LengthBasedPredictor) -> None:
        """Test prediction with zero tokens."""
        info = RequestInfo(
            request_id="req2",
            model_name="Qwen2-7B",
            prompt_length=0,
            max_output_length=0,
        )

        lifetime = predictor.predict(info)

        # Should return gamma (base overhead)
        assert lifetime == 0.5

    def test_predict_large_tokens(self, predictor: LengthBasedPredictor) -> None:
        """Test prediction with very large token counts."""
        info = RequestInfo(
            request_id="req3",
            model_name="Qwen2-7B",
            prompt_length=10000,
            max_output_length=5000,
        )

        lifetime = predictor.predict(info)

        # Expected: 0.0001 * 10000 + 0.05 * 5000 + 0.5 = 1 + 250 + 0.5 = 251.5
        # But clamped to max_lifetime (300.0)
        assert lifetime == 251.5

    def test_predict_clamped_to_max(self, custom_predictor: LengthBasedPredictor) -> None:
        """Test prediction clamped to max_lifetime."""
        info = RequestInfo(
            request_id="req4",
            model_name="Qwen2-7B",
            prompt_length=10000,
            max_output_length=10000,
        )

        lifetime = custom_predictor.predict(info)

        # Expected: 0.0002 * 10000 + 0.1 * 10000 + 1.0 = 2 + 1000 + 1 = 1003
        # But clamped to max_lifetime (100.0)
        assert lifetime == 100.0

    def test_predict_clamped_to_min(self, custom_predictor: LengthBasedPredictor) -> None:
        """Test prediction clamped to min_lifetime."""
        info = RequestInfo(
            request_id="req5",
            model_name="Qwen2-7B",
            prompt_length=1,
            max_output_length=1,
        )

        lifetime = custom_predictor.predict(info)

        # Expected: 0.0002 * 1 + 0.1 * 1 + 1.0 = 1.1002
        # Should not be clamped (min_lifetime = 0.5)
        assert abs(lifetime - 1.1002) < 0.001

    def test_predict_empty_request_id(self, predictor: LengthBasedPredictor) -> None:
        """Test that empty request_id raises ValueError."""
        info = RequestInfo(
            request_id="",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        with pytest.raises(ValueError, match="request_id cannot be empty"):
            predictor.predict(info)

    def test_update(self, predictor: LengthBasedPredictor) -> None:
        """Test update method."""
        predictor.update("req1", 5.0)

        assert predictor._update_count == 1

    def test_update_invalid_lifetime(self, predictor: LengthBasedPredictor) -> None:
        """Test that negative actual_lifetime raises ValueError."""
        with pytest.raises(ValueError, match="actual_lifetime must be >= 0"):
            predictor.update("req1", -1.0)

    def test_get_confidence(self, predictor: LengthBasedPredictor) -> None:
        """Test get_confidence returns 0.7."""
        assert predictor.get_confidence() == 0.7

    def test_get_features_used(self, predictor: LengthBasedPredictor) -> None:
        """Test get_features_used returns correct features."""
        features = predictor.get_features_used()
        assert "prompt_length" in features
        assert "max_output_length" in features
        assert len(features) == 2


class TestHistoricalPredictor:
    """Tests for HistoricalPredictor."""

    @pytest.fixture
    def predictor(self) -> HistoricalPredictor:
        """Create a historical predictor with default config."""
        return HistoricalPredictor()

    @pytest.fixture
    def custom_predictor(self) -> HistoricalPredictor:
        """Create a historical predictor with custom config."""
        config = {
            "window_size": 50,
            "aggregation": "median",
            "split_by_type": False,
            "min_samples": 3,
            "fallback_lifetime": 10.0,
        }
        return HistoricalPredictor(config)

    def test_initialization_default(self, predictor: HistoricalPredictor) -> None:
        """Test initialization with default config."""
        assert predictor.window_size == 100
        assert predictor.aggregation == "mean"
        assert predictor.split_by_type is True
        assert predictor.min_samples == 5
        assert predictor.fallback_lifetime == 5.0

    def test_initialization_custom(self, custom_predictor: HistoricalPredictor) -> None:
        """Test initialization with custom config."""
        assert custom_predictor.window_size == 50
        assert custom_predictor.aggregation == "median"
        assert custom_predictor.split_by_type is False
        assert custom_predictor.min_samples == 3
        assert custom_predictor.fallback_lifetime == 10.0

    def test_invalid_window_size(self) -> None:
        """Test that non-positive window_size raises ValueError."""
        with pytest.raises(ValueError, match="window_size must be > 0"):
            HistoricalPredictor({"window_size": 0})

    def test_invalid_aggregation(self) -> None:
        """Test that invalid aggregation raises ValueError."""
        with pytest.raises(ValueError, match="aggregation must be"):
            HistoricalPredictor({"aggregation": "invalid"})

    def test_invalid_min_samples(self) -> None:
        """Test that min_samples < 1 raises ValueError."""
        with pytest.raises(ValueError, match="min_samples must be >= 1"):
            HistoricalPredictor({"min_samples": 0})

    def test_invalid_fallback_lifetime(self) -> None:
        """Test that non-positive fallback_lifetime raises ValueError."""
        with pytest.raises(ValueError, match="fallback_lifetime must be > 0"):
            HistoricalPredictor({"fallback_lifetime": 0.0})

    def test_predict_insufficient_history(self, predictor: HistoricalPredictor) -> None:
        """Test prediction with insufficient history returns fallback."""
        info = RequestInfo(
            request_id="req1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        lifetime = predictor.predict(info)

        # Should return fallback_lifetime
        assert lifetime == 5.0

    def test_predict_with_history_mean(self, predictor: HistoricalPredictor) -> None:
        """Test prediction with sufficient history (mean aggregation)."""
        predictor.min_samples = 3
        predictor.split_by_type = False  # Use global history

        # Add some history
        for i in range(5):
            predictor.update(f"req{i}", 3.0 + i * 0.5)

        info = RequestInfo(
            request_id="req_new",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        lifetime = predictor.predict(info)

        # Mean of [3.0, 3.5, 4.0, 4.5, 5.0] = 4.0
        assert abs(lifetime - 4.0) < 0.001

    def test_predict_with_history_median(self, custom_predictor: HistoricalPredictor) -> None:
        """Test prediction with sufficient history (median aggregation)."""
        # Add some history
        for i in range(5):
            custom_predictor.update(f"req{i}", 3.0 + i * 0.5)

        info = RequestInfo(
            request_id="req_new",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        lifetime = custom_predictor.predict(info)

        # Median of [3.0, 3.5, 4.0, 4.5, 5.0] = 4.0
        assert abs(lifetime - 4.0) < 0.001

    def test_predict_median_even_count(self, custom_predictor: HistoricalPredictor) -> None:
        """Test median calculation with even number of samples."""
        # Add 4 samples
        for i in range(4):
            custom_predictor.update(f"req{i}", 2.0 + i * 1.0)

        info = RequestInfo(
            request_id="req_new",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        lifetime = custom_predictor.predict(info)

        # Median of [2.0, 3.0, 4.0, 5.0] = (3.0 + 4.0) / 2 = 3.5
        assert abs(lifetime - 3.5) < 0.001

    def test_update_history_window(self) -> None:
        """Test that history respects window size."""
        # Create predictor with small window
        predictor = HistoricalPredictor({"window_size": 5})

        # Add more updates than window size
        for i in range(10):
            predictor.update(f"req{i}", float(i))

        # Global history should have only last 5
        assert len(predictor._global_history) == 5
        assert list(predictor._global_history) == [5.0, 6.0, 7.0, 8.0, 9.0]

    def test_predict_empty_request_id(self, predictor: HistoricalPredictor) -> None:
        """Test that empty request_id raises ValueError."""
        info = RequestInfo(
            request_id="",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        with pytest.raises(ValueError, match="request_id cannot be empty"):
            predictor.predict(info)

    def test_update_invalid_lifetime(self, predictor: HistoricalPredictor) -> None:
        """Test that negative actual_lifetime raises ValueError."""
        with pytest.raises(ValueError, match="actual_lifetime must be >= 0"):
            predictor.update("req1", -1.0)

    def test_get_confidence_low(self, predictor: HistoricalPredictor) -> None:
        """Test confidence with insufficient history."""
        assert predictor.get_confidence() == 0.3

    def test_get_confidence_moderate(self, predictor: HistoricalPredictor) -> None:
        """Test confidence with moderate history."""
        # Add samples up to 30% of window (between min and half)
        for i in range(30):
            predictor.update(f"req{i}", 3.0)

        # Should be moderate confidence (30 < 50 which is window_size//2)
        assert predictor.get_confidence() == 0.6

    def test_get_confidence_high(self, predictor: HistoricalPredictor) -> None:
        """Test confidence with full history."""
        # Add samples to fill window
        for i in range(predictor.window_size):
            predictor.update(f"req{i}", 3.0)

        assert predictor.get_confidence() == 0.9

    def test_get_features_used_split_by_type(self, predictor: HistoricalPredictor) -> None:
        """Test get_features_used when split_by_type is True."""
        features = predictor.get_features_used()
        assert "request_type" in features

    def test_get_features_used_no_split(self, custom_predictor: HistoricalPredictor) -> None:
        """Test get_features_used when split_by_type is False."""
        features = custom_predictor.get_features_used()
        assert len(features) == 0

    def test_get_statistics(self, predictor: HistoricalPredictor) -> None:
        """Test get_statistics returns correct data."""
        # Add some history
        for i in range(5):
            predictor.update(f"req{i}", 3.0 + i * 0.5)

        stats = predictor.get_statistics()

        assert stats["window_size"] == 100
        assert stats["aggregation"] == "mean"
        assert stats["global_history_size"] == 5
        assert "mean_lifetime" in stats
        assert "min_lifetime" in stats
        assert "max_lifetime" in stats

    def test_reset(self, predictor: HistoricalPredictor) -> None:
        """Test reset clears history."""
        # Add some history
        for i in range(5):
            predictor.update(f"req{i}", 3.0)

        predictor.reset()

        assert len(predictor._global_history) == 0
        assert len(predictor._history) == 0
        assert predictor._update_count == 0


class TestHybridHeuristicPredictor:
    """Tests for HybridHeuristicPredictor."""

    @pytest.fixture
    def predictor(self) -> HybridHeuristicPredictor:
        """Create a hybrid predictor with default config."""
        return HybridHeuristicPredictor()

    @pytest.fixture
    def custom_predictor(self) -> HybridHeuristicPredictor:
        """Create a hybrid predictor with custom config."""
        config = {
            "length_weight": 0.3,
            "historical_weight": 0.7,
            "length_config": {"alpha": 0.0002, "beta": 0.1},
            "historical_config": {"window_size": 50, "aggregation": "median"},
        }
        return HybridHeuristicPredictor(config)

    def test_initialization_default(self, predictor: HybridHeuristicPredictor) -> None:
        """Test initialization with default config."""
        # Weights should be normalized to sum to 1
        assert abs(predictor.length_weight - 0.4) < 0.001
        assert abs(predictor.historical_weight - 0.6) < 0.001

    def test_initialization_custom(self, custom_predictor: HybridHeuristicPredictor) -> None:
        """Test initialization with custom config."""
        # Weights should be normalized: 0.3/(0.3+0.7)=0.3, 0.7/(0.3+0.7)=0.7
        assert abs(custom_predictor.length_weight - 0.3) < 0.001
        assert abs(custom_predictor.historical_weight - 0.7) < 0.001

    def test_weight_normalization(self) -> None:
        """Test that weights are normalized to sum to 1."""
        config = {"length_weight": 0.25, "historical_weight": 0.75}
        predictor = HybridHeuristicPredictor(config)

        # Should be normalized: 0.25/(0.25+0.75)=0.25, 0.75/(0.25+0.75)=0.75
        assert abs(predictor.length_weight - 0.25) < 0.001
        assert abs(predictor.historical_weight - 0.75) < 0.001

    def test_invalid_length_weight(self) -> None:
        """Test that invalid length_weight raises ValueError."""
        with pytest.raises(ValueError, match="length_weight must be in"):
            HybridHeuristicPredictor({"length_weight": -0.1})

    def test_invalid_historical_weight(self) -> None:
        """Test that invalid historical_weight raises ValueError."""
        with pytest.raises(ValueError, match="historical_weight must be in"):
            HybridHeuristicPredictor({"historical_weight": -0.1})

    def test_zero_weights(self) -> None:
        """Test that both weights being zero raises ValueError."""
        with pytest.raises(ValueError, match="At least one weight must be > 0"):
            HybridHeuristicPredictor({"length_weight": 0.0, "historical_weight": 0.0})

    def test_predict_insufficient_history(self, predictor: HybridHeuristicPredictor) -> None:
        """Test prediction with insufficient history uses fallback."""
        info = RequestInfo(
            request_id="req1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        lifetime = predictor.predict(info)

        # Length: 0.0001 * 100 + 0.05 * 50 + 0.5 = 3.01
        # Historical: fallback = 5.0
        # Hybrid: 0.4 * 3.01 + 0.6 * 5.0 = 1.204 + 3.0 = 4.204
        assert abs(lifetime - 4.204) < 0.01

    def test_predict_with_history(self, predictor: HybridHeuristicPredictor) -> None:
        """Test prediction with sufficient history."""
        # Set split_by_type to False for consistent behavior
        predictor.historical_predictor.split_by_type = False

        # Add history
        for i in range(10):
            predictor.update(f"req{i}", 4.0)

        info = RequestInfo(
            request_id="req_new",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        lifetime = predictor.predict(info)

        # Length: 0.0001 * 100 + 0.05 * 50 + 0.5 = 3.01
        # Historical: mean of [4.0, ...] = 4.0
        # Hybrid: 0.4 * 3.01 + 0.6 * 4.0 = 1.204 + 2.4 = 3.604
        assert abs(lifetime - 3.604) < 0.01

    def test_predict_empty_request_id(self, predictor: HybridHeuristicPredictor) -> None:
        """Test that empty request_id raises ValueError."""
        info = RequestInfo(
            request_id="",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        with pytest.raises(ValueError, match="request_id cannot be empty"):
            predictor.predict(info)

    def test_update(self, predictor: HybridHeuristicPredictor) -> None:
        """Test update method updates both predictors."""
        predictor.update("req1", 5.0)

        assert predictor._update_count == 1
        assert predictor.length_predictor._update_count == 1
        assert predictor.historical_predictor._update_count == 1

    def test_update_invalid_lifetime(self, predictor: HybridHeuristicPredictor) -> None:
        """Test that negative actual_lifetime raises ValueError."""
        with pytest.raises(ValueError, match="actual_lifetime must be >= 0"):
            predictor.update("req1", -1.0)

    def test_get_confidence(self, predictor: HybridHeuristicPredictor) -> None:
        """Test get_confidence returns weighted average."""
        confidence = predictor.get_confidence()

        # Length confidence: 0.7
        # Historical confidence: 0.3 (no history)
        # Weighted: 0.4 * 0.7 + 0.6 * 0.3 = 0.28 + 0.18 = 0.46
        assert abs(confidence - 0.46) < 0.01

    def test_get_features_used(self, predictor: HybridHeuristicPredictor) -> None:
        """Test get_features_used returns combined features."""
        features = predictor.get_features_used()

        assert "prompt_length" in features
        assert "max_output_length" in features
        assert "request_type" in features

    def test_get_statistics(self, predictor: HybridHeuristicPredictor) -> None:
        """Test get_statistics returns data from both predictors."""
        stats = predictor.get_statistics()

        assert "length_weight" in stats
        assert "historical_weight" in stats
        assert "length_predictor_stats" in stats
        assert "historical_predictor_stats" in stats

    def test_reset(self, predictor: HybridHeuristicPredictor) -> None:
        """Test reset clears both predictors."""
        # Add some history
        predictor.update("req1", 5.0)

        predictor.reset()

        assert predictor._update_count == 0
        assert predictor.length_predictor._update_count == 0
        assert predictor.historical_predictor._update_count == 0
