"""Integration tests for lifetime predictor with scheduling policies."""

from __future__ import annotations

import pytest

from sagellm_control import (
    ControlPlaneManager,
    EngineState,
    ExecutionInstanceType,
    HybridHeuristicPredictor,
    LengthBasedPredictor,
)
from sagellm_control.policies import LifetimeAwarePolicy
from sagellm_control.types import RequestMetadata, RequestPriority


class TestLifetimeAwarePolicyIntegration:
    """Integration tests for LifetimeAwarePolicy."""

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

    @pytest.fixture
    def policy(self, predictor: LengthBasedPredictor) -> LifetimeAwarePolicy:
        """Create a lifetime-aware policy."""
        return LifetimeAwarePolicy(
            predictor=predictor,
            enable_predictor=True,
        )

    def test_policy_initialization(self, policy: LifetimeAwarePolicy) -> None:
        """Test policy initialization."""
        assert policy.name == "lifetime_aware"
        assert policy._enable_predictor is True
        assert policy._predictor is not None

    def test_predict_lifetime_with_predictor(self, policy: LifetimeAwarePolicy) -> None:
        """Test lifetime prediction with predictor enabled."""
        request = RequestMetadata(
            request_id="req1",
            trace_id="trace1",
            model_id="Qwen2-7B",
            prompt="Test prompt",
            max_tokens=50,
            prompt_tokens=10,
        )

        lifetime = policy._predict_lifetime(request)
        assert lifetime > 0
        # Should be approximately: 0.0001 * 10 + 0.05 * 50 + 0.5 = 3.001
        assert 2.5 < lifetime < 3.5

    def test_predict_lifetime_fallback(self) -> None:
        """Test lifetime prediction fallback when predictor is disabled."""
        policy = LifetimeAwarePolicy(enable_predictor=False)

        request = RequestMetadata(
            request_id="req1",
            trace_id="trace1",
            model_id="Qwen2-7B",
            prompt="Test prompt",
            max_tokens=50,
            prompt_tokens=10,
        )

        lifetime = policy._predict_lifetime(request)
        assert lifetime > 0

    def test_compute_priority_score(self, policy: LifetimeAwarePolicy) -> None:
        """Test priority score computation."""
        # Long request with normal priority
        request1 = RequestMetadata(
            request_id="req1",
            model_id="Qwen2-7B",
            priority=RequestPriority.NORMAL,
        )
        score1 = policy._compute_priority_score(request1, 10.0)

        # Short request with normal priority
        request2 = RequestMetadata(
            request_id="req2",
            model_id="Qwen2-7B",
            priority=RequestPriority.NORMAL,
        )
        score2 = policy._compute_priority_score(request2, 1.0)

        # Long request should have higher priority
        assert score1 > score2

    def test_kv_tokens_reservation(self, policy: LifetimeAwarePolicy) -> None:
        """Test KV cache token reservation for long requests."""
        request = RequestMetadata(
            request_id="req1",
            model_id="Qwen2-7B",
            prompt="Test prompt",
            max_tokens=100,
            prompt_tokens=50,
        )

        # Long lifetime (>= 5s) should get extra reservation
        kv_long = policy._reserve_kv_tokens_for_request(request, 10.0)
        kv_short = policy._reserve_kv_tokens_for_request(request, 1.0)

        assert kv_long > kv_short
        # kv_long should be approximately kv_short * 1.2
        assert kv_long == int(kv_short * 1.2)

    def test_reorder_queue_by_lifetime(self, policy: LifetimeAwarePolicy) -> None:
        """Test queue reordering by predicted lifetime."""
        requests = [
            RequestMetadata(
                request_id="req-short",
                model_id="Qwen2-7B",
                max_tokens=10,
                prompt_tokens=5,
            ),
            RequestMetadata(
                request_id="req-long",
                model_id="Qwen2-7B",
                max_tokens=200,
                prompt_tokens=100,
            ),
            RequestMetadata(
                request_id="req-medium",
                model_id="Qwen2-7B",
                max_tokens=50,
                prompt_tokens=25,
            ),
        ]

        reordered = policy.reorder_queue(requests)

        # Long request should be first (highest priority)
        assert reordered[0].request_id == "req-long"
        # Short request should be last (lowest priority)
        assert reordered[2].request_id == "req-short"

    def test_feedback_mechanism(self, policy: LifetimeAwarePolicy) -> None:
        """Test online feedback mechanism."""
        request = RequestMetadata(
            request_id="req1",
            model_id="Qwen2-7B",
            max_tokens=50,
            prompt_tokens=10,
        )

        # Predict lifetime
        policy._predict_lifetime(request)
        assert "req1" in policy._predicted_lifetimes

        # Provide feedback
        policy.feedback_request_completed("req1", 3.2)

        # Prediction should be removed after feedback
        assert "req1" not in policy._predicted_lifetimes

        # Predictor should have been updated
        assert policy._predictor._update_count > 0  # type: ignore[attr-defined]


