"""Tests for LifetimePredictor base class.

This module tests the abstract base class interface that all
lifetime predictors must implement.
"""

from __future__ import annotations

from datetime import datetime

import pytest

from sagellm_control.predictor.base import LifetimePredictor
from sagellm_control.predictor.types import PredictionResult, RequestInfo


class StubPredictor(LifetimePredictor):
    """Stub predictor for testing base class functionality."""

    def __init__(self, config: dict[str, object] | None = None) -> None:
        """Initialize stub predictor."""
        super().__init__(config)
        self._predictions: dict[str, float] = {}
        self._updates: dict[str, float] = {}

    def predict(self, request_info: RequestInfo) -> float:
        """Return a simple prediction based on token lengths."""
        if not request_info:
            raise ValueError("request_info cannot be None")

        # Simple formula: 0.001 * (prompt + max_output)
        lifetime = 0.001 * (request_info.prompt_length + request_info.max_output_length)
        self._predictions[request_info.request_id] = lifetime
        return lifetime

    def update(self, request_id: str, actual_lifetime: float) -> None:
        """Record the actual lifetime."""
        if actual_lifetime < 0:
            raise ValueError(f"actual_lifetime must be >= 0, got {actual_lifetime}")

        self._updates[request_id] = actual_lifetime
        self._update_count += 1

    def get_features_used(self) -> list[str]:
        """Return features used by this predictor."""
        return ["prompt_length", "max_output_length"]


