"""SQLAlchemy storage implementation for fair task queue.

This module includes:
- SQLAlchemy storage implementation
- SQLAlchemy database models
"""

import json
import logging
from datetime import UTC, datetime, timedelta
from typing import Any, NoReturn
from uuid import UUID

from sqlalchemy import (
    JSON,
    Boolean,
    DateTime,
    Float,
    ForeignKey,
    Index,
    Integer,
    LargeBinary,
    String,
    and_,
    delete,
    func,
    or_,
    select,
    update,
)
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import Mapped, mapped_column

from haolib.background.fair.abstract import (
    ComponentHealthStatus,
    FairnessAlgorithm,
    FairTask,
    GroupMetadata,
    GroupQueueState,
    GroupStates,
    HealthDetails,
    RateLimit,
    TaskMetadata,
    TaskStatus,
)
from haolib.background.fair.exceptions import (
    FairQueueIdempotencyConflictError,
    FairQueueStorageError,
)
from haolib.database.models.base.sqlalchemy import SQLAlchemyBaseModel
from haolib.database.models.mixins.sqlalchemy import SQLAlchemyDateTimeModel, SQLAlchemyIdModel

logger = logging.getLogger(__name__)

# ============================================================================
# SQLAlchemy Database Models
# ============================================================================


class FairTaskModel(SQLAlchemyBaseModel, SQLAlchemyIdModel[UUID], SQLAlchemyDateTimeModel):
    """Database model for fair queue tasks.

    Stores task data, status, and metadata for the fair task queue.
    """

    __tablename__ = "fair_queue_tasks"

    # Task identification
    task_id: Mapped[str] = mapped_column(String(255), unique=True, nullable=False, index=True)
    group_id: Mapped[str] = mapped_column(String(255), nullable=False, index=True)
    idempotency_key: Mapped[str | None] = mapped_column(String(255), unique=True, nullable=True, index=True)

    # Task data
    task_data: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
    status: Mapped[str] = mapped_column(String(50), nullable=False, index=True, default="pending")

    # Task metadata
    priority: Mapped[int] = mapped_column(Integer, nullable=False, default=0, index=True)
    retry_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    max_retries: Mapped[int] = mapped_column(Integer, nullable=False, default=3)

    # Scheduling
    scheduled_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
    expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)
    enqueued_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now())

    # Dependencies
    depends_on: Mapped[list[str] | None] = mapped_column(JSON, nullable=True)

    # Metadata for recovery and tracking (renamed from 'metadata' to avoid conflict with DeclarativeBase.metadata)
    task_metadata: Mapped[dict[str, Any] | None] = mapped_column("metadata", JSON, nullable=True)

    # Processing tracking
    processing_started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)

    # Dead letter queue flag
    in_dlq: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False, index=True)

    # Composite index for efficient queries
    __table_args__ = (
        Index("idx_group_status", "group_id", "status"),
        Index("idx_group_priority_scheduled", "group_id", "priority", "scheduled_at"),
        Index("idx_status_enqueued", "status", "enqueued_at"),
    )


class GroupQueueStateModel(SQLAlchemyBaseModel):
    """Database model for group queue state.

    Stores fairness algorithm state and metadata for each group.
    """

    __tablename__ = "fair_queue_group_states"

    group_id: Mapped[str] = mapped_column(String(255), primary_key=True, nullable=False)
    queue_length: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    virtual_finish_time: Mapped[float] = mapped_column(Float, nullable=False, default=0.0)
    weight: Mapped[float] = mapped_column(Float, nullable=False, default=1.0)
    last_processed_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
    # Metadata for group state (renamed from 'metadata' to avoid conflict with DeclarativeBase.metadata)
    group_metadata: Mapped[dict[str, Any] | None] = mapped_column("metadata", JSON, nullable=True)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
    )


class RateLimitModel(SQLAlchemyBaseModel):
    """Database model for rate limiting.

    Stores rate limit state for each group.
    """

    __tablename__ = "fair_queue_rate_limits"

    group_id: Mapped[str] = mapped_column(
        String(255), ForeignKey("fair_queue_group_states.group_id"), primary_key=True, nullable=False
    )
    max_tasks: Mapped[int] = mapped_column(Integer, nullable=False)
    per_time_seconds: Mapped[float] = mapped_column(Float, nullable=False)
    current_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    window_start: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, default=func.now())
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
    )


class TaskResultModel(SQLAlchemyBaseModel, SQLAlchemyIdModel[UUID], SQLAlchemyDateTimeModel):
    """Database model for task results.

    Stores task results with optional TTL.
    """

    __tablename__ = "fair_queue_task_results"

    task_id: Mapped[str] = mapped_column(String(255), unique=True, nullable=False, index=True)
    group_id: Mapped[str] = mapped_column(String(255), nullable=False, index=True)
    result_data: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
    expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True, index=True)

    __table_args__ = (Index("idx_result_expires", "expires_at"),)


class GlobalQueueSizeModel(SQLAlchemyBaseModel):
    """Database model for global queue size tracking.

    Stores the current global queue size atomically.
    """

    __tablename__ = "fair_queue_global_size"

    id: Mapped[int] = mapped_column(Integer, primary_key=True, default=1)
    current_size: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
    updated_at: Mapped[datetime] = mapped_column(
        DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now()
    )


# ============================================================================
# SQLAlchemy Storage Implementation
# ============================================================================

# Default weight for group queues
DEFAULT_GROUP_WEIGHT = 1.0

# Task result storage constants
DEFAULT_TASK_RESULT_TTL_SECONDS = 86400  # 24 hours default TTL for task results


