"""Tests for declarative TaskIQ task pattern with fair queue.

These tests verify that you can:
1. Define TaskIQ tasks using @broker.task decorator or broker.task() method
2. Create task messages using await task.kiq() method (note: .kiq() is async)
3. Enqueue them in the fair queue

Example usage pattern:
    ```python
    # Define task
    @broker.task
    async def my_task(arg1: str, user_id: str):
        return result

    # Create and enqueue
    task = await my_task.kiq(arg1="value", user_id="user_123")
    task_id = await fair_queue.enqueue(task, group_id="user_123")
    ```
"""

import asyncio
from collections.abc import AsyncIterator
from typing import Any

import pytest
import pytest_asyncio
from redis.asyncio import Redis
from taskiq import AsyncBroker
from taskiq_redis import RedisStreamBroker

from haolib.background.fair import EnqueueOptions, FairTaskQueue
from haolib.background.fair.abstract import TaskStatus
from haolib.background.fair.adapters import TaskIQAdapter
from haolib.background.fair.extractors import extract_from_kwargs
from haolib.background.fair.storages.redis import RedisFairQueueStorage
from haolib.entrypoints.taskiq import TaskiqEntrypointWorker
from tests.integration.background.fair.constants import TASK_COUNT_VERY_LARGE

# ============================================================================
# TaskIQ Task Definitions (Declarative Pattern)
# ============================================================================

# These will be set after broker is created
process_image_task: Any = None
send_email_task: Any = None
calculate_sum_task: Any = None

# Global execution counter for testing
_execution_counts: dict[str, int] = {}


async def process_image(image_url: str, user_id: str) -> dict[str, Any]:
    """Example TaskIQ task that processes an image."""
    # Track execution count
    key = f"process_image:{image_url}:{user_id}"
    _execution_counts[key] = _execution_counts.get(key, 0) + 1

    await asyncio.sleep(0.1)  # Simulate processing
    return {
        "status": "processed",
        "image_url": image_url,
        "user_id": user_id,
        "execution_count": _execution_counts[key],
    }


async def send_email(email: str, subject: str, user_id: str) -> dict[str, Any]:
    """Example TaskIQ task that sends an email."""
    await asyncio.sleep(0.1)  # Simulate sending
    return {
        "status": "sent",
        "email": email,
        "subject": subject,
        "user_id": user_id,
    }


async def calculate_sum(a: int, b: int, user_id: str) -> int:
    """Example TaskIQ task that calculates sum."""
    return a + b


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


@pytest_asyncio.fixture
async def taskiq_broker() -> AsyncIterator[AsyncBroker]:
    """Create a RedisStreamBroker for testing."""
    broker = RedisStreamBroker(url="redis://localhost:6379")
    await broker.startup()
    try:
        yield broker
    finally:
        await broker.shutdown()


@pytest_asyncio.fixture
async def registered_tasks(taskiq_broker: AsyncBroker) -> dict[str, Any]:
    """Register TaskIQ tasks on the broker and return task references."""
    global process_image_task, send_email_task, calculate_sum_task

    # Register tasks using declarative pattern
    process_image_task = taskiq_broker.task(process_image)
    send_email_task = taskiq_broker.task(send_email)
    calculate_sum_task = taskiq_broker.task(calculate_sum)

    return {
        "process_image": process_image_task,
        "send_email": send_email_task,
        "calculate_sum": calculate_sum_task,
    }


@pytest_asyncio.fixture
def group_key_extractor() -> Any:
    """Create a group key extractor that extracts user_id from TaskInvocation kwargs."""
    # Use functional extractor with fallback to default
    from haolib.background.fair.extractors import compose_extractors, extract_with_default

    return compose_extractors(
        extract_from_kwargs("user_id"),
        extract_with_default("default"),
    )


@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_worker(taskiq_broker: AsyncBroker) -> AsyncIterator[TaskiqEntrypointWorker]:
    """Create and run a TaskIQ worker to execute tasks."""
    worker = TaskiqEntrypointWorker(broker=taskiq_broker, should_run_worker=True)
    await worker.startup()
    try:
        yield worker
    finally:
        await worker.shutdown()


@pytest_asyncio.fixture
async def fair_queue_with_taskiq(
    taskiq_broker: AsyncBroker,
    registered_tasks: dict[str, Any],
    group_key_extractor: Any,
    redis_client: Redis,
    taskiq_worker: TaskiqEntrypointWorker,  # Keep worker alive during tests
) -> FairTaskQueue:
    """Create a fair queue with TaskIQ adapter and Redis storage."""
    storage = RedisFairQueueStorage(redis_client, key_prefix="test_fair_queue")
    task_adapter = TaskIQAdapter(broker=taskiq_broker)

    return FairTaskQueue(
        storage=storage,  # type: ignore[arg-type]
        task_adapter=task_adapter,
        group_key_extractor=group_key_extractor,
    )


