"""Tests for PrefixGrouper."""

from __future__ import annotations

import time

import pytest

from sagellm_control.grouping import PrefixGroup, PrefixGrouper
from sagellm_control.types import RequestMetadata, RequestType


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


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 TestPrefixGroup:
    """Tests for PrefixGroup dataclass."""

    def test_prefix_group_creation(self) -> None:
        """Test creating a prefix group."""
        group = PrefixGroup(prefix="Hello world", prefix_tokens=2)
        assert group.prefix == "Hello world"
        assert group.prefix_tokens == 2
        assert group.size == 0
        assert len(group.requests) == 0

    def test_prefix_group_with_requests(self) -> None:
        """Test prefix group with requests."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")
        group = PrefixGroup(prefix="Hello world", requests=[req1, req2])
        assert group.size == 2
        assert group.requests == [req1, req2]

    def test_prefix_group_age(self) -> None:
        """Test age calculation."""
        group = PrefixGroup(prefix="test")
        time.sleep(0.05)  # 50ms
        assert group.age_ms >= 45  # Allow some tolerance


class TestPrefixGrouper:
    """Tests for PrefixGrouper."""

    def test_initialization(self, grouper: PrefixGrouper) -> None:
        """Test grouper initialization."""
        assert grouper.min_prefix_len == 10
        assert grouper.wait_timeout_ms == 100
        assert grouper.min_group_size == 2
        assert not grouper.has_pending_requests()
        assert grouper.pending_count() == 0

    def test_add_request_too_short(self, grouper: PrefixGrouper) -> None:
        """Test adding a request with too short prompt."""
        req = create_request("req1", "Hi")
        grouper.add_request(req)
        assert grouper.pending_count() == 1
        ungrouped = grouper.get_ungrouped_requests()
        assert len(ungrouped) == 1
        assert ungrouped[0].request_id == "req1"

    def test_add_request_creates_group(self, grouper: PrefixGrouper) -> None:
        """Test that adding a request creates a new group."""
        req = create_request("req1", "Hello world, this is a test prompt")
        grouper.add_request(req)
        assert grouper.has_pending_requests()
        assert grouper.pending_count() == 1

    def test_matching_requests_grouped(self, grouper: PrefixGrouper) -> None:
        """Test that matching requests are grouped together."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")
        req3 = create_request("req3", "Hello world, nice weather!")

        grouper.add_request(req1)
        grouper.add_request(req2)
        grouper.add_request(req3)

        assert grouper.pending_count() == 3

        # Force get all groups
        groups = grouper.get_ready_groups(force_all=True)
        assert len(groups) == 1
        assert groups[0].size == 3

    def test_different_prefixes_separate_groups(self, grouper: PrefixGrouper) -> None:
        """Test that different prefixes create separate groups."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")
        req3 = create_request("req3", "Goodbye world, see you later!")
        req4 = create_request("req4", "Goodbye world, take care!")

        grouper.add_request(req1)
        grouper.add_request(req2)
        grouper.add_request(req3)
        grouper.add_request(req4)

        assert grouper.pending_count() == 4

        groups = grouper.get_ready_groups(force_all=True)
        assert len(groups) == 2
        # Each group should have 2 requests
        assert all(g.size == 2 for g in groups)

    def test_get_ready_groups_by_size(self, grouper: PrefixGrouper) -> None:
        """Test that groups become ready when reaching min_group_size."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")

        grouper.add_request(req1)
        # Only 1 request, not ready yet
        ready = grouper.get_ready_groups()
        assert len(ready) == 0

        grouper.add_request(req2)
        # Now 2 requests, should be ready
        ready = grouper.get_ready_groups()
        assert len(ready) == 1
        assert ready[0].size == 2

    def test_get_ready_groups_by_timeout(self, grouper: PrefixGrouper) -> None:
        """Test that groups become ready after timeout."""
        req = create_request("req1", "Hello world, this is a test")
        grouper.add_request(req)

        # Not ready immediately
        ready = grouper.get_ready_groups()
        assert len(ready) == 0

        # Wait for timeout
        time.sleep(0.15)  # 150ms > 100ms timeout

        # Should be ready now
        ready = grouper.get_ready_groups()
        assert len(ready) == 1
        assert ready[0].size == 1

    def test_longest_common_prefix(self) -> None:
        """Test the LCP algorithm."""
        lcp = PrefixGrouper._longest_common_prefix("Hello world", "Hello there")
        assert lcp == "Hello "

        lcp = PrefixGrouper._longest_common_prefix("abcdef", "abcxyz")
        assert lcp == "abc"

        lcp = PrefixGrouper._longest_common_prefix("test", "testing")
        assert lcp == "test"

        lcp = PrefixGrouper._longest_common_prefix("abc", "xyz")
        assert lcp == ""

    def test_adjust_timeout_high_load(self, grouper: PrefixGrouper) -> None:
        """Test timeout adjustment under high load."""
        initial_timeout = grouper.wait_timeout_ms
        grouper.adjust_timeout(load_factor=0.9)
        assert grouper.wait_timeout_ms < initial_timeout

    def test_adjust_timeout_low_load(self, grouper: PrefixGrouper) -> None:
        """Test timeout adjustment under low load."""
        initial_timeout = grouper.wait_timeout_ms
        grouper.adjust_timeout(load_factor=0.2)
        assert grouper.wait_timeout_ms > initial_timeout

    def test_clear(self, grouper: PrefixGrouper) -> None:
        """Test clearing all groups."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Short")
        grouper.add_request(req1)
        grouper.add_request(req2)
        assert grouper.has_pending_requests()

        grouper.clear()
        assert not grouper.has_pending_requests()
        assert grouper.pending_count() == 0

    def test_get_stats(self, grouper: PrefixGrouper) -> None:
        """Test getting grouper statistics."""
        req1 = create_request("req1", "Hello world, how are you?")
        req2 = create_request("req2", "Hello world, what's up?")
        req3 = create_request("req3", "Short")

        grouper.add_request(req1)
        grouper.add_request(req2)
        grouper.add_request(req3)

        stats = grouper.get_stats()
        assert stats["num_groups"] == 1
        assert stats["num_ungrouped"] == 1
        assert stats["total_pending"] == 3
        assert stats["avg_group_size"] == 2.0
        assert stats["wait_timeout_ms"] == 100.0
