"""Tests for Feature Extractor.

This module tests the FeatureExtractor class to ensure correct
feature extraction and handling of missing values.
"""

from __future__ import annotations

from datetime import datetime

import pytest

from sagellm_control.predictor.feature_extractor import ExtractedFeatures, FeatureExtractor
from sagellm_control.predictor.types import RequestInfo


class TestExtractedFeatures:
    """Tests for ExtractedFeatures dataclass."""

    def test_to_dict(self) -> None:
        """Test converting features to dictionary."""
        features = ExtractedFeatures(
            input_token_length=100,
            max_output_tokens=50,
            request_type="llm_generate",
            user_id="user123",
            timestamp=1234567890.0,
            hour_of_day=15,
            day_of_week=2,
            is_weekend=False,
            time_period="afternoon",
            priority=2,
            temperature=1.0,
            top_p=0.9,
            model_name="Qwen2-7B",
            metadata={"extra": "value"},
        )

        feature_dict = features.to_dict()

        assert feature_dict["input_token_length"] == 100
        assert feature_dict["max_output_tokens"] == 50
        assert feature_dict["request_type"] == "llm_generate"
        assert feature_dict["user_id"] == "user123"
        assert feature_dict["timestamp"] == 1234567890.0
        assert feature_dict["hour_of_day"] == 15
        assert feature_dict["day_of_week"] == 2
        assert feature_dict["is_weekend"] is False
        assert feature_dict["time_period"] == "afternoon"
        assert feature_dict["priority"] == 2
        assert feature_dict["temperature"] == 1.0
        assert feature_dict["top_p"] == 0.9
        assert feature_dict["model_name"] == "Qwen2-7B"
        assert feature_dict["extra"] == "value"

    def test_get_feature_names(self) -> None:
        """Test getting feature names."""
        features = ExtractedFeatures(
            input_token_length=100,
            max_output_tokens=50,
            request_type="llm_generate",
            user_id="user123",
            timestamp=1234567890.0,
            hour_of_day=15,
            day_of_week=2,
            is_weekend=False,
            time_period="afternoon",
            priority=2,
            temperature=1.0,
            top_p=0.9,
            model_name="Qwen2-7B",
            metadata={"extra1": "value1", "extra2": "value2"},
        )

        names = features.get_feature_names()

        assert "input_token_length" in names
        assert "max_output_tokens" in names
        assert "request_type" in names
        assert "user_id" in names
        assert "timestamp" in names
        assert "hour_of_day" in names
        assert "day_of_week" in names
        assert "is_weekend" in names
        assert "time_period" in names
        assert "priority" in names
        assert "temperature" in names
        assert "top_p" in names
        assert "model_name" in names
        assert "extra1" in names
        assert "extra2" in names