class TestLifetimePredictorBase:
    """Tests for LifetimePredictor base class."""

    @pytest.fixture
    def predictor(self) -> StubPredictor:
        """Create a stub predictor for testing."""
        return StubPredictor()

    @pytest.fixture
    def request_info(self) -> RequestInfo:
        """Create a sample request info."""
        return RequestInfo(
            request_id="req1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

    def test_initialization_default(self) -> None:
        """Test initialization with default config."""
        predictor = StubPredictor()
        assert predictor.config == {}
        assert predictor._prediction_count == 0
        assert predictor._update_count == 0

    def test_initialization_with_config(self) -> None:
        """Test initialization with custom config."""
        config = {"param1": "value1", "param2": 42}
        predictor = StubPredictor(config)
        assert predictor.config == config
        assert predictor.config["param1"] == "value1"
        assert predictor.config["param2"] == 42

    def test_predict_basic(self, predictor: StubPredictor, request_info: RequestInfo) -> None:
        """Test basic predict() method."""
        lifetime = predictor.predict(request_info)

        assert isinstance(lifetime, float)
        assert lifetime > 0
        # Expected: 0.001 * (100 + 50) = 0.15
        assert abs(lifetime - 0.15) < 0.001

    def test_predict_increments_count(self, predictor: StubPredictor) -> None:
        """Test that predict_detailed() increments prediction count."""
        info1 = RequestInfo(
            request_id="req1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )
        info2 = RequestInfo(
            request_id="req2",
            model_name="Qwen2-7B",
            prompt_length=200,
            max_output_length=100,
        )

        assert predictor._prediction_count == 0
        predictor.predict_detailed(info1)
        assert predictor._prediction_count == 1
        predictor.predict_detailed(info2)
        assert predictor._prediction_count == 2

    def test_predict_detailed(self, predictor: StubPredictor, request_info: RequestInfo) -> None:
        """Test predict_detailed() returns PredictionResult."""
        result = predictor.predict_detailed(request_info)

        assert isinstance(result, PredictionResult)
        assert result.request_id == "req1"
        assert result.predicted_lifetime > 0
        assert 0.0 <= result.confidence <= 1.0
        assert isinstance(result.predicted_at, datetime)
        assert isinstance(result.features_used, list)
        assert "predictor_type" in result.metadata
        assert result.metadata["predictor_type"] == "StubPredictor"

    def test_update_basic(self, predictor: StubPredictor) -> None:
        """Test basic update() method."""
        predictor.update("req1", 5.0)

        assert predictor._update_count == 1
        assert "req1" in predictor._updates
        assert predictor._updates["req1"] == 5.0

    def test_update_multiple(self, predictor: StubPredictor) -> None:
        """Test multiple updates."""
        predictor.update("req1", 5.0)
        predictor.update("req2", 10.0)
        predictor.update("req3", 15.0)

        assert predictor._update_count == 3
        assert len(predictor._updates) == 3

    def test_update_negative_lifetime_raises(self, predictor: StubPredictor) -> 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_default(self, predictor: StubPredictor) -> None:
        """Test default confidence score."""
        confidence = predictor.get_confidence()
        assert confidence == 1.0

    def test_get_features_used(self, predictor: StubPredictor) -> None:
        """Test get_features_used() method."""
        features = predictor.get_features_used()
        assert isinstance(features, list)
        assert "prompt_length" in features
        assert "max_output_length" in features

    def test_get_statistics(self, predictor: StubPredictor, request_info: RequestInfo) -> None:
        """Test get_statistics() method."""
        # Make some predictions and updates
        predictor.predict_detailed(request_info)
        predictor.update("req1", 5.0)

        stats = predictor.get_statistics()
        assert isinstance(stats, dict)
        assert stats["prediction_count"] == 1
        assert stats["update_count"] == 1
        assert stats["predictor_type"] == "StubPredictor"

    def test_reset(self, predictor: StubPredictor, request_info: RequestInfo) -> None:
        """Test reset() method."""
        # Make predictions and updates
        predictor.predict_detailed(request_info)
        predictor.update("req1", 5.0)

        assert predictor._prediction_count > 0
        assert predictor._update_count > 0

        # Reset
        predictor.reset()

        assert predictor._prediction_count == 0
        assert predictor._update_count == 0


class TestLifetimePredictorInterface:
    """Tests to ensure abstract interface is enforced."""

    def test_cannot_instantiate_abstract_class(self) -> None:
        """Test that LifetimePredictor cannot be instantiated directly."""
        with pytest.raises(TypeError, match="Can't instantiate abstract class"):
            LifetimePredictor()  # type: ignore

    def test_subclass_must_implement_predict(self) -> None:
        """Test that subclass must implement predict()."""

        class IncompletePredictor(LifetimePredictor):
            def update(self, request_id: str, actual_lifetime: float) -> None:
                pass

        with pytest.raises(TypeError, match="Can't instantiate abstract class"):
            IncompletePredictor()  # type: ignore

    def test_subclass_must_implement_update(self) -> None:
        """Test that subclass must implement update()."""

        class IncompletePredictor(LifetimePredictor):
            def predict(self, request_info: RequestInfo) -> float:
                return 1.0

        with pytest.raises(TypeError, match="Can't instantiate abstract class"):
            IncompletePredictor()  # type: ignore


class TestRequestInfo:
    """Tests for RequestInfo dataclass."""

    def test_initialization_minimal(self) -> None:
        """Test initialization with minimal required fields."""
        info = RequestInfo(
            request_id="req1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        assert info.request_id == "req1"
        assert info.model_name == "Qwen2-7B"
        assert info.prompt_length == 100
        assert info.max_output_length == 50
        assert info.priority == 2  # Default
        assert info.request_type == "llm_generate"  # Default

    def test_initialization_full(self) -> None:
        """Test initialization with all fields."""
        arrival_time = datetime.now()
        info = RequestInfo(
            request_id="req1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            priority=1,
            arrival_time=arrival_time,
            request_type="llm_chat",
            temperature=0.8,
            top_p=0.9,
            top_k=50,
            batch_size=2,
            context_length=1024,
            num_beams=4,
            metadata={"user": "test_user"},
        )

        assert info.priority == 1
        assert info.arrival_time == arrival_time
        assert info.request_type == "llm_chat"
        assert info.temperature == 0.8
        assert info.top_p == 0.9
        assert info.top_k == 50
        assert info.batch_size == 2
        assert info.context_length == 1024
        assert info.num_beams == 4
        assert info.metadata["user"] == "test_user"

    def test_negative_prompt_length_raises(self) -> None:
        """Test that negative prompt_length raises ValueError."""
        with pytest.raises(ValueError, match="prompt_length must be >= 0"):
            RequestInfo(
                request_id="req1",
                model_name="Qwen2-7B",
                prompt_length=-10,
                max_output_length=50,
            )

    def test_negative_max_output_length_raises(self) -> None:
        """Test that negative max_output_length raises ValueError."""
        with pytest.raises(ValueError, match="max_output_length must be >= 0"):
            RequestInfo(
                request_id="req1",
                model_name="Qwen2-7B",
                prompt_length=100,
                max_output_length=-50,
            )

    def test_invalid_priority_raises(self) -> None:
        """Test that invalid priority raises ValueError."""
        with pytest.raises(ValueError, match="priority must be 0-4"):
            RequestInfo(
                request_id="req1",
                model_name="Qwen2-7B",
                prompt_length=100,
                max_output_length=50,
                priority=10,
            )

    def test_invalid_temperature_raises(self) -> None:
        """Test that invalid temperature raises ValueError."""
        with pytest.raises(ValueError, match="temperature must be 0.0-2.0"):
            RequestInfo(
                request_id="req1",
                model_name="Qwen2-7B",
                prompt_length=100,
                max_output_length=50,
                temperature=3.0,
            )

    def test_invalid_top_p_raises(self) -> None:
        """Test that invalid top_p raises ValueError."""
        with pytest.raises(ValueError, match="top_p must be 0.0-1.0"):
            RequestInfo(
                request_id="req1",
                model_name="Qwen2-7B",
                prompt_length=100,
                max_output_length=50,
                top_p=1.5,
            )


class TestPredictionResult:
    """Tests for PredictionResult dataclass."""

    def test_initialization_minimal(self) -> None:
        """Test initialization with minimal fields."""
        result = PredictionResult(
            request_id="req1",
            predicted_lifetime=5.0,
        )

        assert result.request_id == "req1"
        assert result.predicted_lifetime == 5.0
        assert result.confidence == 1.0  # Default
        assert isinstance(result.predicted_at, datetime)
        assert result.features_used == []  # Default
        assert result.metadata == {}  # Default

    def test_initialization_full(self) -> None:
        """Test initialization with all fields."""
        predicted_at = datetime.now()
        result = PredictionResult(
            request_id="req1",
            predicted_lifetime=5.0,
            confidence=0.8,
            predicted_at=predicted_at,
            features_used=["prompt_length", "max_output_length"],
            metadata={"predictor": "test"},
        )

        assert result.confidence == 0.8
        assert result.predicted_at == predicted_at
        assert len(result.features_used) == 2
        assert result.metadata["predictor"] == "test"

    def test_negative_predicted_lifetime_raises(self) -> None:
        """Test that negative predicted_lifetime raises ValueError."""
        with pytest.raises(ValueError, match="predicted_lifetime must be >= 0"):
            PredictionResult(
                request_id="req1",
                predicted_lifetime=-5.0,
            )

    def test_invalid_confidence_raises(self) -> None:
        """Test that invalid confidence raises ValueError."""
        with pytest.raises(ValueError, match="confidence must be 0.0-1.0"):
            PredictionResult(
                request_id="req1",
                predicted_lifetime=5.0,
                confidence=1.5,
            )

    def test_zero_predicted_lifetime_valid(self) -> None:
        """Test that zero predicted_lifetime is valid."""
        result = PredictionResult(
            request_id="req1",
            predicted_lifetime=0.0,
        )
        assert result.predicted_lifetime == 0.0
