"""Tests for Lifetime Predictor (Task 2.5).

Tests cover:
- Feature extraction
- Historical predictor (main implementation)
- Prediction accuracy and performance
"""

from __future__ import annotations

import time

import pytest

from sagellm_kv_cache.models import KVHandle
from sagellm_kv_cache.predictor import (
    FeatureExtractor,
    HistoricalPredictor,
    PredictorFeatures,
)


class TestPredictorFeatures:
    """Test feature extraction."""

    def test_from_handle_basic(self) -> None:
        """Test basic feature extraction from KVHandle."""
        handle = KVHandle.create(num_tokens=128, device="cuda:0")
        handle.access_count = 5
        handle.last_access = time.time()

        features = PredictorFeatures.from_handle(handle, hit_len=64)

        assert features.handle_id == handle.handle_id
        assert features.num_tokens == 128
        assert features.access_count == 5
        assert features.hit_len == 64
        assert features.device == "cuda:0"
        assert features.is_pinned is False

    def test_from_handle_with_session(self) -> None:
        """Test feature extraction with session ID."""
        handle = KVHandle.create(
            num_tokens=256,
            session_id="session-123",
            request_id="req-456",
        )

        features = PredictorFeatures.from_handle(handle)

        assert features.session_id == "session-123"
        assert features.request_id == "req-456"

    def test_to_dict_and_from_dict(self) -> None:
        """Test serialization and deserialization."""
        handle = KVHandle.create(num_tokens=64)
        features = PredictorFeatures.from_handle(handle)

        data = features.to_dict()
        restored = PredictorFeatures.from_dict(data)

        assert restored.handle_id == features.handle_id
        assert restored.num_tokens == features.num_tokens


class TestFeatureExtractor:
    """Test batch feature extraction."""

    def test_extract_batch(self) -> None:
        """Test extracting features from multiple handles."""
        handles = [KVHandle.create(num_tokens=i * 100) for i in range(1, 4)]
        extractor = FeatureExtractor()

        features = extractor.extract_batch(handles)

        assert len(features) == 3
        assert features[0].num_tokens == 100
        assert features[1].num_tokens == 200
        assert features[2].num_tokens == 300

    def test_extract_batch_with_hit_lens(self) -> None:
        """Test extraction with prefix hit lengths."""
        handles = [KVHandle.create(num_tokens=200) for _ in range(3)]
        hit_lens = [10, 20, 30]
        extractor = FeatureExtractor()

        features = extractor.extract_batch(handles, hit_lens=hit_lens)

        assert features[0].hit_len == 10
        assert features[1].hit_len == 20
        assert features[2].hit_len == 30

    def test_extract_batch_hit_lens_mismatch(self) -> None:
        """Test error when hit_lens length doesn't match handles."""
        handles = [KVHandle.create(num_tokens=100) for _ in range(3)]
        hit_lens = [10, 20]  # Wrong length
        extractor = FeatureExtractor()

        with pytest.raises(ValueError, match="hit_lens length"):
            extractor.extract_batch(handles, hit_lens=hit_lens)


class TestHistoricalPredictor:
    """Test historical predictor (main implementation)."""

    def test_predict_batch_cold_start(self) -> None:
        """Test prediction with no training data (cold start)."""
        predictor = HistoricalPredictor()
        handles = [KVHandle.create(num_tokens=100) for _ in range(3)]
        features = [PredictorFeatures.from_handle(h) for h in handles]

        predictions = predictor.predict_batch(features)

        assert len(predictions) == 3
        for pred in predictions:
            assert 0 <= pred.reuse_prob <= 1
            assert pred.remaining_lifetime > 0
            assert pred.evict_score > 0
            # Cold start should have lower confidence
            assert pred.confidence < 1.0

    def test_predict_batch_empty(self) -> None:
        """Test error on empty input."""
        predictor = HistoricalPredictor()

        with pytest.raises(ValueError, match="must not be empty"):
            predictor.predict_batch([])

    def test_update_and_predict(self) -> None:
        """Test that predictor learns from updates."""
        predictor = HistoricalPredictor(alpha=0.5)  # Fast learning

        # Create a handle and extract features
        handle = KVHandle.create(num_tokens=200, session_id="session-1")
        features = PredictorFeatures.from_handle(handle)

        # Get initial prediction
        initial_pred = predictor.predict_batch([features])[0]
        initial_lifetime = initial_pred.remaining_lifetime

        # Simulate observing a very long lifetime
        predictor.update(features, actual_lifetime=10000.0)

        # Predict again - should reflect the update
        updated_pred = predictor.predict_batch([features])[0]

        # After observing long lifetime, prediction should increase
        assert updated_pred.remaining_lifetime > initial_lifetime
        # Confidence should increase with more samples
        assert updated_pred.confidence > initial_pred.confidence

    def test_session_aware_prediction(self) -> None:
        """Test that predictor maintains per-session statistics."""
        predictor = HistoricalPredictor()

        # Train on session-1 with long lifetimes
        h1 = KVHandle.create(num_tokens=100, session_id="session-1")
        f1 = PredictorFeatures.from_handle(h1)
        for _ in range(5):
            predictor.update(f1, actual_lifetime=5000.0)

        # Train on session-2 with short lifetimes
        h2 = KVHandle.create(num_tokens=100, session_id="session-2")
        f2 = PredictorFeatures.from_handle(h2)
        for _ in range(5):
            predictor.update(f2, actual_lifetime=100.0)

        # Predict for both sessions
        predictions = predictor.predict_batch([f1, f2])

        # Session-1 should have longer predicted lifetime
        assert predictions[0].remaining_lifetime > predictions[1].remaining_lifetime

    def test_token_size_effect(self) -> None:
        """Test that token count affects predictions."""
        predictor = HistoricalPredictor()

        # Create handles with different token counts
        h_small = KVHandle.create(num_tokens=50)
        h_large = KVHandle.create(num_tokens=500)

        features = [
            PredictorFeatures.from_handle(h_small),
            PredictorFeatures.from_handle(h_large),
        ]

        predictions = predictor.predict_batch(features)

        # Larger KV should generally have longer predicted lifetime
        assert predictions[1].remaining_lifetime > predictions[0].remaining_lifetime

    def test_recency_effect(self) -> None:
        """Test that recent access increases reuse probability."""
        predictor = HistoricalPredictor()

        current_time = time.time()

        # Create two handles: one recently accessed, one not
        h_recent = KVHandle.create(num_tokens=100)
        h_recent.last_access = current_time - 10  # 10 seconds ago

        h_old = KVHandle.create(num_tokens=100)
        h_old.last_access = current_time - 3600  # 1 hour ago

        features = [
            PredictorFeatures.from_handle(h_recent),
            PredictorFeatures.from_handle(h_old),
        ]

        predictions = predictor.predict_batch(features)

        # Recently accessed should have higher reuse probability
        assert predictions[0].reuse_prob > predictions[1].reuse_prob

    def test_get_statistics(self) -> None:
        """Test statistics retrieval."""
        predictor = HistoricalPredictor()

        # Add some training data
        handle = KVHandle.create(num_tokens=100, session_id="session-1")
        features = PredictorFeatures.from_handle(handle)
        predictor.update(features, actual_lifetime=1000.0)

        stats = predictor.get_statistics()

        assert "global" in stats
        assert "session-1" in stats
        assert stats["global"]["sample_count"] > 0
        assert stats["session-1"]["sample_count"] > 0

    def test_get_model_version(self) -> None:
        """Test model version retrieval."""
        predictor = HistoricalPredictor()
        version = predictor.get_model_version()

        assert "historical" in version.lower()
        assert len(version) > 0

    def test_window_size_limit(self) -> None:
        """Test that predictor respects window size limit."""
        predictor = HistoricalPredictor(window_size=3)

        # Create many sessions
        for i in range(10):
            handle = KVHandle.create(num_tokens=100, session_id=f"session-{i}")
            features = PredictorFeatures.from_handle(handle)
            predictor.update(features, actual_lifetime=1000.0)

        stats = predictor.get_statistics()

        # Should not exceed window size + 1 (for global)
        assert len(stats) <= predictor.window_size + 1