class TestControlPlaneManagerPredictorIntegration:
    """Integration tests for ControlPlaneManager with predictor."""

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

    @pytest.fixture
    def manager(self, predictor: HybridHeuristicPredictor) -> ControlPlaneManager:
        """Create a ControlPlaneManager with predictor."""
        return ControlPlaneManager(
            scheduling_policy="lifetime_aware",
            routing_strategy="least_loaded",
            predictor=predictor,
            enable_predictor=True,
        )

    def test_manager_initialization_with_predictor(self, manager: ControlPlaneManager) -> None:
        """Test manager initialization with predictor."""
        assert manager._scheduling_policy_name == "lifetime_aware"
        assert manager._enable_predictor is True
        assert manager._predictor is not None

    def test_auto_create_predictor(self) -> None:
        """Test auto-creation of predictor when not provided."""
        manager = ControlPlaneManager(
            scheduling_policy="lifetime_aware",
            enable_predictor=True,
            predictor_config={"length_weight": 0.5, "historical_weight": 0.5},
        )

        # Policy should be created with auto predictor
        policy = manager._get_policy()
        assert isinstance(policy, LifetimeAwarePolicy)
        assert policy._predictor is not None

    @pytest.mark.asyncio
    async def test_schedule_with_predictor(self, manager: ControlPlaneManager) -> None:
        """Test request scheduling with predictor."""
        # Register an engine
        manager.register_engine(
            engine_id="engine-001",
            model_id="Qwen2-7B",
            host="localhost",
            port=8001,
            metadata={"instance_type": ExecutionInstanceType.GENERAL.value},
        )
        manager.update_engine_state("engine-001", EngineState.READY)

        # Schedule a request
        decision = await manager.schedule_request(
            request_id="req1",
            trace_id="trace1",
            model_id="Qwen2-7B",
            prompt="Test prompt",
            max_tokens=50,
        )

        assert decision.is_scheduled
        assert decision.engine_id == "engine-001"

        # Request start time should be tracked
        assert "req1" in manager._request_start_times

        # Cleanup
        manager.unregister_engine("engine-001")

    @pytest.mark.asyncio
    async def test_complete_with_feedback(self, manager: ControlPlaneManager) -> None:
        """Test request completion with predictor feedback."""
        # Register an engine
        manager.register_engine(
            engine_id="engine-001",
            model_id="Qwen2-7B",
            host="localhost",
            port=8001,
            metadata={"instance_type": ExecutionInstanceType.GENERAL.value},
        )
        manager.update_engine_state("engine-001", EngineState.READY)

        # Ensure policy is created
        policy = manager._get_policy()
        assert isinstance(policy, LifetimeAwarePolicy)

        # Schedule a request
        await manager.schedule_request(
            request_id="req1",
            trace_id="trace1",
            model_id="Qwen2-7B",
            prompt="Test prompt",
            max_tokens=50,
        )

        # Complete the request with feedback
        manager.complete_request("req1", "engine-001", actual_lifetime=3.5)

        # Predictor should have received feedback
        # Check the policy's predictor instead of manager's
        assert policy._predictor._update_count > 0  # type: ignore[attr-defined]

        # Request start time should be cleaned up
        assert "req1" not in manager._request_start_times

        # Cleanup
        manager.unregister_engine("engine-001")

    @pytest.mark.asyncio
    async def test_multiple_requests_with_priority(self, manager: ControlPlaneManager) -> None:
        """Test multiple requests with lifetime-based priority."""
        # Register engines
        manager.register_engine(
            engine_id="engine-001",
            model_id="Qwen2-7B",
            host="localhost",
            port=8001,
            metadata={"instance_type": ExecutionInstanceType.GENERAL.value},
        )
        manager.update_engine_state("engine-001", EngineState.READY)

        # Schedule multiple requests with different sizes
        decisions = []
        for i, max_tokens in enumerate([10, 100, 50]):
            decision = await manager.schedule_request(
                request_id=f"req{i}",
                trace_id=f"trace{i}",
                model_id="Qwen2-7B",
                prompt="Test prompt" * (i + 1),
                max_tokens=max_tokens,
            )
            decisions.append(decision)

        # All should be scheduled
        assert all(d.is_scheduled for d in decisions)

        # Cleanup
        manager.unregister_engine("engine-001")


def test_policy_without_predictor() -> None:
    """Test LifetimeAwarePolicy without predictor (disabled mode)."""
    policy = LifetimeAwarePolicy(enable_predictor=False)

    assert policy.name == "lifetime_aware"
    assert policy._enable_predictor is False

    # Should still work with heuristic fallback
    request = RequestMetadata(
        request_id="req1",
        model_id="Qwen2-7B",
        max_tokens=50,
        prompt_tokens=10,
    )

    lifetime = policy._predict_lifetime(request)
    assert lifetime > 0
