"""Performance benchmarks for fair task queue."""

import asyncio
import time
from typing import Any

import pytest

from haolib.background.fair import FairTaskQueue
from haolib.background.fair.storages.memory import InMemoryFairQueueStorage
from tests.integration.background.fair.conftest import (
    create_task_invocation,
    group_key_extractor,
    mock_task_adapter,
)
from tests.integration.background.fair.constants import (
    THROUGHPUT_HIGH,
    THROUGHPUT_MEDIUM,
    THROUGHPUT_MINIMUM,
)


class TestFairTaskQueuePerformance:
    """Performance benchmarks for fair task queue."""

    @pytest.mark.asyncio
    @pytest.mark.benchmark
    async def test_enqueue_throughput(
        self, in_memory_storage: InMemoryFairQueueStorage, mock_task_adapter: Any
    ) -> None:
        """Benchmark enqueue throughput."""
        queue = FairTaskQueue(
            storage=in_memory_storage,  # type: ignore[arg-type]
            task_adapter=mock_task_adapter,
            group_key_extractor=group_key_extractor,  # type: ignore[arg-type]
        )

        async with queue:
            num_tasks = 1000
            tasks = [create_task_invocation(kwargs={"user_id": f"user_{i % 10}"}) for i in range(num_tasks)]

            start_time = time.time()
            task_ids = await asyncio.gather(
                *[queue.enqueue(task, group_id=f"user_{i % 10}") for i, task in enumerate(tasks)]
            )
            elapsed = time.time() - start_time

            assert len(task_ids) == num_tasks
            throughput = num_tasks / elapsed

            # Should handle at least THROUGHPUT_MINIMUM tasks/sec
            assert throughput > THROUGHPUT_MINIMUM

    @pytest.mark.asyncio
    @pytest.mark.benchmark
    async def test_batch_enqueue_throughput(
        self, in_memory_storage: InMemoryFairQueueStorage, mock_task_adapter: Any
    ) -> None:
        """Benchmark batch enqueue throughput."""
        queue = FairTaskQueue(
            storage=in_memory_storage,  # type: ignore[arg-type]
            task_adapter=mock_task_adapter,
            group_key_extractor=group_key_extractor,  # type: ignore[arg-type]
        )

        async with queue:
            num_tasks = 1000
            tasks = [create_task_invocation(kwargs={"user_id": "user_1"}) for _ in range(num_tasks)]

            start_time = time.time()
            task_ids = await queue.enqueue_batch(tasks, group_id="user_1")
            elapsed = time.time() - start_time

            assert len(task_ids) == num_tasks
            throughput = num_tasks / elapsed

            # Batch should be faster than individual enqueues
            assert throughput > THROUGHPUT_HIGH

    @pytest.mark.asyncio
    @pytest.mark.benchmark
    async def test_dequeue_throughput(
        self, in_memory_storage: InMemoryFairQueueStorage, mock_task_adapter: Any
    ) -> None:
        """Benchmark dequeue throughput."""
        queue = FairTaskQueue(
            storage=in_memory_storage,  # type: ignore[arg-type]
            task_adapter=mock_task_adapter,
            group_key_extractor=group_key_extractor,  # type: ignore[arg-type]
        )

        async with queue:
            num_tasks = 1000
            tasks = [create_task_invocation(kwargs={"user_id": "user_1"}) for _ in range(num_tasks)]

            # Enqueue all tasks
            await queue.enqueue_batch(tasks, group_id="user_1")

            # Dequeue all tasks
            start_time = time.time()
            dequeued = []
            for _ in range(num_tasks):
                task = await in_memory_storage.dequeue_task("user_1")
                if task:
                    dequeued.append(task)
            elapsed = time.time() - start_time

            assert len(dequeued) == num_tasks
            throughput = num_tasks / elapsed

            # Should handle at least THROUGHPUT_MEDIUM tasks/sec
            assert throughput > THROUGHPUT_MEDIUM

    @pytest.mark.asyncio
    @pytest.mark.benchmark
    async def test_concurrent_operations(
        self, in_memory_storage: InMemoryFairQueueStorage, mock_task_adapter: Any
    ) -> None:
        """Benchmark concurrent enqueue/dequeue operations."""
        queue = FairTaskQueue(
            storage=in_memory_storage,  # type: ignore[arg-type]
            task_adapter=mock_task_adapter,
            group_key_extractor=group_key_extractor,  # type: ignore[arg-type]
        )

        async with queue:
            num_tasks = 500
            num_groups = 10

            async def enqueue_tasks(group_id: str) -> list[str]:
                """Enqueue tasks for a group."""
                tasks = [create_task_invocation(kwargs={"user_id": group_id}) for _ in range(num_tasks // num_groups)]
                return await queue.enqueue_batch(tasks, group_id=group_id)

            start_time = time.time()
            task_ids_lists = await asyncio.gather(*[enqueue_tasks(f"group_{i}") for i in range(num_groups)])
            elapsed = time.time() - start_time

            total_tasks = sum(len(ids) for ids in task_ids_lists)
            throughput = total_tasks / elapsed

            assert total_tasks == num_tasks
            assert throughput > THROUGHPUT_MEDIUM

    @pytest.mark.asyncio
    @pytest.mark.benchmark
    async def test_memory_usage(self, in_memory_storage: InMemoryFairQueueStorage, mock_task_adapter: Any) -> None:
        """Test memory usage with large number of tasks."""
        import sys

        queue = FairTaskQueue(
            storage=in_memory_storage,
            task_adapter=mock_task_adapter,
            group_key_extractor=group_key_extractor,
            max_pending_results=1000,
            max_result_events=1000,
        )

        async with queue:
            num_tasks = 10000
            tasks = [create_task_invocation(kwargs={"user_id": f"user_{i % 100}"}) for i in range(num_tasks)]

            # Measure memory before
            # Note: This is approximate, actual memory measurement would require psutil
            await queue.enqueue_batch(tasks[:1000], group_id="user_0")

            # Memory should be bounded by max_pending_results
            scheduler = queue.scheduler
            assert scheduler.pending_results_count <= queue.max_pending_results
            assert scheduler.result_events_count <= queue.max_result_events
