"""Tests for PrefixGroupingPolicy."""

from __future__ import annotations

import pytest

from sagellm_control.policies.prefix_grouping import PrefixGroupingPolicy
from sagellm_control.types import (
    EngineInfo,
    EngineState,
    RequestMetadata,
    RequestType,
)


@pytest.fixture
def policy() -> PrefixGroupingPolicy:
    """Create a PrefixGroupingPolicy instance."""
    return PrefixGroupingPolicy(min_prefix_len=10, wait_timeout_ms=100, min_group_size=2)


@pytest.fixture
def test_engine() -> EngineInfo:
    """Create a test engine."""
    return EngineInfo(
        engine_id="engine-1",
        model_id="test-model",
        host="localhost",
        port=8000,
        state=EngineState.READY,
        active_requests=0,
    )


def create_request(request_id: str, prompt: str) -> RequestMetadata:
    """Helper to create a test request."""
    return RequestMetadata(
        request_id=request_id,
        trace_id=f"trace-{request_id}",
        model_id="test-model",
        prompt=prompt,
        request_type=RequestType.LLM_CHAT,
    )


class TestPrefixGroupingPolicy:
    """Tests for PrefixGroupingPolicy."""

    def test_initialization(self, policy: PrefixGroupingPolicy) -> None:
        """Test policy initialization."""
        assert policy.name == "prefix_grouping"
        assert policy.grouper.min_prefix_len == 10
        assert policy.grouper.wait_timeout_ms == 100
        assert policy.grouper.min_group_size == 2

    def test_schedule_no_candidates(self, policy: PrefixGroupingPolicy) -> None:
        """Test scheduling with no available candidates."""
        req = create_request("req1", "Hello world, how are you?")
        decision = policy.schedule(req, [], [])
        assert not decision.is_scheduled
        assert decision.reason == "no_candidates"

    def test_schedule_no_healthy_engines(self, policy: PrefixGroupingPolicy) -> None:
        """Test scheduling with no healthy engines."""
        req = create_request("req1", "Hello world, how are you?")
        unhealthy_engine = EngineInfo(
            engine_id="engine-1",
            model_id="test-model",
            host="localhost",
            port=8000,
            state=EngineState.ERROR,
        )
        decision = policy.schedule(req, [unhealthy_engine], [])
        assert not decision.is_scheduled
        assert decision.reason == "no_healthy_engines"

    def test_schedule_waiting_for_group(
        self, policy: PrefixGroupingPolicy, test_engine: EngineInfo
    ) -> None:
        """Test that policy waits for more requests to form a group."""
        req = create_request("req1", "Hello world, how are you?")
        decision = policy.schedule(req, [test_engine], [])
        # Should wait since we don't have enough requests yet
        assert not decision.is_scheduled
        assert decision.reason == "waiting_for_group"

    def test_schedule_ready_group(
        self, policy: PrefixGroupingPolicy, test_engine: EngineInfo
    ) -> None:
        """Test scheduling when a group is ready."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")

        # First request - should wait
        decision1 = policy.schedule(req1, [test_engine], [])
        assert not decision1.is_scheduled

        # Second request - group should be ready now
        decision2 = policy.schedule(req2, [test_engine], [req1])
        # Group is ready (size >= min_group_size)
        if decision2.is_scheduled:
            assert decision2.engine_id == "engine-1"
        else:
            # May still be waiting depending on timing
            assert decision2.reason == "waiting_for_group"

    def test_schedule_selects_least_loaded(self, policy: PrefixGroupingPolicy) -> None:
        """Test that policy selects the least loaded engine."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")

        engine1 = EngineInfo(
            engine_id="engine-1",
            model_id="test-model",
            host="localhost",
            port=8000,
            state=EngineState.READY,
            active_requests=5,
        )
        engine2 = EngineInfo(
            engine_id="engine-2",
            model_id="test-model",
            host="localhost",
            port=8001,
            state=EngineState.READY,
            active_requests=2,  # Less loaded
        )

        policy.schedule(req1, [engine1, engine2], [])
        decision = policy.schedule(req2, [engine1, engine2], [req1])

        if decision.is_scheduled:
            # Should prefer engine-2 (less loaded)
            assert decision.engine_id == "engine-2"

    def test_reorder_queue(self, policy: PrefixGroupingPolicy) -> None:
        """Test reordering requests by prefix grouping."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Goodbye world, see you later!")
        req3 = create_request("req3", "Hello world, what's up?")
        req4 = create_request("req4", "Goodbye world, take care!")

        requests = [req1, req2, req3, req4]
        reordered = policy.reorder_queue(requests)

        # All 4 requests should be present
        assert len(reordered) == 4
        # Requests should be grouped (order within groups may vary)
        request_ids = [r.request_id for r in reordered]
        assert set(request_ids) == {"req1", "req2", "req3", "req4"}

    def test_reorder_queue_prioritizes_large_groups(self, policy: PrefixGroupingPolicy) -> None:
        """Test that larger groups are prioritized in reordering."""
        # Create a larger group
        large_group_reqs = [
            create_request(f"large-{i}", f"Common prefix for large group {i}") for i in range(5)
        ]
        # Create a smaller group
        small_group_reqs = [
            create_request(f"small-{i}", f"Different prefix for small group {i}") for i in range(2)
        ]

        all_requests = large_group_reqs + small_group_reqs
        reordered = policy.reorder_queue(all_requests)

        assert len(reordered) == 7
        # The larger group should come first (flexible check)
        first_five_ids = [r.request_id for r in reordered[:5]]
        large_group_ids = [r.request_id for r in large_group_reqs]
        # At least some requests from the large group should be prioritized
        assert len(set(first_five_ids) & set(large_group_ids)) >= 3

    def test_get_group_stats(self, policy: PrefixGroupingPolicy) -> None:
        """Test getting grouping statistics."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")

        # Add requests to internal grouper
        policy.grouper.add_request(req1)
        policy.grouper.add_request(req2)

        stats = policy.get_group_stats()
        assert "num_groups" in stats
        assert "total_pending" in stats
        assert stats["total_pending"] == 2
