"""Comprehensive integration tests for new TaskInvocation design with real Redis.

This test suite covers all edge cases and scenarios for the new implementation-agnostic
task queue design using real Redis storage.
"""

import asyncio
import contextlib
from collections.abc import AsyncIterator
from datetime import UTC, datetime, timedelta
from typing import Any
from unittest.mock import AsyncMock, MagicMock

import pytest
import pytest_asyncio
from redis.asyncio import Redis
from taskiq import AsyncBroker, InMemoryBroker

from haolib.background.fair import (
    EnqueueOptions,
    FairTaskQueue,
    FairTaskQueueBuilder,
    SerializationStrategy,
    TaskInvocation,
    TaskInvocationGroupKeyExtractor,
    TaskIQAdapter,
    TaskRegistration,
)
from haolib.background.fair.abstract import TaskStatus
from haolib.background.fair.exceptions import (
    FairQueueInvalidTaskError,
    FairQueueQueueFullError,
    FairQueueSerializationError,
)
from haolib.background.fair.storages.redis import RedisFairQueueStorage
from tests.integration.background.fair.constants import (
    BATCH_SIZE_LARGE_2,
    BATCH_SIZE_MEDIUM,
    BATCH_SIZE_SMALL,
    DELAY_MEDIUM,
    DELAY_TINY,
    TASK_COUNT_SMALL,
    TASK_COUNT_VERY_LARGE,
)

# ============================================================================
# Test Task Functions
# ============================================================================


async def simple_task(value: int, user_id: str) -> dict[str, Any]:
    """Simple test task."""
    await asyncio.sleep(DELAY_TINY)
    return {"result": value * 2, "user_id": user_id}