class TestPredictorPerformance:
    """Test predictor performance requirements."""

    def test_prediction_latency(self) -> None:
        """Test that prediction meets <1ms latency requirement."""
        predictor = HistoricalPredictor()

        # Create a batch of handles
        handles = [KVHandle.create(num_tokens=100) for _ in range(10)]
        features = [PredictorFeatures.from_handle(h) for h in handles]

        # Measure prediction time
        start = time.time()
        predictions = predictor.predict_batch(features)
        elapsed = time.time() - start

        # Should complete in < 10ms for batch of 10 (generous margin)
        assert elapsed < 0.01
        assert len(predictions) == 10

    def test_update_performance(self) -> None:
        """Test that updates are efficient."""
        predictor = HistoricalPredictor()

        handle = KVHandle.create(num_tokens=100)
        features = PredictorFeatures.from_handle(handle)

        # Measure update time for 100 updates
        start = time.time()
        for _ in range(100):
            predictor.update(features, actual_lifetime=1000.0)
        elapsed = time.time() - start

        # Should complete quickly (< 10ms for 100 updates)
        assert elapsed < 0.01


class TestPredictorIntegration:
    """Integration tests with eviction policies."""

    def test_predictions_usable_by_eviction(self) -> None:
        """Test that predictions can be consumed by eviction policies."""
        predictor = HistoricalPredictor()

        # Create handles with different characteristics
        handles = []
        for i in range(5):
            h = KVHandle.create(num_tokens=100 * (i + 1))
            h.access_count = i
            handles.append(h)

        # Get predictions
        features = [PredictorFeatures.from_handle(h) for h in handles]
        predictions = predictor.predict_batch(features)

        # Sort by evict_score (higher = more evictable)
        sorted_preds = sorted(predictions, key=lambda p: p.evict_score, reverse=True)

        # Verify we can rank handles by evictability
        assert len(sorted_preds) == 5
        assert all(p.evict_score > 0 for p in sorted_preds)

        # First in sorted list should be most evictable
        # (typically low access count)
        most_evictable = sorted_preds[0]
        least_evictable = sorted_preds[-1]

        assert most_evictable.evict_score > least_evictable.evict_score

    def test_predictions_contain_required_fields(self) -> None:
        """Test that predictions contain all Protocol-required fields."""
        predictor = HistoricalPredictor()

        handle = KVHandle.create(num_tokens=100, request_id="req-123")
        features = PredictorFeatures.from_handle(handle)

        predictions = predictor.predict_batch([features])
        pred = predictions[0]

        # Verify all required Protocol fields
        assert pred.handle_id == handle.handle_id
        assert pred.request_id == "req-123"
        assert isinstance(pred.reuse_prob, float)
        assert isinstance(pred.remaining_lifetime, float)
        assert isinstance(pred.evict_score, float)
        assert isinstance(pred.confidence, float)
        assert isinstance(pred.features_summary, dict)
        assert isinstance(pred.predicted_at, float)

        # Verify value ranges
        assert 0 <= pred.reuse_prob <= 1
        assert pred.remaining_lifetime >= 0
        assert 0 <= pred.confidence <= 1