class DatabaseFairQueueStorage:
    """Database storage implementation for fair task queue.

    This implementation uses SQLAlchemy with async sessions and provides
    atomic operations using database transactions and row-level locking.

    Note:
        This implementation requires:
        1. Database models to be created (use Alembic migrations)
        2. Proper database indexes for performance
        3. Transaction support in the database
        4. Proper connection pooling configuration

    Connection Pooling:
        The AsyncSession should be configured with appropriate connection pool
        settings for your workload. Recommended settings:
        - pool_size: 5-20 (depending on concurrent operations)
        - max_overflow: 10-30 (additional connections beyond pool_size)
        - pool_timeout: 30 (seconds to wait for connection)
        - pool_recycle: 3600 (recycle connections after 1 hour)

    Example:
        ```python
        from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
        from sqlalchemy.orm import sessionmaker

        engine = create_async_engine(
            "postgresql+asyncpg://...",
            pool_size=10,
            max_overflow=20,
            pool_timeout=30,
            pool_recycle=3600,
        )
        async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
        storage = DatabaseFairQueueStorage(session=async_session())
        ```

    Transaction Isolation:
        By default, operations use the session's isolation level. For high-concurrency
        scenarios, consider using READ COMMITTED or SERIALIZABLE isolation level
        depending on your consistency requirements.

    """

    def __init__(
        self,
        session: AsyncSession,
        operation_timeout: timedelta | None = None,
    ) -> None:
        """Initialize database storage.

        Args:
            session: SQLAlchemy async session
            operation_timeout: Optional timeout for storage operations (not used for DB, kept for API compatibility)

        """
        self._session = session
        self._operation_timeout = operation_timeout or timedelta(seconds=5)

    def _raise_idempotency_error(self, existing_task_id: str, idempotency_key: str) -> NoReturn:
        """Raise error for idempotency conflict.

        Args:
            existing_task_id: Existing task ID
            idempotency_key: Idempotency key that caused the conflict

        Raises:
            FairQueueIdempotencyConflictError: Always raises with idempotency conflict message

        """
        raise FairQueueIdempotencyConflictError(task_id=existing_task_id, idempotency_key=idempotency_key) from None

    async def enqueue_task(self, task: FairTask) -> None:
        """Enqueue a task to storage.

        Args:
            task: Task to enqueue

        Raises:
            FairQueueStorageError: If storage operation fails
            FairQueueIdempotencyConflictError: If idempotency key conflict occurs

        """
        try:
            # Check idempotency if provided
            if task.idempotency_key:
                existing_task = await self._session.execute(
                    select(FairTaskModel.task_id).where(FairTaskModel.idempotency_key == task.idempotency_key)
                )
                existing = existing_task.scalar_one_or_none()
                if existing and task.idempotency_key:
                    self._raise_idempotency_error(existing, task.idempotency_key)

            # Serialize task data
            if isinstance(task.task_data, bytes):
                task_data_bytes = task.task_data
            else:
                task_data_bytes = json.dumps(task.task_data, default=str).encode()

            # Create task model
            task_model = FairTaskModel(
                task_id=task.task_id,
                group_id=task.group_id,
                task_data=task_data_bytes,
                status=task.status.value,
                priority=task.priority,
                scheduled_at=task.scheduled_at,
                expires_at=task.expires_at,
                idempotency_key=task.idempotency_key,
                max_retries=task.max_retries,
                retry_count=task.retry_count,
                depends_on=task.depends_on if task.depends_on else None,
                task_metadata=task.metadata.model_dump() if task.metadata else None,
                enqueued_at=task.enqueued_at if task.enqueued_at else datetime.now(UTC),
            )

            self._session.add(task_model)

            # Update or create group state
            await self._update_group_state_internal(task.group_id, increment_queue_length=True)

            # Increment global queue size
            await self._increment_global_queue_size()

            await self._session.commit()

        except FairQueueIdempotencyConflictError:
            await self._session.rollback()
            raise
        except IntegrityError as e:
            await self._session.rollback()
            if "idempotency_key" in str(e).lower() and task.idempotency_key:
                # Try to get existing task
                existing_task = await self._session.execute(
                    select(FairTaskModel.task_id).where(FairTaskModel.idempotency_key == task.idempotency_key)
                )
                existing = existing_task.scalar_one_or_none()
                if existing and task.idempotency_key:
                    raise FairQueueIdempotencyConflictError(
                        task_id=existing, idempotency_key=task.idempotency_key
                    ) from e
            message = f"Failed to enqueue task: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task.task_id,
                group_id=task.group_id,
                operation="enqueue_task",
            ) from e
        except Exception as e:
            await self._session.rollback()
            message = f"Failed to enqueue task: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task.task_id,
                group_id=task.group_id,
                operation="enqueue_task",
            ) from e

    async def enqueue_tasks_batch(self, tasks: list[FairTask]) -> list[tuple[str, bool, str | None]]:
        """Enqueue multiple tasks atomically.

        Args:
            tasks: List of tasks to enqueue

        Returns:
            List of (task_id, success, error) tuples

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        if not tasks:
            return []

        results: list[tuple[str, bool, str | None]] = []

        try:
            # Check idempotency keys first
            idempotency_keys = [t.idempotency_key for t in tasks if t.idempotency_key]
            if idempotency_keys:
                existing_tasks = await self._session.execute(
                    select(FairTaskModel.task_id, FairTaskModel.idempotency_key).where(
                        FairTaskModel.idempotency_key.in_(idempotency_keys)
                    )
                )
                existing_map = {row.idempotency_key: row.task_id for row in existing_tasks}

                # Prepare task models
            task_models: list[FairTaskModel] = []
            for task in tasks:
                try:
                    # Check idempotency
                    if task.idempotency_key and task.idempotency_key in existing_map:
                        results.append((existing_map[task.idempotency_key], False, "idempotency_conflict"))
                        continue

                    # Serialize task data
                    if isinstance(task.task_data, bytes):
                        task_data_bytes = task.task_data
                    else:
                        task_data_bytes = json.dumps(task.task_data, default=str).encode()

                    task_model = FairTaskModel(
                        task_id=task.task_id,
                        group_id=task.group_id,
                        task_data=task_data_bytes,
                        status=task.status.value,
                        priority=task.priority,
                        scheduled_at=task.scheduled_at,
                        expires_at=task.expires_at,
                        idempotency_key=task.idempotency_key,
                        max_retries=task.max_retries,
                        retry_count=task.retry_count,
                        depends_on=task.depends_on if task.depends_on else None,
                        metadata=task.metadata.model_dump() if task.metadata else None,
                        enqueued_at=task.enqueued_at if task.enqueued_at else datetime.now(UTC),
                    )
                    task_models.append(task_model)
                    results.append((task.task_id, True, None))

                except Exception as e:
                    results.append((task.task_id, False, str(e)))

            if task_models:
                # Bulk insert tasks
                self._session.add_all(task_models)

                # Update group states (group by group_id)
                group_ids = {t.group_id for t in tasks}
                for group_id in group_ids:
                    await self._update_group_state_internal(group_id, increment_queue_length=True)

                # Increment global queue size
                await self._increment_global_queue_size(len(task_models))

                await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to enqueue tasks batch: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="enqueue_tasks_batch",
            ) from e

        return results

    async def dequeue_task(self, group_id: str) -> FairTask | None:
        """Dequeue next task from a group's queue.

        Args:
            group_id: Group ID to dequeue from

        Returns:
            Dequeued task or None if queue is empty

        Raises:
            FairQueueStorageEirror: If storage operation fails

        """
        try:
            current_time = datetime.now(UTC)

            # Find next available task (not expired, not scheduled for future, highest priority)
            # Use SELECT FOR UPDATE to lock the row
            stmt = (
                select(FairTaskModel)
                .where(
                    and_(
                        FairTaskModel.group_id == group_id,
                        FairTaskModel.status == TaskStatus.PENDING.value,
                        or_(
                            FairTaskModel.scheduled_at.is_(None),
                            FairTaskModel.scheduled_at <= current_time,
                        ),
                        or_(
                            FairTaskModel.expires_at.is_(None),
                            FairTaskModel.expires_at > current_time,
                        ),
                    )
                )
                .order_by(
                    FairTaskModel.priority.desc(),
                    FairTaskModel.enqueued_at.asc(),
                )
                .limit(1)
                .with_for_update(skip_locked=True)
            )

            result = await self._session.execute(stmt)
            task_model = result.scalar_one_or_none()

            if not task_model:
                return None

            # Check expiration one more time
            if task_model.expires_at and task_model.expires_at < current_time:
                task_model.status = TaskStatus.EXPIRED.value
                await self._session.commit()
                # Try next task
                return await self.dequeue_task(group_id)

            # Check if scheduled for future
            if task_model.scheduled_at and task_model.scheduled_at > current_time:
                # Re-enqueue and try next
                return await self.dequeue_task(group_id)

            # Update status to PROCESSING
            task_model.status = TaskStatus.PROCESSING.value
            task_model.processing_started_at = current_time

            # Update group state
            await self._update_group_state_internal(group_id, increment_queue_length=False)

            # Decrement global queue size
            await self._decrement_global_queue_size()

            await self._session.commit()

            # Convert to FairTask
            return self._model_to_fair_task(task_model)

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to dequeue task from group {group_id}: {e}"
            raise FairQueueStorageError(
                detail=message,
                group_id=group_id,
                operation="dequeue_task",
            ) from e

    async def get_next_group(self, algorithm: FairnessAlgorithm) -> str | None:
        """Get next group to process based on fairness algorithm.

        Optimized to use EXISTS subquery instead of JOIN for better performance
        on large datasets. Only loads group states that actually have pending tasks.

        Args:
            algorithm: Fairness algorithm to use

        Returns:
            Group ID to process next, or None if no groups available

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Optimized: Use EXISTS subquery instead of JOIN for better performance
            # This avoids loading all task rows and only checks for existence
            subquery = (
                select(1)
                .where(
                    and_(
                        FairTaskModel.group_id == GroupQueueStateModel.group_id,
                        FairTaskModel.status == TaskStatus.PENDING.value,
                    )
                )
                .exists()
            )

            # Get only group states that have pending tasks
            stmt = select(GroupQueueStateModel).where(subquery)

            result = await self._session.execute(stmt)
            state_models = result.scalars().all()

            if not state_models:
                return None

            # Build group states dict and verify queue length
            group_states_dict: dict[str, GroupQueueState] = {}
            for state_model in state_models:
                state = await self._model_to_group_state(state_model)
                # Double-check queue length (may have changed since query)
                if state and state.queue_length > 0:
                    group_states_dict[state.group_id] = state

            if not group_states_dict:
                return None

            # Use algorithm to select next group
            group_states = GroupStates(groups=group_states_dict)
            return algorithm.select_next_group(group_states)

        except Exception as e:
            message = f"Failed to get next group: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="get_next_group",
            ) from e

    async def get_group_queue_length(self, group_id: str) -> int:
        """Get queue length for a group.

        Uses atomic COUNT query to ensure accurate queue length even under
        concurrent enqueue/dequeue operations. This method is safe to call
        from multiple concurrent transactions.

        Args:
            group_id: Group ID

        Returns:
            Queue length (number of pending tasks for this group)

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Use atomic COUNT query - this is safe for concurrent access
            # The database ensures consistency at the transaction level
            stmt = select(func.count(FairTaskModel.task_id)).where(
                and_(
                    FairTaskModel.group_id == group_id,
                    FairTaskModel.status == TaskStatus.PENDING.value,
                )
            )
            result = await self._session.execute(stmt)
            return result.scalar() or 0
        except Exception as e:
            message = f"Failed to get queue length: {e}"
            raise FairQueueStorageError(
                detail=message,
                group_id=group_id,
                operation="get_group_queue_length",
            ) from e

    async def get_group_state(self, group_id: str) -> GroupQueueState | None:
        """Get current state of a group's queue.

        Args:
            group_id: Group ID

        Returns:
            Group queue state or None if group doesn't exist

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(GroupQueueStateModel).where(GroupQueueStateModel.group_id == group_id)
            result = await self._session.execute(stmt)
            state_model = result.scalar_one_or_none()

            if not state_model:
                # Check if group has any tasks
                queue_length = await self.get_group_queue_length(group_id)
                if queue_length == 0:
                    return None

                # Create default state
                return GroupQueueState(
                    group_id=group_id,
                    queue_length=queue_length,
                    virtual_finish_time=0.0,
                    weight=DEFAULT_GROUP_WEIGHT,
                )

            return await self._model_to_group_state(state_model)

        except Exception as e:
            message = f"Failed to get group state: {e}"
            raise FairQueueStorageError(
                detail=message,
                group_id=group_id,
                operation="get_group_state",
            ) from e

    async def get_all_groups(self) -> list[str]:
        """Get all active group IDs.

        Returns:
            List of all active group IDs

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(GroupQueueStateModel.group_id).distinct()
            result = await self._session.execute(stmt)
            return [row[0] for row in result.all()]
        except Exception as e:
            message = f"Failed to get all groups: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="get_all_groups",
            ) from e

    async def update_group_state(self, state: GroupQueueState) -> None:
        """Update group queue state.

        Args:
            state: Group state to update

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Update or insert group state (database-agnostic)
            existing = await self._session.get(GroupQueueStateModel, state.group_id)
            if existing:
                existing.queue_length = state.queue_length
                existing.virtual_finish_time = state.virtual_finish_time
                existing.weight = state.weight
                existing.last_processed_at = state.last_processed_at
                # Convert GroupMetadata to dict for database storage
                existing.group_metadata = state.metadata.model_dump() if state.metadata else None
            else:
                new_state = GroupQueueStateModel(
                    group_id=state.group_id,
                    queue_length=state.queue_length,
                    virtual_finish_time=state.virtual_finish_time,
                    weight=state.weight,
                    last_processed_at=state.last_processed_at,
                    group_metadata=state.metadata.model_dump() if state.metadata else None,
                )
                self._session.add(new_state)

            # Update rate limit if provided
            if state.rate_limit:
                existing_rate_limit = await self._session.get(RateLimitModel, state.group_id)
                if existing_rate_limit:
                    existing_rate_limit.max_tasks = state.rate_limit.max_tasks
                    existing_rate_limit.per_time_seconds = state.rate_limit.per_time.total_seconds()
                    existing_rate_limit.current_count = state.rate_limit.current_count
                    existing_rate_limit.window_start = state.rate_limit.window_start
                else:
                    new_rate_limit = RateLimitModel(
                        group_id=state.group_id,
                        max_tasks=state.rate_limit.max_tasks,
                        per_time_seconds=state.rate_limit.per_time.total_seconds(),
                        current_count=state.rate_limit.current_count,
                        window_start=state.rate_limit.window_start,
                    )
                    self._session.add(new_rate_limit)

            await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to update group state: {e}"
            raise FairQueueStorageError(
                detail=message,
                group_id=state.group_id,
                operation="update_group_state",
            ) from e

    async def get_task_by_idempotency_key(self, idempotency_key: str) -> str | None:
        """Get task ID by idempotency key.

        Args:
            idempotency_key: Idempotency key

        Returns:
            Task ID if found, None otherwise

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(FairTaskModel.task_id).where(FairTaskModel.idempotency_key == idempotency_key)
            result = await self._session.execute(stmt)
            return result.scalar_one_or_none()
        except Exception as e:
            message = f"Failed to get task by idempotency key: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="get_task_by_idempotency_key",
            ) from e

    async def get_task_status(self, task_id: str, group_id: str | None = None) -> TaskStatus | None:
        """Get task status by task ID and optional group ID.

        Args:
            task_id: Task ID
            group_id: Optional group ID (will be derived if not provided)

        Returns:
            Task status or None if task not found

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(FairTaskModel.status)
            if group_id:
                stmt = stmt.where(
                    and_(
                        FairTaskModel.task_id == task_id,
                        FairTaskModel.group_id == group_id,
                    )
                )
            else:
                stmt = stmt.where(FairTaskModel.task_id == task_id)

            result = await self._session.execute(stmt)
            status_str = result.scalar_one_or_none()

            if not status_str:
                return None

            return TaskStatus(status_str)

        except ValueError:
            return None
        except Exception as e:
            message = f"Failed to get task status: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="get_task_status",
            ) from e

    async def update_task_status(
        self, task_id: str, group_id: str, status: TaskStatus, metadata: TaskMetadata | None = None
    ) -> None:
        """Update task status and optionally update metadata.

        Args:
            task_id: Task ID
            group_id: Group ID
            status: New task status
            metadata: Optional metadata to update in task data

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = (
                update(FairTaskModel)
                .where(
                    and_(
                        FairTaskModel.task_id == task_id,
                        FairTaskModel.group_id == group_id,
                    )
                )
                .values(status=status.value)
            )

            if metadata:
                # Get current metadata and update
                current_task = await self._session.execute(
                    select(FairTaskModel.task_metadata).where(
                        and_(
                            FairTaskModel.task_id == task_id,
                            FairTaskModel.group_id == group_id,
                        )
                    )
                )
                current_metadata_dict = current_task.scalar_one_or_none() or {}
                if isinstance(current_metadata_dict, dict):
                    # Merge with new metadata
                    current_metadata = (
                        TaskMetadata.model_validate(current_metadata_dict) if current_metadata_dict else TaskMetadata()
                    )
                    # Update fields from new metadata
                    if metadata.processing_started_at:
                        current_metadata.processing_started_at = metadata.processing_started_at
                    if metadata.processing_completed_at:
                        current_metadata.processing_completed_at = metadata.processing_completed_at
                    if metadata.error_message:
                        current_metadata.error_message = metadata.error_message
                    if metadata.error_type:
                        current_metadata.error_type = metadata.error_type
                    current_metadata.custom_data.update(metadata.custom_data)
                    # Convert back to dict for storage
                    stmt = stmt.values(task_metadata=current_metadata.model_dump())

            await self._session.execute(stmt)
            await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to update task status: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="update_task_status",
            ) from e

    async def remove_task(self, task_id: str, group_id: str) -> None:
        """Remove a task from the queue.

        Args:
            task_id: Task ID to remove
            group_id: Group ID of the task

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = delete(FairTaskModel).where(
                and_(
                    FairTaskModel.task_id == task_id,
                    FairTaskModel.group_id == group_id,
                )
            )
            await self._session.execute(stmt)

            # Update group state
            await self._update_group_state_internal(group_id, increment_queue_length=False)

            # Decrement global queue size
            await self._decrement_global_queue_size()

            await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to remove task: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="remove_task",
            ) from e

    async def mark_task_completed(self, task_id: str, group_id: str) -> None:
        """Mark task as completed and clean up task data.

        Args:
            task_id: Task ID
            group_id: Group ID

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = (
                update(FairTaskModel)
                .where(
                    and_(
                        FairTaskModel.task_id == task_id,
                        FairTaskModel.group_id == group_id,
                    )
                )
                .values(status=TaskStatus.COMPLETED.value)
            )
            await self._session.execute(stmt)

            # Decrement global queue size
            await self._decrement_global_queue_size()

            await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to mark task as completed: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="mark_task_completed",
            ) from e

    async def mark_task_failed(self, task_id: str, group_id: str) -> None:
        """Mark task as failed and preserve task data for DLQ.

        Args:
            task_id: Task ID
            group_id: Group ID

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = (
                update(FairTaskModel)
                .where(
                    and_(
                        FairTaskModel.task_id == task_id,
                        FairTaskModel.group_id == group_id,
                    )
                )
                .values(status=TaskStatus.FAILED.value, in_dlq=True)
            )
            await self._session.execute(stmt)
            await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to mark task as failed: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="mark_task_failed",
            ) from e

    async def add_to_dead_letter_queue(self, task: FairTask) -> None:
        """Add a task to the dead letter queue.

        Args:
            task: Failed task to add to DLQ

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Task should already be marked as failed, just ensure in_dlq is set
            stmt = (
                update(FairTaskModel)
                .where(
                    and_(
                        FairTaskModel.task_id == task.task_id,
                        FairTaskModel.group_id == task.group_id,
                    )
                )
                .values(in_dlq=True, status=TaskStatus.FAILED.value)
            )
            await self._session.execute(stmt)
            await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to add task to dead letter queue: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task.task_id,
                group_id=task.group_id,
                operation="add_to_dead_letter_queue",
            ) from e

    async def store_task_result(self, task_id: str, group_id: str, result: Any, ttl_seconds: int | None = None) -> None:
        """Store task result.

        Args:
            task_id: Task ID
            group_id: Group ID
            result: Task result to store
            ttl_seconds: Optional TTL for result in seconds

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            result_data = json.dumps(result, default=str).encode()
            expires_at = None
            if ttl_seconds:
                expires_at = datetime.now(UTC) + timedelta(seconds=ttl_seconds)

            # Use upsert to handle existing results (database-agnostic)
            existing_result = await self._session.execute(
                select(TaskResultModel).where(TaskResultModel.task_id == task_id)
            )
            existing = existing_result.scalar_one_or_none()
            if existing:
                existing.result_data = result_data
                existing.expires_at = expires_at
            else:
                new_result = TaskResultModel(
                    task_id=task_id,
                    group_id=group_id,
                    result_data=result_data,
                    expires_at=expires_at,
                )
                self._session.add(new_result)
            await self._session.commit()

        except Exception as e:
            await self._session.rollback()
            message = f"Failed to store task result: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="store_task_result",
            ) from e

    async def get_task_result(self, task_id: str, group_id: str | None = None) -> Any | None:
        """Get task result.

        Args:
            task_id: Task ID
            group_id: Optional group ID (not strictly needed, kept for API consistency)

        Returns:
            Task result or None if not found

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(TaskResultModel.result_data, TaskResultModel.expires_at).where(
                TaskResultModel.task_id == task_id
            )
            if group_id:
                stmt = stmt.where(TaskResultModel.group_id == group_id)

            result = await self._session.execute(stmt)
            row = result.first()

            if not row:
                return None

            result_data, expires_at = row

            # Check expiration
            if expires_at and expires_at < datetime.now(UTC):
                # Result expired, delete it
                await self._session.execute(delete(TaskResultModel).where(TaskResultModel.task_id == task_id))
                await self._session.commit()
                return None

            return json.loads(result_data.decode())

        except json.JSONDecodeError as e:
            message = f"Failed to decode task result: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="get_task_result",
            ) from e
        except Exception as e:
            message = f"Failed to get task result: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="get_task_result",
            ) from e

    async def get_task(self, task_id: str, group_id: str) -> FairTask | None:
        """Get task by ID and group ID.

        Args:
            task_id: Task ID
            group_id: Group ID

        Returns:
            FairTask if found, None otherwise

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(FairTaskModel).where(
                and_(
                    FairTaskModel.task_id == task_id,
                    FairTaskModel.group_id == group_id,
                )
            )
            result = await self._session.execute(stmt)
            task_model = result.scalar_one_or_none()

            if not task_model:
                return None

            return self._model_to_fair_task(task_model)

        except Exception as e:
            message = f"Failed to get task: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                group_id=group_id,
                operation="get_task",
            ) from e

    async def get_group_id_by_task_id(self, task_id: str) -> str | None:
        """Get group ID by task ID.

        Args:
            task_id: Task ID

        Returns:
            Group ID if found, None otherwise

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(FairTaskModel.group_id).where(FairTaskModel.task_id == task_id)
            result = await self._session.execute(stmt)
            return result.scalar_one_or_none()
        except Exception as e:
            message = f"Failed to get group ID for task {task_id}: {e}"
            raise FairQueueStorageError(
                detail=message,
                task_id=task_id,
                operation="get_group_id_by_task_id",
            ) from e

    async def purge_expired_tasks(self, before: datetime) -> int:
        """Remove all expired tasks before the given timestamp.

        Args:
            before: Remove tasks expired before this timestamp

        Returns:
            Number of tasks removed

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Get count before deletion
            count_stmt = select(func.count(FairTaskModel.task_id)).where(
                and_(
                    FairTaskModel.expires_at.isnot(None),
                    FairTaskModel.expires_at < before,
                )
            )
            count_result = await self._session.execute(count_stmt)
            count = count_result.scalar() or 0

            # Delete expired tasks
            stmt = delete(FairTaskModel).where(
                and_(
                    FairTaskModel.expires_at.isnot(None),
                    FairTaskModel.expires_at < before,
                )
            )
            await self._session.execute(stmt)
            await self._session.commit()
        except Exception as e:
            await self._session.rollback()
            message = f"Failed to purge expired tasks: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="purge_expired_tasks",
            ) from e
        else:
            return count

    async def get_dead_letter_queue(
        self,
        group_id: str | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> list[FairTask]:
        """Get tasks from dead letter queue with pagination.

        Args:
            group_id: Optional group ID filter
            limit: Maximum number of tasks to return
            offset: Number of tasks to skip

        Returns:
            List of failed tasks

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = (
                select(FairTaskModel)
                .where(
                    and_(
                        FairTaskModel.in_dlq.is_(True),
                        FairTaskModel.status == TaskStatus.FAILED.value,
                    )
                )
                .order_by(FairTaskModel.enqueued_at.desc())
                .limit(limit)
                .offset(offset)
            )

            if group_id:
                stmt = stmt.where(FairTaskModel.group_id == group_id)

            result = await self._session.execute(stmt)
            task_models = result.scalars().all()

            return [self._model_to_fair_task(model) for model in task_models]

        except Exception as e:
            message = f"Failed to get dead letter queue: {e}"
            raise FairQueueStorageError(
                detail=message,
                group_id=group_id,
                operation="get_dead_letter_queue",
            ) from e

    async def get_tasks_by_status(
        self,
        status: TaskStatus,
        group_id: str | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> list[tuple[str, str]]:
        """Get task IDs by status with optional group filter.

        Args:
            status: Task status to filter by
            group_id: Optional group ID filter
            limit: Maximum number of tasks to return
            offset: Number of tasks to skip

        Returns:
            List of (task_id, group_id) tuples

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = (
                select(FairTaskModel.task_id, FairTaskModel.group_id)
                .where(FairTaskModel.status == status.value)
                .limit(limit)
                .offset(offset)
            )

            if group_id:
                stmt = stmt.where(FairTaskModel.group_id == group_id)

            result = await self._session.execute(stmt)
            return [(row[0], row[1]) for row in result.all()]

        except Exception as e:
            message = f"Failed to get tasks by status: {e}"
            raise FairQueueStorageError(
                detail=message,
                group_id=group_id,
                operation="get_tasks_by_status",
            ) from e

    async def health_check(self) -> bool:
        """Check if storage is healthy.

        Returns:
            True if storage is healthy, False otherwise

        """
        try:
            # Simple query to check database connectivity
            await self._session.execute(select(1))
        except Exception:
            return False
        else:
            return True

    async def get_health_status(self) -> ComponentHealthStatus:
        """Get detailed health status.

        Returns:
            ComponentHealthStatus model with health status details

        """
        try:
            # Check database connectivity
            await self._session.execute(select(1))

            # Get some basic stats
            task_count = await self._session.execute(select(func.count(FairTaskModel.task_id)))
            total_tasks = task_count.scalar() or 0

            group_count = await self._session.execute(select(func.count(GroupQueueStateModel.group_id)))
            total_groups = group_count.scalar() or 0

            return ComponentHealthStatus(
                healthy=True,
                component="storage",
                details=HealthDetails(
                    custom_data={
                        "total_tasks": total_tasks,
                        "total_groups": total_groups,
                    }
                ),
                error=None,
                timestamp=datetime.now(UTC),
            )
        except Exception as e:
            return ComponentHealthStatus(
                healthy=False,
                component="storage",
                details=HealthDetails(custom_data={}),
                error=str(e),
                timestamp=datetime.now(UTC),
            )

    async def check_and_increment_rate_limit(
        self,
        group_id: str,
        max_tasks: int,
        per_time: timedelta,
    ) -> tuple[bool, datetime, int]:
        """Atomically check and increment rate limit.

        Uses database transactions with row-level locking for atomicity.

        Args:
            group_id: Group ID
            max_tasks: Maximum tasks allowed
            per_time: Time window

        Returns:
            Tuple of (allowed, window_start, current_count)

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            current_time = datetime.now(UTC)
            per_time_seconds = per_time.total_seconds()

            # Get or create rate limit with lock
            stmt = select(RateLimitModel).where(RateLimitModel.group_id == group_id).with_for_update()
            result = await self._session.execute(stmt)
            rate_limit_model = result.scalar_one_or_none()

            if not rate_limit_model:
                # Create new rate limit
                rate_limit_model = RateLimitModel(
                    group_id=group_id,
                    max_tasks=max_tasks,
                    per_time_seconds=per_time_seconds,
                    current_count=0,
                    window_start=current_time,
                )
                self._session.add(rate_limit_model)
                allowed = True
                window_start = current_time
                current_count = 1
            else:
                # Check if window has expired
                window_age = (current_time - rate_limit_model.window_start).total_seconds()
                if window_age >= per_time_seconds:
                    # Reset window
                    rate_limit_model.window_start = current_time
                    rate_limit_model.current_count = 1
                    allowed = True
                    window_start = current_time
                    current_count = 1
                # Check if limit exceeded
                elif rate_limit_model.current_count >= max_tasks:
                    allowed = False
                    window_start = rate_limit_model.window_start
                    current_count = rate_limit_model.current_count
                else:
                    # Increment count
                    rate_limit_model.current_count += 1
                    allowed = True
                    window_start = rate_limit_model.window_start
                    current_count = rate_limit_model.current_count

            await self._session.commit()
        except Exception as e:
            await self._session.rollback()
            message = f"Failed to check rate limit: {e}"
            raise FairQueueStorageError(
                detail=message,
                group_id=group_id,
                operation="check_and_increment_rate_limit",
            ) from e
        else:
            return (allowed, window_start, current_count)

    async def check_and_increment_global_queue_size(
        self,
        max_size: int,
    ) -> tuple[bool, int]:
        """Atomically check and increment global queue size.

        Args:
            max_size: Maximum global queue size

        Returns:
            Tuple of (allowed, current_size)

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Get or create global size with lock
            stmt = select(GlobalQueueSizeModel).where(GlobalQueueSizeModel.id == 1).with_for_update()
            result = await self._session.execute(stmt)
            size_model = result.scalar_one_or_none()

            if not size_model:
                size_model = GlobalQueueSizeModel(id=1, current_size=1)
                self._session.add(size_model)
                allowed = max_size is None or max_size >= 1
                current_size = 1
            else:
                new_size = size_model.current_size + 1
                allowed = max_size is None or new_size <= max_size
                if allowed:
                    size_model.current_size = new_size
                current_size = size_model.current_size

            await self._session.commit()
        except Exception as e:
            await self._session.rollback()
            message = f"Failed to check and increment global queue size: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="check_and_increment_global_queue_size",
            ) from e
        else:
            return (allowed, current_size)

    async def check_and_increment_global_queue_size_batch(
        self,
        max_size: int,
        batch_size: int,
    ) -> tuple[bool, int]:
        """Atomically check and increment global queue size for a batch of tasks.

        Args:
            max_size: Maximum global queue size
            batch_size: Number of tasks to add in the batch

        Returns:
            Tuple of (allowed, current_size)

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Get or create global size with lock
            stmt = select(GlobalQueueSizeModel).where(GlobalQueueSizeModel.id == 1).with_for_update()
            result = await self._session.execute(stmt)
            size_model = result.scalar_one_or_none()

            if not size_model:
                size_model = GlobalQueueSizeModel(id=1, current_size=batch_size)
                self._session.add(size_model)
                allowed = max_size is None or batch_size <= max_size
                current_size = batch_size
            else:
                new_size = size_model.current_size + batch_size
                allowed = max_size is None or new_size <= max_size
                if allowed:
                    size_model.current_size = new_size
                current_size = size_model.current_size

            await self._session.commit()
        except Exception as e:
            await self._session.rollback()
            message = f"Failed to check and increment global queue size for batch: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="check_and_increment_global_queue_size_batch",
            ) from e
        else:
            return (allowed, current_size)

    async def get_global_queue_size(self) -> int:
        """Get current global queue size without incrementing.

        Returns:
            Current global queue size

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            stmt = select(GlobalQueueSizeModel).where(GlobalQueueSizeModel.id == 1)
            result = await self._session.execute(stmt)
            size_model = result.scalar_one_or_none()
        except Exception as e:
            message = f"Failed to get global queue size: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="get_global_queue_size",
            ) from e
        else:
            if size_model:
                return size_model.current_size
            return 0

    async def decrement_global_queue_size(self) -> None:
        """Atomically decrement global queue size.

        Uses SELECT FOR UPDATE to ensure atomic decrement and prevent
        negative values. This method is safe to call concurrently.

        Raises:
            FairQueueStorageError: If storage operation fails

        """
        try:
            # Use SELECT FOR UPDATE to ensure atomic operation
            stmt = select(GlobalQueueSizeModel).where(GlobalQueueSizeModel.id == 1).with_for_update(nowait=True)
            result = await self._session.execute(stmt)
            existing = result.scalar_one_or_none()

            if existing:
                # Ensure we don't go negative (defensive programming)
                existing.current_size = max(0, existing.current_size - 1)
            # If it doesn't exist, that's fine - it means size is already 0
            # No need to create a record with negative value

            await self._session.commit()
        except Exception as e:
            await self._session.rollback()
            message = f"Failed to decrement global queue size: {e}"
            raise FairQueueStorageError(
                detail=message,
                operation="decrement_global_queue_size",
            ) from e

    # Helper methods

    def _model_to_fair_task(self, model: FairTaskModel) -> FairTask:
        """Convert database model to FairTask.

        Args:
            model: Database model

        Returns:
            FairTask object

        """
        # Deserialize task data
        try:
            task_data = json.loads(model.task_data.decode())
        except (json.JSONDecodeError, UnicodeDecodeError):
            # Fallback: return as bytes
            task_data = model.task_data

        return FairTask(
            task_id=model.task_id,
            group_id=model.group_id,
            task_data=task_data,
            enqueued_at=model.enqueued_at,
            scheduled_at=model.scheduled_at,
            expires_at=model.expires_at,
            priority=model.priority,
            status=TaskStatus(model.status),
            retry_count=model.retry_count,
            max_retries=model.max_retries,
            metadata=TaskMetadata.model_validate(model.task_metadata) if model.task_metadata else TaskMetadata(),
            idempotency_key=model.idempotency_key,
            depends_on=model.depends_on or [],
        )

    async def _model_to_group_state(self, model: GroupQueueStateModel) -> GroupQueueState:
        """Convert database model to GroupQueueState.

        Args:
            model: Database model

        Returns:
            GroupQueueState object

        """
        # Get rate limit if exists
        rate_limit = None
        rate_limit_stmt = select(RateLimitModel).where(RateLimitModel.group_id == model.group_id)
        rate_limit_result = await self._session.execute(rate_limit_stmt)
        rate_limit_model = rate_limit_result.scalar_one_or_none()

        if rate_limit_model:
            rate_limit = RateLimit(
                max_tasks=rate_limit_model.max_tasks,
                per_time=timedelta(seconds=rate_limit_model.per_time_seconds),
                current_count=rate_limit_model.current_count,
                window_start=rate_limit_model.window_start,
            )

        # Get actual queue length
        queue_length = await self.get_group_queue_length(model.group_id)

        return GroupQueueState(
            group_id=model.group_id,
            queue_length=queue_length,
            virtual_finish_time=model.virtual_finish_time,
            weight=model.weight,
            last_processed_at=model.last_processed_at,
            rate_limit=rate_limit,
            metadata=GroupMetadata.model_validate(model.group_metadata) if model.group_metadata else GroupMetadata(),
        )

    async def _update_group_state_internal(
        self,
        group_id: str,
        increment_queue_length: bool,  # noqa: ARG002
    ) -> None:
        """Internal method to update group state queue length.

        Uses atomic database operations with proper transaction isolation
        to prevent race conditions. Queue length is calculated directly
        from the database using COUNT within the same transaction,
        ensuring consistency even under concurrent access.

        Args:
            group_id: Group ID
            increment_queue_length: If True, increment queue length; if False, decrement
                Note: This parameter is kept for API compatibility but queue_length
                is always recalculated atomically from the database.

        """
        # Use SELECT FOR UPDATE NOWAIT to prevent deadlocks and ensure atomicity
        # This ensures we get an exclusive lock on the group state row
        stmt = (
            select(GroupQueueStateModel)
            .where(GroupQueueStateModel.group_id == group_id)
            .with_for_update(nowait=True, skip_locked=False)
        )
        result = await self._session.execute(stmt)
        existing = result.scalar_one_or_none()

        # Get actual queue length atomically from database within the same transaction
        # This prevents race conditions by using a database-level COUNT query
        # which is atomic and always reflects the current state
        # We do this AFTER acquiring the lock to ensure consistency
        queue_length = await self.get_group_queue_length(group_id)

        if existing:
            existing.queue_length = queue_length
        else:
            new_state = GroupQueueStateModel(
                group_id=group_id,
                queue_length=queue_length,
                virtual_finish_time=0.0,
                weight=DEFAULT_GROUP_WEIGHT,
            )
            self._session.add(new_state)

    async def _increment_global_queue_size(self, amount: int = 1) -> None:
        """Increment global queue size atomically.

        Uses SELECT FOR UPDATE to prevent race conditions in concurrent scenarios.

        Args:
            amount: Amount to increment by

        """
        # Use SELECT FOR UPDATE NOWAIT to prevent deadlocks
        stmt = select(GlobalQueueSizeModel).where(GlobalQueueSizeModel.id == 1).with_for_update(nowait=True)
        result = await self._session.execute(stmt)
        existing = result.scalar_one_or_none()

        if existing:
            existing.current_size += amount
        else:
            new_size = GlobalQueueSizeModel(id=1, current_size=amount)
            self._session.add(new_size)

    async def _decrement_global_queue_size(self) -> None:
        """Decrement global queue size."""
        await self.decrement_global_queue_size()