async def complex_task(
    data: dict[str, Any],
    items: list[str],
    user_id: str,
    metadata: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """Complex test task with various argument types."""
    await asyncio.sleep(DELAY_TINY)
    return {
        "processed_data": data,
        "items_count": len(items),
        "user_id": user_id,
        "metadata": metadata or {},
    }


async def task_with_no_kwargs(value: str) -> str:
    """Task with only positional args."""
    await asyncio.sleep(DELAY_TINY)
    return f"processed_{value}"


# ============================================================================
# Fixtures
# ============================================================================


@pytest_asyncio.fixture
async def redis_client() -> AsyncIterator[Redis]:
    """Create a real Redis client for testing."""
    redis = Redis(host="localhost", port=6379, db=15, decode_responses=False)
    # Clean up before test
    await redis.flushdb()
    yield redis
    # Clean up after test
    await redis.flushdb()
    await redis.aclose()


@pytest_asyncio.fixture
async def taskiq_broker() -> AsyncIterator[AsyncBroker]:
    """Create a TaskIQ broker for testing."""
    broker = InMemoryBroker()
    await broker.startup()
    yield broker
    await broker.shutdown()


@pytest.fixture
def group_key_extractor() -> TaskInvocationGroupKeyExtractor:
    """Create group key extractor for TaskInvocation."""
    return TaskInvocationGroupKeyExtractor(
        group_id_key="user_id",
        fallback_keys=["group_id", "tenant_id"],
    )


@pytest_asyncio.fixture
async def fair_queue_redis(
    redis_client: Redis,
    taskiq_broker: AsyncBroker,
    group_key_extractor: TaskInvocationGroupKeyExtractor,
) -> AsyncIterator[FairTaskQueue]:
    """Create a FairTaskQueue with Redis storage."""
    storage = RedisFairQueueStorage(redis_client, key_prefix="test_fair_queue")
    adapter = TaskIQAdapter(broker=taskiq_broker)
    queue = FairTaskQueue(
        storage=storage,  # type: ignore[arg-type]
        task_adapter=adapter,
        group_key_extractor=group_key_extractor,
        scheduler_poll_interval=timedelta(seconds=0.1),
    )
    await queue.__aenter__()
    yield queue
    await queue.__aexit__(None, None, None)


# ============================================================================
# Test Classes
# ============================================================================


class TestNewDesignBasic:
    """Basic tests for new TaskInvocation design."""

    @pytest.mark.asyncio
    async def test_enqueue_with_function(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test enqueueing a task using a function."""
        queue = fair_queue_redis

        # Register task with broker
        taskiq_broker.task(simple_task)

        # Enqueue using function
        task_id = await queue.enqueue(
            simple_task,
            value=10,
            user_id="user_123",
            group_id="user_123",
        )

        assert task_id is not None
        assert isinstance(task_id, str)

        # Verify task status
        status = await queue.get_task_status(task_id, "user_123")
        assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_enqueue_with_task_registration(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test enqueueing using TaskRegistration."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Register task
        registration = await adapter.register_task(simple_task, task_name="simple_task")

        # Enqueue using registration
        task_id = await queue.enqueue(
            registration,
            value=20,
            user_id="user_456",
            group_id="user_456",
        )

        assert task_id is not None

        # Verify task status
        status = await queue.get_task_status(task_id, "user_456")
        assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_enqueue_with_task_invocation(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test enqueueing using pre-created TaskInvocation."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Register and create invocation
        registration = await adapter.register_task(simple_task)
        invocation = adapter.create_invocation(
            registration,
            value=30,
            user_id="user_789",
        )

        # Enqueue invocation
        task_id = await queue.enqueue(
            invocation,
            options=EnqueueOptions(group_id="user_789", priority=5),
        )

        assert task_id is not None
        assert task_id == invocation.task_id

        # Verify task status
        status = await queue.get_task_status(task_id, "user_789")
        assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_auto_group_extraction_from_kwargs(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test automatic group ID extraction from task kwargs."""
        queue = fair_queue_redis

        # Enqueue without explicit group_id (should be extracted from user_id)
        task_id = await queue.enqueue(
            simple_task,
            value=40,
            user_id="user_auto",
        )

        assert task_id is not None

        # Verify group_id was extracted
        task_info = await queue.get_task_info(task_id, "user_auto")
        assert task_info is not None
        assert task_info.group_id == "user_auto"

    @pytest.mark.asyncio
    async def test_auto_group_extraction_from_args(
        self,
        fair_queue_redis: FairTaskQueue,
        group_key_extractor: TaskInvocationGroupKeyExtractor,
    ) -> None:
        """Test automatic group ID extraction from task args."""
        # Create extractor that checks args
        extractor = TaskInvocationGroupKeyExtractor(
            extract_from_args=True,
            args_index=1,  # Second arg is user_id
        )

        # Access storage and adapter for testing
        storage = fair_queue_redis.storage
        adapter = fair_queue_redis.task_adapter
        # Check if they are the expected types
        assert isinstance(storage, RedisFairQueueStorage)
        assert isinstance(adapter, TaskIQAdapter)
        # Get Redis client and broker from public properties
        storage_attr = storage.redis
        adapter_attr = adapter.broker
        assert storage_attr is not None
        assert adapter_attr is not None
        storage = RedisFairQueueStorage(
            storage_attr,
            key_prefix="test_fair_queue_args",
        )
        adapter = TaskIQAdapter(broker=adapter_attr)
        queue = FairTaskQueue(
            storage=storage,  # type: ignore[arg-type]
            task_adapter=adapter,
            group_key_extractor=extractor,
        )
        await queue.__aenter__()

        try:
            # Enqueue with user_id as second positional arg
            # Note: simple_task signature is (value: int, user_id: str)
            task_id = await queue.enqueue(
                simple_task,
                50,  # value as first positional arg
                "user_args",  # user_id as second positional arg
            )

            assert task_id is not None
            task_info = await queue.get_task_info(task_id, "user_args")
            assert task_info is not None
            assert task_info.group_id == "user_args"
        finally:
            await queue.__aexit__(None, None, None)


class TestSerializationStrategies:
    """Tests for different serialization strategies."""

    @pytest.mark.asyncio
    async def test_json_serialization(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test JSON serialization strategy."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Register with JSON strategy
        registration = await adapter.register_task(
            simple_task,
            default_serialization_strategy=SerializationStrategy.JSON,
        )

        # Create invocation with JSON
        invocation = adapter.create_invocation(
            registration,
            value=100,
            user_id="user_json",
            serialization_strategy=SerializationStrategy.JSON,
        )

        task_id = await queue.enqueue(invocation, group_id="user_json")
        assert task_id is not None

        # Verify task can be retrieved
        task_info = await queue.get_task_info(task_id, "user_json")
        assert task_info is not None

    @pytest.mark.asyncio
    async def test_pickle_serialization(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test PICKLE serialization strategy."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Register with PICKLE strategy
        registration = await adapter.register_task(
            complex_task,
            default_serialization_strategy=SerializationStrategy.PICKLE,
        )

        # Create invocation with complex data (requires pickle)
        complex_data = {
            "nested": {"key": "value", "number": 42},
            "list": [1, 2, 3],
        }
        invocation = adapter.create_invocation(
            registration,
            data=complex_data,
            items=["a", "b", "c"],
            user_id="user_pickle",
            serialization_strategy=SerializationStrategy.PICKLE,
        )

        task_id = await queue.enqueue(invocation, group_id="user_pickle")
        assert task_id is not None

        # Verify task can be retrieved
        task_info = await queue.get_task_info(task_id, "user_pickle")
        assert task_info is not None

    @pytest.mark.asyncio
    async def test_serialization_failure_handling(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test handling of serialization failures."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Create a task with non-serializable data (function object)
        registration = await adapter.register_task(
            simple_task,
            default_serialization_strategy=SerializationStrategy.JSON,
        )

        # JSON cannot serialize functions
        invocation = adapter.create_invocation(
            registration,
            value=200,
            user_id="user_fail",
            serialization_strategy=SerializationStrategy.JSON,
        )

        # This should work (value is JSON-serializable)
        task_id = await queue.enqueue(invocation, group_id="user_fail")
        assert task_id is not None

        # But if we try to serialize a function with JSON, it should fail
        # (This is tested implicitly - if we pass a function, it will fail during serialization)


class TestEdgeCases:
    """Tests for edge cases and error scenarios."""

    @pytest.mark.asyncio
    async def test_empty_args_kwargs(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test enqueueing task with empty args and kwargs."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        async def no_args_task() -> str:
            return "done"

        registration = await adapter.register_task(no_args_task)
        invocation = adapter.create_invocation(registration)

        task_id = await queue.enqueue(invocation, group_id="empty_task")
        assert task_id is not None

    @pytest.mark.asyncio
    async def test_missing_group_id(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test error when group_id cannot be extracted."""
        queue = fair_queue_redis

        async def task_without_user_id(value: int) -> int:
            return value * 2

        # Should raise error - no user_id in kwargs and no group_id provided
        with pytest.raises(FairQueueInvalidTaskError) as exc_info:
            await queue.enqueue(
                task_without_user_id,
                value=300,
            )

        assert "group_id_extraction_failed" in str(
            exc_info.value.validation_error
        ) or "Failed to extract group_id" in str(exc_info.value.detail)

    @pytest.mark.asyncio
    async def test_invalid_task_type(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test error when invalid task type is provided."""
        queue = fair_queue_redis

        # Pass invalid type
        with pytest.raises(FairQueueInvalidTaskError) as exc_info:
            await queue.enqueue(
                12345,  # type: ignore[arg-type]
                group_id="test",
            )

        assert "invalid_task_type" in str(exc_info.value.validation_error)

    @pytest.mark.asyncio
    async def test_priority_validation(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test priority range validation."""
        queue = fair_queue_redis

        # Test priority too high
        with pytest.raises(FairQueueInvalidTaskError) as exc_info:
            await queue.enqueue(
                simple_task,
                value=400,
                user_id="user_priority",
                options=EnqueueOptions(
                    group_id="user_priority",
                    priority=10000,  # Exceeds max_priority (default 1000)
                ),
            )

        assert "priority out of range" in str(exc_info.value.validation_error)

    @pytest.mark.asyncio
    async def test_group_id_length_validation(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test group ID length validation."""
        queue = fair_queue_redis

        # Create very long group_id
        long_group_id = "a" * 300  # Exceeds default max (255)

        with pytest.raises(FairQueueInvalidTaskError) as exc_info:
            await queue.enqueue(
                simple_task,
                value=500,
                user_id="user_long",
                options=EnqueueOptions(group_id=long_group_id),
            )

        assert "group_id too long" in str(exc_info.value.validation_error)

    @pytest.mark.asyncio
    async def test_task_id_generation(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test that task IDs are generated when not provided."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        registration = await adapter.register_task(simple_task)

        # Create invocation without task_id (should be generated)
        invocation1 = adapter.create_invocation(
            registration,
            value=600,
            user_id="user_gen1",
        )

        invocation2 = adapter.create_invocation(
            registration,
            value=700,
            user_id="user_gen2",
        )

        # Task IDs should be different
        assert invocation1.task_id != invocation2.task_id
        assert len(invocation1.task_id) > 0
        assert len(invocation2.task_id) > 0

    @pytest.mark.asyncio
    async def test_custom_task_id(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test using custom task ID."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        registration = await adapter.register_task(simple_task)

        # Create invocation with custom task_id
        custom_id = "custom_task_12345"
        invocation = adapter.create_invocation(
            registration,
            value=800,
            user_id="user_custom",
            task_id=custom_id,
        )

        assert invocation.task_id == custom_id

        task_id = await queue.enqueue(invocation, group_id="user_custom")
        assert task_id == custom_id


class TestBatchOperations:
    """Tests for batch enqueue operations."""

    @pytest.mark.asyncio
    async def test_batch_enqueue_functions(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test batch enqueueing multiple functions."""
        queue = fair_queue_redis

        tasks = []
        for i in range(TASK_COUNT_VERY_LARGE):
            tasks.append(
                (
                    simple_task,
                    EnqueueOptions(
                        group_id=f"user_{i % 3}",  # Distribute across 3 groups
                        priority=i,
                    ),
                )
            )

        task_ids = await queue.enqueue_batch(tasks)

        assert len(task_ids) == TASK_COUNT_VERY_LARGE
        assert all(tid is not None for tid in task_ids)

        # Verify all tasks are enqueued
        for i, task_id in enumerate(task_ids):
            group_id = f"user_{i % 3}"
            status = await queue.get_task_status(task_id, group_id)
            assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_batch_enqueue_invocations(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test batch enqueueing pre-created invocations."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        registration = await adapter.register_task(simple_task)

        tasks = []
        for i in range(BATCH_SIZE_MEDIUM):
            invocation = adapter.create_invocation(
                registration,
                value=i * 100,
                user_id=f"user_batch_{i}",
            )
            tasks.append(
                (
                    invocation,
                    EnqueueOptions(group_id=f"user_batch_{i}", priority=i),
                )
            )

        task_ids = await queue.enqueue_batch(tasks)

        assert len(task_ids) == BATCH_SIZE_MEDIUM
        assert all(tid is not None for tid in task_ids)


class TestTaskExecution:
    """Tests for task execution and dispatch."""

    @pytest.mark.asyncio
    async def test_task_dispatch_and_execution(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test that tasks are properly dispatched and executed."""
        queue = fair_queue_redis

        # Register task with broker
        taskiq_broker.task(simple_task)

        # Enqueue task
        task_id = await queue.enqueue(
            simple_task,
            value=1000,
            user_id="user_exec",
            group_id="user_exec",
        )

        # Wait for scheduler to process
        await asyncio.sleep(0.5)

        # Check task status (should be processing or completed)
        status = await queue.get_task_status(task_id, "user_exec")
        assert status in (TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.COMPLETED)

    @pytest.mark.asyncio
    async def test_task_with_dependencies(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test task with dependencies."""
        queue = fair_queue_redis

        # Enqueue first task
        task_id1 = await queue.enqueue(
            simple_task,
            value=1100,
            user_id="user_dep1",
            group_id="user_dep1",
        )

        # Enqueue second task that depends on first
        task_id2 = await queue.enqueue(
            simple_task,
            value=1200,
            user_id="user_dep2",
            group_id="user_dep2",
            options=EnqueueOptions(depends_on=[task_id1]),
        )

        assert task_id2 is not None

        # Second task should be pending until first completes
        status2 = await queue.get_task_status(task_id2, "user_dep2")
        assert status2 == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_task_with_priority(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test that tasks with higher priority are processed first."""
        queue = fair_queue_redis

        # Enqueue low priority task first
        await queue.enqueue(
            simple_task,
            value=1300,
            user_id="user_prio",
            group_id="user_prio",
            options=EnqueueOptions(priority=0),
        )

        # Enqueue high priority task second
        await queue.enqueue(
            simple_task,
            value=1400,
            user_id="user_prio",
            group_id="user_prio",
            options=EnqueueOptions(priority=10),
        )

        # Wait a bit
        await asyncio.sleep(0.2)

        # High priority task should be dequeued first (if scheduler processed)
        # This is verified by checking which task was processed
        # (In practice, high priority should be dequeued first)


class TestErrorRecovery:
    """Tests for error recovery and resilience."""

    @pytest.mark.asyncio
    async def test_serialization_error_recovery(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test handling of serialization errors."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Create invocation with JSON strategy (should work for simple data)
        registration = await adapter.register_task(simple_task)
        invocation = adapter.create_invocation(
            registration,
            value=1500,
            user_id="user_ser",
            serialization_strategy=SerializationStrategy.JSON,
        )

        # This should succeed
        task_id = await queue.enqueue(invocation, group_id="user_ser")
        assert task_id is not None

    @pytest.mark.asyncio
    async def test_queue_full_with_backpressure(
        self,
        redis_client: Redis,
        taskiq_broker: AsyncBroker,
        group_key_extractor: TaskInvocationGroupKeyExtractor,
    ) -> None:
        """Test queue full scenario with backpressure."""
        # Create completely isolated queue with small limit
        storage = RedisFairQueueStorage(
            redis_client,
            key_prefix="test_backpressure",
        )
        adapter = TaskIQAdapter(broker=taskiq_broker)
        queue = FairTaskQueue(
            storage=storage,  # type: ignore[arg-type]
            task_adapter=adapter,
            group_key_extractor=group_key_extractor,
            max_queue_size_per_group=2,  # Very small limit
        )
        # Don't start the scheduler at all - use __aenter__ but stop immediately
        await queue.__aenter__()

        # Stop scheduler immediately to prevent processing
        await queue.stop_scheduler(graceful=False)
        # Small delay to ensure scheduler is fully stopped
        await asyncio.sleep(0.2)

        try:
            # Fill up the queue (scheduler is stopped, so tasks won't be processed)
            await queue.enqueue(simple_task, value=1, user_id="user_bp", group_id="user_bp")
            await queue.enqueue(simple_task, value=2, user_id="user_bp", group_id="user_bp")

            # Verify queue is full
            queue_length = await queue.get_queue_length("user_bp")
            assert queue_length >= TASK_COUNT_SMALL, f"Queue should be full, but length is {queue_length}"

            # Try to enqueue with backpressure (should wait and then fail)
            start_time = datetime.now(UTC)
            with pytest.raises(FairQueueQueueFullError):
                await queue.enqueue(
                    simple_task,
                    value=3,
                    user_id="user_bp",
                    group_id="user_bp",
                    options=EnqueueOptions(
                        backpressure=True,
                        backpressure_max_wait=timedelta(seconds=1),
                        backpressure_retry_interval=timedelta(seconds=0.1),
                    ),
                )
            elapsed = (datetime.now(UTC) - start_time).total_seconds()
            # Should have waited at least some time (backpressure)
            assert elapsed >= DELAY_MEDIUM, f"Backpressure should wait, but only waited {elapsed}s"
        finally:
            await queue.__aexit__(None, None, None)

    @pytest.mark.asyncio
    async def test_idempotency(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test idempotency key functionality."""
        queue = fair_queue_redis

        idempotency_key = "idempotent_task_123"

        # Enqueue first time
        task_id1 = await queue.enqueue(
            simple_task,
            value=1600,
            user_id="user_idem",
            group_id="user_idem",
            options=EnqueueOptions(idempotency_key=idempotency_key),
        )

        # Enqueue second time with same idempotency key
        task_id2 = await queue.enqueue(
            simple_task,
            value=1700,
            user_id="user_idem",
            group_id="user_idem",
            options=EnqueueOptions(idempotency_key=idempotency_key),
        )

        # Should return same task_id
        assert task_id1 == task_id2


class TestComplexScenarios:
    """Tests for complex real-world scenarios."""

    @pytest.mark.asyncio
    async def test_multiple_groups_fairness(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test fair distribution across multiple groups."""
        queue = fair_queue_redis

        # Enqueue tasks for multiple groups
        for group_num in range(5):
            for task_num in range(3):
                await queue.enqueue(
                    simple_task,
                    value=group_num * 100 + task_num,
                    user_id=f"group_{group_num}",
                    group_id=f"group_{group_num}",
                )

        # Verify all tasks are enqueued
        for group_num in range(5):
            group_id = f"group_{group_num}"
            length = await queue.storage.get_group_queue_length(group_id)
            assert length == BATCH_SIZE_SMALL

    @pytest.mark.asyncio
    async def test_rate_limiting(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test rate limiting per group."""
        queue = fair_queue_redis

        from haolib.background.fair.abstract import RateLimit

        # Enqueue with rate limit
        task_id = await queue.enqueue(
            simple_task,
            value=1800,
            user_id="user_rate",
            group_id="user_rate",
            options=EnqueueOptions(
                rate_limit=RateLimit(max_tasks=2, per_time=timedelta(seconds=1)),
            ),
        )

        assert task_id is not None

    @pytest.mark.asyncio
    async def test_scheduled_tasks(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test scheduled (delayed) task execution."""
        queue = fair_queue_redis

        # Schedule task for future
        future_time = datetime.now(UTC) + timedelta(seconds=10)

        task_id = await queue.enqueue(
            simple_task,
            value=1900,
            user_id="user_scheduled",
            group_id="user_scheduled",
            options=EnqueueOptions(scheduled_at=future_time),
        )

        assert task_id is not None

        # Task should be pending (not yet scheduled)
        status = await queue.get_task_status(task_id, "user_scheduled")
        assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_expiring_tasks(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test task expiration."""
        queue = fair_queue_redis

        # Create task that expires very soon (before scheduler can process it)
        # Scheduler polls every 0.1 seconds, so expire in 0.05 seconds to ensure
        # expiration is detected before processing
        expire_time = datetime.now(UTC) + timedelta(seconds=0.05)

        task_id = await queue.enqueue(
            simple_task,
            value=2000,
            user_id="user_expire",
            group_id="user_expire",
            options=EnqueueOptions(expires_at=expire_time),
        )

        assert task_id is not None

        # Wait for expiration (longer than expiration time to ensure it's detected)
        await asyncio.sleep(0.15)

        # Task should be expired (or pending if scheduler hasn't tried to dequeue yet)
        # The deterministic approach: expiration time is shorter than poll interval,
        # so scheduler will detect expiration when it tries to dequeue
        status = await queue.get_task_status(task_id, "user_expire")
        assert status in (TaskStatus.PENDING, TaskStatus.EXPIRED)


class TestRegistrationCaching:
    """Tests for task registration caching."""

    @pytest.mark.asyncio
    async def test_registration_idempotency(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test that registering the same task multiple times returns same registration."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Register same task twice
        reg1 = await adapter.register_task(simple_task, task_name="test_task")
        reg2 = await adapter.register_task(simple_task, task_name="test_task")

        # Should return equivalent registrations (same task_name)
        assert reg1.task_name == reg2.task_name

    @pytest.mark.asyncio
    async def test_different_task_names(
        self,
        fair_queue_redis: FairTaskQueue,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test registering same function with different names."""
        queue = fair_queue_redis
        adapter = queue.task_adapter

        # Register with different names
        reg1 = await adapter.register_task(simple_task, task_name="task_v1")
        reg2 = await adapter.register_task(simple_task, task_name="task_v2")

        # Should have different task names
        assert reg1.task_name != reg2.task_name
        assert reg1.task_name == "task_v1"
        assert reg2.task_name == "task_v2"


class TestIntegrationWithBuilder:
    """Tests for integration with FairTaskQueueBuilder."""

    @pytest.mark.asyncio
    async def test_builder_with_new_design(
        self,
        redis_client: Redis,
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test using builder with new design."""
        queue = (
            FairTaskQueueBuilder()
            .with_redis_storage(redis_client, key_prefix="test_builder")
            .with_taskiq_adapter(taskiq_broker)
            .with_group_key_extractor(TaskInvocationGroupKeyExtractor(group_id_key="user_id"))
            .build()
        )

        await queue.__aenter__()

        try:
            # Enqueue task
            task_id = await queue.enqueue(
                simple_task,
                value=2100,
                user_id="user_builder",
            )

            assert task_id is not None
        finally:
            await queue.__aexit__(None, None, None)


class TestConcurrentOperations:
    """Tests for concurrent operations and race conditions."""

    @pytest.mark.asyncio
    async def test_concurrent_enqueue(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test concurrent enqueue operations."""
        queue = fair_queue_redis

        # Enqueue multiple tasks concurrently
        tasks = [
            queue.enqueue(
                simple_task,
                value=i,
                user_id=f"user_concurrent_{i % 3}",
                group_id=f"user_concurrent_{i % 3}",
            )
            for i in range(BATCH_SIZE_LARGE_2)
        ]

        task_ids = await asyncio.gather(*tasks)

        assert len(task_ids) == BATCH_SIZE_LARGE_2
        assert all(tid is not None for tid in task_ids)

        # Verify all tasks are enqueued
        for i, task_id in enumerate(task_ids):
            group_id = f"user_concurrent_{i % 3}"
            status = await queue.get_task_status(task_id, group_id)
            assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_concurrent_batch_enqueue(
        self,
        fair_queue_redis: FairTaskQueue,
    ) -> None:
        """Test concurrent batch enqueue operations."""
        queue = fair_queue_redis

        # Create multiple batches
        batches = []
        for batch_num in range(5):
            batch_tasks = [
                (
                    simple_task,
                    EnqueueOptions(
                        group_id=f"batch_{batch_num}",
                        priority=i,
                    ),
                )
                for i in range(TASK_COUNT_VERY_LARGE)
            ]
            batches.append(queue.enqueue_batch(batch_tasks))

        results = await asyncio.gather(*batches)

        # Verify all batches succeeded
        for batch_result in results:
            assert len(batch_result) == TASK_COUNT_VERY_LARGE
            assert all(tid is not None for tid in batch_result)


class TestStoragePersistence:
    """Tests for storage persistence and recovery."""

    @pytest.mark.asyncio
    async def test_task_persistence_across_restarts(
        self,
        redis_client: Redis,
        taskiq_broker: AsyncBroker,
        group_key_extractor: TaskInvocationGroupKeyExtractor,
    ) -> None:
        """Test that tasks persist across queue restarts."""
        storage = RedisFairQueueStorage(redis_client, key_prefix="test_persistence")
        adapter = TaskIQAdapter(broker=taskiq_broker)

        # Create and start queue
        queue1 = FairTaskQueue(
            storage=storage,  # type: ignore[arg-type]
            task_adapter=adapter,
            group_key_extractor=group_key_extractor,
        )
        await queue1.__aenter__()

        try:
            # Enqueue task
            task_id = await queue1.enqueue(
                simple_task,
                value=2200,
                user_id="user_persist",
                group_id="user_persist",
            )

            # Stop queue
            await queue1.__aexit__(None, None, None)

            # Create new queue instance (simulating restart)
            queue2 = FairTaskQueue(
                storage=storage,  # type: ignore[arg-type]
                task_adapter=adapter,
                group_key_extractor=group_key_extractor,
            )
            await queue2.__aenter__()

            try:
                # Task should still be there
                status = await queue2.get_task_status(task_id, "user_persist")
                assert status == TaskStatus.PENDING

                # Task info should be retrievable
                task_info = await queue2.get_task_info(task_id, "user_persist")
                assert task_info is not None
                assert task_info.task_id == task_id
            finally:
                await queue2.__aexit__(None, None, None)
        finally:
            # Cleanup in case of error
            with contextlib.suppress(Exception):
                await queue1.__aexit__(None, None, None)
