"""Conftest for fair task queue integration tests."""

import asyncio
import uuid
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from typing import Any
from unittest.mock import AsyncMock

import pytest_asyncio
from redis.asyncio import Redis
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from taskiq import AsyncBroker

from haolib.background.fair.abstract import GroupQueueState, RateLimit, SerializationStrategy, TaskInvocation
from haolib.background.fair.adapters import TaskIQAdapter
from haolib.background.fair.extractors import (
    compose_extractors,
    extract_from_kwargs,
    extract_with_default,
)
from haolib.background.fair.storages.memory import InMemoryFairQueueStorage
from haolib.background.fair.storages.redis import RedisFairQueueStorage
from haolib.background.fair.storages.sqlalchemy import (
    DatabaseFairQueueStorage,
    FairTaskModel,
    GlobalQueueSizeModel,
    GroupQueueStateModel,
    RateLimitModel,
    TaskResultModel,
)
from tests.integration.conftest import MockAppConfig

# Test constants
GROUP_ID_1 = "group_1"
GROUP_ID_2 = "group_2"
GROUP_ID_3 = "group_3"
TASK_ID_1 = "task_1"
TASK_ID_2 = "task_2"
TASK_ID_3 = "task_3"


def create_group_state(
    group_id: str,
    queue_length: int = 0,
    virtual_finish_time: float = 0.0,
    weight: float = 1.0,
    last_processed_at: datetime | None = None,
    rate_limit: RateLimit | None = None,
) -> GroupQueueState:
    """Create a test group state.

    Args:
        group_id: Group ID
        queue_length: Queue length
        virtual_finish_time: Virtual finish time
        weight: Group weight
        last_processed_at: Last processed timestamp
        rate_limit: Optional rate limit

    Returns:
        GroupQueueState instance

    """
    return GroupQueueState(
        group_id=group_id,
        queue_length=queue_length,
        virtual_finish_time=virtual_finish_time,
        weight=weight,
        last_processed_at=last_processed_at or datetime.now(UTC),
        rate_limit=rate_limit,
    )


async def simple_test_task(**kwargs: Any) -> dict[str, Any]:
    """Simple test task function for testing.

    Args:
        **kwargs: Task arguments

    Returns:
        Task result

    """
    await asyncio.sleep(0.01)  # Simulate some work
    return {"status": "completed", **kwargs}


def create_task_invocation(
    task_id: str | None = None,
    task_name: str = "test_task",
    kwargs: dict[str, Any] | None = None,
    args: tuple[Any, ...] | None = None,
) -> TaskInvocation:
    """Create a test TaskInvocation.

    Args:
        task_id: Optional task ID
        task_name: Task name
        kwargs: Task kwargs
        args: Task args

    Returns:
        TaskInvocation instance

    """

    return TaskInvocation(
        task_name=task_name,
        args=args or (),
        kwargs=kwargs or {},
        task_id=task_id or uuid.uuid4().hex,
        serialization_strategy=SerializationStrategy.JSON,
        invocation_metadata={},
    )


@pytest_asyncio.fixture
async def mock_broker() -> AsyncMock:
    """Create a mock TaskIQ broker."""
    broker = AsyncMock(spec=AsyncBroker)
    broker.kick = AsyncMock()
    broker.result_backend = None
    return broker


@pytest_asyncio.fixture
async def mock_redis() -> AsyncMock:
    """Create a mock Redis client."""
    redis = AsyncMock(spec=Redis)
    redis.flushdb = AsyncMock()
    redis.get = AsyncMock(return_value=None)
    redis.set = AsyncMock(return_value=True)
    redis.delete = AsyncMock(return_value=1)
    redis.exists = AsyncMock(return_value=0)
    redis.eval = AsyncMock(return_value=[1, 0, 0])
    redis.zadd = AsyncMock(return_value=1)
    redis.zrem = AsyncMock(return_value=1)
    redis.zrange = AsyncMock(return_value=[])
    redis.hgetall = AsyncMock(return_value={})
    redis.hset = AsyncMock(return_value=1)
    redis.hincrby = AsyncMock(return_value=1)
    redis.incr = AsyncMock(return_value=1)
    redis.decr = AsyncMock(return_value=1)
    redis.expire = AsyncMock(return_value=True)
    return redis


@pytest_asyncio.fixture
async def in_memory_storage() -> InMemoryFairQueueStorage:
    """Create in-memory storage for testing."""
    return InMemoryFairQueueStorage()


@pytest_asyncio.fixture
async def redis_storage(mock_app_config: MockAppConfig) -> RedisFairQueueStorage:
    """Create Redis storage for testing."""
    return RedisFairQueueStorage(redis=Redis.from_url(str(mock_app_config.redis.url)), key_prefix="test_fair_queue")


@pytest_asyncio.fixture
async def sqlalchemy_storage(mock_app_config: MockAppConfig) -> AsyncIterator[DatabaseFairQueueStorage]:
    """Create SQLAlchemy storage for testing."""
    # Use in-memory SQLite for testing
    engine = create_async_engine(str(mock_app_config.sqlalchemy.url), echo=False)
    async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    # Create tables

    def create_tables(sync_conn: Any) -> None:
        """Create all tables."""
        FairTaskModel.__table__.create(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        GroupQueueStateModel.__table__.create(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        RateLimitModel.__table__.create(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        TaskResultModel.__table__.create(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        GlobalQueueSizeModel.__table__.create(sync_conn, checkfirst=True)  # type: ignore[attr-defined]

    async with engine.begin() as conn:
        await conn.run_sync(create_tables)

    session = async_session()
    storage = DatabaseFairQueueStorage(session=session)

    yield storage

    # Cleanup
    await session.close()

    def drop_tables(sync_conn: Any) -> None:
        """Drop all tables."""
        FairTaskModel.__table__.drop(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        GroupQueueStateModel.__table__.drop(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        RateLimitModel.__table__.drop(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        TaskResultModel.__table__.drop(sync_conn, checkfirst=True)  # type: ignore[attr-defined]
        GlobalQueueSizeModel.__table__.drop(sync_conn, checkfirst=True)  # type: ignore[attr-defined]

    async with engine.begin() as conn:
        await conn.run_sync(drop_tables)
    await engine.dispose()


@pytest_asyncio.fixture
async def mock_task_adapter(mock_broker: AsyncMock) -> Any:
    """Create a mock task adapter."""

    return TaskIQAdapter(broker=mock_broker)


@pytest_asyncio.fixture
def group_key_extractor() -> Any:
    """Create a simple group key extractor."""

    # Use functional extractor with fallback to default
    return compose_extractors(
        extract_from_kwargs("user_id"),
        extract_with_default("default"),
    )