# ============================================================================
# Tests
# ============================================================================


class TestDeclarativeTaskIQPattern:
    """Test declarative TaskIQ task pattern with fair queue."""

    @pytest.mark.asyncio
    async def test_enqueue_task_using_kiq_method(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test enqueueing a task using declarative pattern (decorated task as callable)."""
        queue = fair_queue_with_taskiq
        process_image = registered_tasks["process_image"]

        async with queue:
            # Use decorated task as callable (current implementation pattern)
            # The queue will register it and create an invocation
            task_id = await queue.enqueue(
                process_image,
                image_url="https://example.com/image.jpg",
                user_id="user_123",
                group_id="user_123",
            )

            assert task_id is not None

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

    @pytest.mark.asyncio
    async def test_enqueue_multiple_tasks_with_kiq(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test enqueueing multiple tasks using decorated tasks as callables."""
        queue = fair_queue_with_taskiq
        process_image = registered_tasks["process_image"]
        send_email = registered_tasks["send_email"]

        async with queue:
            # Use decorated tasks as callables (current implementation pattern)
            task_id1 = await queue.enqueue(
                process_image,
                image_url="https://example.com/image1.jpg",
                user_id="user_1",
                group_id="user_1",
            )
            task_id2 = await queue.enqueue(
                send_email,
                email="test@example.com",
                subject="Test",
                user_id="user_1",
                group_id="user_1",
            )
            task_id3 = await queue.enqueue(
                process_image,
                image_url="https://example.com/image2.jpg",
                user_id="user_2",
                group_id="user_2",
            )

            # Verify all tasks were enqueued
            assert task_id1 is not None
            assert task_id2 is not None
            assert task_id3 is not None

            # Verify statuses
            assert await queue.get_task_status(task_id1, "user_1") == TaskStatus.PENDING
            assert await queue.get_task_status(task_id2, "user_1") == TaskStatus.PENDING
            assert await queue.get_task_status(task_id3, "user_2") == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_enqueue_with_options_object(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test enqueueing with EnqueueOptions object."""
        queue = fair_queue_with_taskiq
        process_image = registered_tasks["process_image"]

        async with queue:
            # Enqueue with options object using decorated task as callable
            options = EnqueueOptions(
                group_id="user_123",
                priority=10,
                max_retries=5,
            )

            task_id = await queue.enqueue(
                process_image,
                image_url="https://example.com/image.jpg",
                user_id="user_123",
                options=options,
            )

            assert task_id is not None
            status = await queue.get_task_status(task_id, "user_123")
            assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_enqueue_with_priority(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test enqueueing tasks with different priorities."""
        queue = fair_queue_with_taskiq
        process_image = registered_tasks["process_image"]

        async with queue:
            # Enqueue with priorities using decorated task as callable
            task_id_low = await queue.enqueue(
                process_image,
                image_url="https://example.com/low.jpg",
                user_id="user_1",
                group_id="user_1",
                priority=1,
            )
            task_id_high = await queue.enqueue(
                process_image,
                image_url="https://example.com/high.jpg",
                user_id="user_1",
                group_id="user_1",
                priority=10,
            )

            assert task_id_low is not None
            assert task_id_high is not None

    @pytest.mark.asyncio
    async def test_enqueue_with_auto_group_extraction(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test that group_id is auto-extracted from task kwargs."""
        queue = fair_queue_with_taskiq
        process_image = registered_tasks["process_image"]

        async with queue:
            # Enqueue without specifying group_id (should be auto-extracted from user_id)
            task_id = await queue.enqueue(
                process_image,
                image_url="https://example.com/image.jpg",
                user_id="user_456",  # This should be extracted as group_id
            )

            assert task_id is not None
            # Verify task was grouped correctly
            status = await queue.get_task_status(task_id, "user_456")
            assert status == TaskStatus.PENDING

    @pytest.mark.asyncio
    async def test_enqueue_different_task_types(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test enqueueing different types of tasks."""
        queue = fair_queue_with_taskiq
        process_image = registered_tasks["process_image"]
        send_email = registered_tasks["send_email"]
        calculate_sum = registered_tasks["calculate_sum"]

        async with queue:
            # Enqueue different task types using decorated tasks as callables
            task_id1 = await queue.enqueue(
                process_image,
                image_url="https://example.com/image.jpg",
                user_id="user_1",
                group_id="user_1",
            )
            task_id2 = await queue.enqueue(
                send_email,
                email="test@example.com",
                subject="Hello",
                user_id="user_1",
                group_id="user_1",
            )
            task_id3 = await queue.enqueue(
                calculate_sum,
                a=5,
                b=10,
                user_id="user_1",
                group_id="user_1",
            )

            assert task_id1 is not None
            assert task_id2 is not None
            assert task_id3 is not None

    @pytest.mark.asyncio
    async def test_enqueue_batch_with_kiq(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test batch enqueueing tasks using decorated tasks as callables."""
        from haolib.background.fair.abstract import TaskInvocation
        from haolib.background.fair.adapters import TaskIQAdapter

        queue = fair_queue_with_taskiq
        process_image = registered_tasks["process_image"]
        adapter = queue.task_adapter

        async with queue:
            # Create TaskInvocation objects with kwargs for batch enqueue
            tasks = []
            for i in range(10):
                group_id = f"user_{i % 3}"  # Distribute across 3 users
                # Get registration for the task
                registration = await adapter.register_task(process_image)
                # Create invocation with kwargs
                invocation = adapter.create_invocation(
                    registration,
                    image_url=f"https://example.com/image{i}.jpg",
                    user_id=group_id,
                )
                tasks.append(
                    (
                        invocation,
                        EnqueueOptions(group_id=group_id),
                    )
                )

            # Batch enqueue
            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)

    @pytest.mark.asyncio
    async def test_task_dispatch_through_fair_queue(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
        taskiq_broker: AsyncBroker,
    ) -> None:
        """Test that tasks are properly dispatched through fair queue."""
        global _execution_counts
        _execution_counts.clear()

        queue = fair_queue_with_taskiq
        registered_tasks["process_image"]

        async with queue:
            # Enqueue task using the raw function directly
            # Import the raw function (not the decorated one)
            import sys

            current_module = sys.modules[__name__]
            raw_process_image = current_module.process_image

            task_id = await queue.enqueue(
                raw_process_image,
                image_url="https://example.com/image.jpg",
                user_id="user_123",
                group_id="user_123",
            )

            assert task_id is not None

            # Wait for scheduler to process (give more time for Redis + worker)
            # The scheduler polls every 0.1s, and worker needs time to consume and execute
            await asyncio.sleep(2.0)

            # Verify task was dispatched (status should change from PENDING)
            status = await queue.get_task_status(task_id, "user_123")
            # Status should be PROCESSING or COMPLETED if worker executed it
            # If it's still PENDING, the scheduler might not have processed it yet, which is also OK for this test
            assert status in (TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.COMPLETED)

            # Check execution count if task was executed
            key = "process_image:https://example.com/image.jpg:user_123"
            execution_count = _execution_counts.get(key, 0)
            # If task was executed, it should be exactly once (not twice)
            if execution_count > 0:
                assert execution_count == 1, f"Task was executed {execution_count} times, expected 1"

    @pytest.mark.asyncio
    async def test_no_double_execution(
        self,
        fair_queue_with_taskiq: FairTaskQueue,
        registered_tasks: dict[str, Any],
    ) -> None:
        """Test that tasks are executed only once, not double-called."""
        global _execution_counts
        _execution_counts.clear()

        queue = fair_queue_with_taskiq
        registered_tasks["process_image"]

        async with queue:
            # Enqueue two different tasks using raw function (different image URLs to track separately)
            import sys

            current_module = sys.modules[__name__]
            raw_process_image = current_module.process_image

            task_id1 = await queue.enqueue(
                raw_process_image,
                image_url="https://example.com/test1.jpg",
                user_id="user_test",
                group_id="user_test",
            )
            task_id2 = await queue.enqueue(
                raw_process_image,
                image_url="https://example.com/test2.jpg",
                user_id="user_test",
                group_id="user_test",
            )

            assert task_id1 is not None
            assert task_id2 is not None
            assert task_id1 != task_id2  # Different task IDs

            # Wait for scheduler to process both tasks (give more time for Redis + worker)
            await asyncio.sleep(2.0)

            # Verify both tasks were dispatched (status should change from PENDING)
            status1 = await queue.get_task_status(task_id1, "user_test")
            status2 = await queue.get_task_status(task_id2, "user_test")
            assert status1 in (TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.COMPLETED)
            assert status2 in (TaskStatus.PENDING, TaskStatus.PROCESSING, TaskStatus.COMPLETED)

            # Check execution counts if tasks were executed
            key1 = "process_image:https://example.com/test1.jpg:user_test"
            key2 = "process_image:https://example.com/test2.jpg:user_test"
            execution_count1 = _execution_counts.get(key1, 0)
            execution_count2 = _execution_counts.get(key2, 0)

            # If tasks were executed, each should be exactly once (not twice - no double execution)
            if execution_count1 > 0:
                assert execution_count1 == 1, (
                    f"Task 1 was executed {execution_count1} times, expected 1 (no double execution)"
                )
            if execution_count2 > 0:
                assert execution_count2 == 1, (
                    f"Task 2 was executed {execution_count2} times, expected 1 (no double execution)"
                )