class TestFeatureExtractor:
    """Tests for FeatureExtractor class."""

    @pytest.fixture
    def extractor(self) -> FeatureExtractor:
        """Create a feature extractor for testing."""
        return FeatureExtractor()

    @pytest.fixture
    def sample_request(self) -> RequestInfo:
        """Create a sample request for testing."""
        return RequestInfo(
            request_id="req_123",
            model_name="Qwen2-7B",
            prompt_length=200,
            max_output_length=100,
            priority=2,
            arrival_time=datetime(2024, 3, 15, 14, 30, 0),  # Friday 2:30 PM
            request_type="llm_generate",
            temperature=0.8,
            top_p=0.95,
            metadata={"user_id": "user_456"},
        )

    def test_extract_basic_features(
        self, extractor: FeatureExtractor, sample_request: RequestInfo
    ) -> None:
        """Test extraction of basic features."""
        features = extractor.extract(sample_request)

        assert features.input_token_length == 200
        assert features.max_output_tokens == 100
        assert features.request_type == "llm_generate"
        assert features.user_id == "user_456"
        assert features.timestamp > 0
        assert features.priority == 2
        assert features.temperature == 0.8
        assert features.top_p == 0.95
        assert features.model_name == "Qwen2-7B"

    def test_extract_time_features_afternoon(
        self, extractor: FeatureExtractor, sample_request: RequestInfo
    ) -> None:
        """Test extraction of time features for afternoon."""
        features = extractor.extract(sample_request)

        assert features.hour_of_day == 14
        assert features.day_of_week == 4  # Friday (0=Monday)
        assert features.is_weekend is False
        assert features.time_period == "afternoon"

    def test_extract_time_features_morning(self, extractor: FeatureExtractor) -> None:
        """Test time features for morning."""
        request = RequestInfo(
            request_id="req_morning",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            arrival_time=datetime(2024, 3, 16, 9, 0, 0),  # Saturday 9:00 AM
        )

        features = extractor.extract(request)

        assert features.hour_of_day == 9
        assert features.time_period == "morning"
        assert features.day_of_week == 5  # Saturday
        assert features.is_weekend is True

    def test_extract_time_features_evening(self, extractor: FeatureExtractor) -> None:
        """Test time features for evening."""
        request = RequestInfo(
            request_id="req_evening",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            arrival_time=datetime(2024, 3, 17, 20, 0, 0),  # Sunday 8:00 PM
        )

        features = extractor.extract(request)

        assert features.hour_of_day == 20
        assert features.time_period == "evening"
        assert features.day_of_week == 6  # Sunday
        assert features.is_weekend is True

    def test_extract_time_features_night(self, extractor: FeatureExtractor) -> None:
        """Test time features for night."""
        request = RequestInfo(
            request_id="req_night",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            arrival_time=datetime(2024, 3, 18, 2, 0, 0),  # Monday 2:00 AM
        )

        features = extractor.extract(request)

        assert features.hour_of_day == 2
        assert features.time_period == "night"
        assert features.day_of_week == 0  # Monday
        assert features.is_weekend is False

    def test_user_id_from_metadata(self, extractor: FeatureExtractor) -> None:
        """Test user_id extraction from metadata."""
        request = RequestInfo(
            request_id="req_123",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            metadata={"user_id": "test_user"},
        )

        features = extractor.extract(request)

        assert features.user_id == "test_user"

    def test_user_id_from_request_id(self, extractor: FeatureExtractor) -> None:
        """Test user_id extraction from request_id prefix."""
        request = RequestInfo(
            request_id="user789_req_456",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        features = extractor.extract(request)

        assert features.user_id == "user789"

    def test_user_id_default_unknown(self, extractor: FeatureExtractor) -> None:
        """Test user_id defaults to 'unknown' when not provided."""
        request = RequestInfo(
            request_id="req_123",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
        )

        features = extractor.extract(request)

        assert features.user_id == "unknown"

    def test_additional_metadata_extraction(self, extractor: FeatureExtractor) -> None:
        """Test extraction of additional metadata fields."""
        request = RequestInfo(
            request_id="req_123",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            context_length=150,
            batch_size=4,
            num_beams=2,
            top_k=50,
        )

        features = extractor.extract(request)

        assert features.metadata["context_length"] == 150
        assert features.metadata["batch_size"] == 4
        assert features.metadata["num_beams"] == 2
        assert features.metadata["top_k"] == 50

    def test_metadata_not_added_for_default_values(self, extractor: FeatureExtractor) -> None:
        """Test that default values are not added to metadata."""
        request = RequestInfo(
            request_id="req_123",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            context_length=0,  # Default
            batch_size=1,  # Default
            num_beams=1,  # Default
            top_k=-1,  # Default
        )

        features = extractor.extract(request)

        assert "context_length" not in features.metadata
        assert "batch_size" not in features.metadata
        assert "num_beams" not in features.metadata
        assert "top_k" not in features.metadata

    def test_extract_batch(self, extractor: FeatureExtractor) -> None:
        """Test batch feature extraction."""
        requests = [
            RequestInfo(
                request_id=f"req_{i}",
                model_name="Qwen2-7B",
                prompt_length=100 + i * 10,
                max_output_length=50 + i * 5,
            )
            for i in range(5)
        ]

        features_list = extractor.extract_batch(requests)

        assert len(features_list) == 5
        for i, features in enumerate(features_list):
            assert features.input_token_length == 100 + i * 10
            assert features.max_output_tokens == 50 + i * 5

    def test_get_feature_count(self, extractor: FeatureExtractor) -> None:
        """Test getting feature count."""
        count = extractor.get_feature_count()
        assert count == 13  # Base feature count

    def test_get_extraction_count(
        self, extractor: FeatureExtractor, sample_request: RequestInfo
    ) -> None:
        """Test tracking extraction count."""
        assert extractor.get_extraction_count() == 0

        extractor.extract(sample_request)
        assert extractor.get_extraction_count() == 1

        extractor.extract(sample_request)
        assert extractor.get_extraction_count() == 2

        extractor.extract_batch([sample_request] * 3)
        assert extractor.get_extraction_count() == 5

    def test_extractor_with_config(self) -> None:
        """Test creating extractor with custom config."""
        config = {"some_option": "value"}
        extractor = FeatureExtractor(config=config)

        assert extractor.config == config

    def test_extractor_without_config(self) -> None:
        """Test creating extractor without config."""
        extractor = FeatureExtractor()

        assert extractor.config == {}

    def test_feature_extraction_performance(
        self, extractor: FeatureExtractor, sample_request: RequestInfo
    ) -> None:
        """Test that feature extraction is fast enough (< 0.05ms = 50 microseconds).

        This is a smoke test to ensure we meet the performance requirement.
        """
        import time

        # Warmup
        for _ in range(10):
            extractor.extract(sample_request)

        # Measure
        iterations = 1000
        start = time.perf_counter()
        for _ in range(iterations):
            extractor.extract(sample_request)
        end = time.perf_counter()

        avg_time_ms = ((end - start) / iterations) * 1000
        print(f"\nAverage extraction time: {avg_time_ms:.4f} ms")

        # Should be well under 0.05ms (50 microseconds)
        assert avg_time_ms < 0.05, f"Extraction too slow: {avg_time_ms:.4f} ms"

    def test_all_features_present(
        self, extractor: FeatureExtractor, sample_request: RequestInfo
    ) -> None:
        """Test that all expected features are present after extraction."""
        features = extractor.extract(sample_request)

        # Check all basic features
        assert hasattr(features, "input_token_length")
        assert hasattr(features, "max_output_tokens")
        assert hasattr(features, "request_type")
        assert hasattr(features, "user_id")
        assert hasattr(features, "timestamp")

        # Check all time features
        assert hasattr(features, "hour_of_day")
        assert hasattr(features, "day_of_week")
        assert hasattr(features, "is_weekend")
        assert hasattr(features, "time_period")

        # Check additional features
        assert hasattr(features, "priority")
        assert hasattr(features, "temperature")
        assert hasattr(features, "top_p")
        assert hasattr(features, "model_name")
        assert hasattr(features, "metadata")

    def test_feature_types(self, extractor: FeatureExtractor, sample_request: RequestInfo) -> None:
        """Test that extracted features have correct types."""
        features = extractor.extract(sample_request)

        assert isinstance(features.input_token_length, int)
        assert isinstance(features.max_output_tokens, int)
        assert isinstance(features.request_type, str)
        assert isinstance(features.user_id, str)
        assert isinstance(features.timestamp, float)
        assert isinstance(features.hour_of_day, int)
        assert isinstance(features.day_of_week, int)
        assert isinstance(features.is_weekend, bool)
        assert isinstance(features.time_period, str)
        assert isinstance(features.priority, int)
        assert isinstance(features.temperature, float)
        assert isinstance(features.top_p, float)
        assert isinstance(features.model_name, str)
        assert isinstance(features.metadata, dict)

    def test_time_period_boundaries(self, extractor: FeatureExtractor) -> None:
        """Test time period assignment at boundaries."""
        # Test boundary: 6:00 AM (start of morning)
        request_morning = RequestInfo(
            request_id="req_1",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            arrival_time=datetime(2024, 3, 15, 6, 0, 0),
        )
        features = extractor.extract(request_morning)
        assert features.time_period == "morning"

        # Test boundary: 12:00 PM (start of afternoon)
        request_afternoon = RequestInfo(
            request_id="req_2",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            arrival_time=datetime(2024, 3, 15, 12, 0, 0),
        )
        features = extractor.extract(request_afternoon)
        assert features.time_period == "afternoon"

        # Test boundary: 6:00 PM (start of evening)
        request_evening = RequestInfo(
            request_id="req_3",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            arrival_time=datetime(2024, 3, 15, 18, 0, 0),
        )
        features = extractor.extract(request_evening)
        assert features.time_period == "evening"

        # Test boundary: 10:00 PM (start of night)
        request_night = RequestInfo(
            request_id="req_4",
            model_name="Qwen2-7B",
            prompt_length=100,
            max_output_length=50,
            arrival_time=datetime(2024, 3, 15, 22, 0, 0),
        )
        features = extractor.extract(request_night)
        assert features.time_period == "night"
