"""In-memory storage implementation for fair task queue.

This implementation is useful for:
- Testing without external dependencies
- Development and prototyping
- Single-instance deployments where persistence is not required

Note: This implementation is NOT suitable for production use as it:
- Does not persist data across restarts
- Is not thread-safe for concurrent access
- Does not support distributed deployments
"""

import asyncio
from collections import defaultdict
from datetime import UTC, datetime, timedelta
from typing import Any

from haolib.background.fair.abstract import (
    ComponentHealthStatus,
    FairnessAlgorithm,
    FairTask,
    GroupQueueState,
    GroupStates,
    HealthDetails,
    TaskMetadata,
    TaskStatus,
)
from haolib.background.fair.exceptions import (
    FairQueueIdempotencyConflictError,
)


class InMemoryFairQueueStorage:
    """In-memory storage implementation for fair task queue.

    This is a simple, thread-safe in-memory implementation suitable for
    testing and development. It does not persist data across restarts.

    Example:
        ```python
        from haolib.background.fair.storages.memory import InMemoryFairQueueStorage

        storage = InMemoryFairQueueStorage()
        queue = FairTaskQueue(
            storage=storage,
            task_adapter=task_adapter,
            group_key_extractor=extractor,
        )
        ```

    """

    def __init__(self) -> None:
        """Initialize in-memory storage."""
        self._lock = asyncio.Lock()
        # Task storage: task_id -> FairTask
        self._tasks: dict[str, FairTask] = {}
        # Task status: task_id -> TaskStatus
        self._task_status: dict[str, TaskStatus] = {}
        # Group queues: group_id -> list[task_id] (sorted by priority)
        self._group_queues: dict[str, list[str]] = defaultdict(list)
        # Group states: group_id -> GroupQueueState
        self._group_states: dict[str, GroupQueueState] = {}
        # All groups set
        self._groups: set[str] = set()
        # Idempotency mapping: idempotency_key -> task_id
        self._idempotency_keys: dict[str, str] = {}
        # Task results: task_id -> result
        self._task_results: dict[str, Any] = {}
        # Dead letter queue: group_id -> list[task_id]
        self._dead_letter_queues: dict[str, list[str]] = defaultdict(list)
        # Global queue size counter
        self._global_queue_size: int = 0
        # Rate limit tracking: group_id -> (window_start, current_count)
        self._rate_limits: dict[str, tuple[datetime, int]] = {}

    def _enqueue_task_locked(self, task: FairTask) -> None:
        """Enqueue a task to storage (assumes lock is already held).

        Args:
            task: Task to enqueue

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

        """
        # Check idempotency
        if task.idempotency_key:
            if task.idempotency_key in self._idempotency_keys:
                existing_task_id = self._idempotency_keys[task.idempotency_key]
                raise FairQueueIdempotencyConflictError(
                    task_id=existing_task_id,
                    idempotency_key=task.idempotency_key,
                    group_id=task.group_id,
                )
            self._idempotency_keys[task.idempotency_key] = task.task_id

        # Store task
        self._tasks[task.task_id] = task
        self._task_status[task.task_id] = task.status

        # Add to group queue (sorted by priority, then enqueued_at)
        if task.task_id not in self._group_queues[task.group_id]:
            self._group_queues[task.group_id].append(task.task_id)
            # Sort by priority (descending), then by enqueued_at (ascending)
            self._group_queues[task.group_id].sort(
                key=lambda tid: (
                    -self._tasks[tid].priority,
                    self._tasks[tid].enqueued_at.timestamp(),
                )
            )

        # Add group to set
        self._groups.add(task.group_id)

        # Update group state (lock already held)
        self._update_group_state_internal_locked(task.group_id, increment_queue_length=True)

        # Increment global queue size
        self._global_queue_size += 1

    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

        """
        async with self._lock:
            self._enqueue_task_locked(task)

    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

        """
        results: list[tuple[str, bool, str | None]] = []
        async with self._lock:
            for task in tasks:
                try:
                    # Check idempotency
                    if task.idempotency_key and task.idempotency_key in self._idempotency_keys:
                        existing_task_id = self._idempotency_keys[task.idempotency_key]
                        results.append((existing_task_id, True, None))
                        continue

                    # Enqueue task (lock already held)
                    self._enqueue_task_locked(task)
                    results.append((task.task_id, True, None))
                except FairQueueIdempotencyConflictError as e:
                    if e.task_id:
                        results.append((e.task_id, False, "idempotency_conflict"))
                    else:
                        results.append((task.task_id, False, "idempotency_conflict"))
                except Exception as e:
                    results.append((task.task_id, False, str(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

        """
        async with self._lock:
            if group_id not in self._group_queues or not self._group_queues[group_id]:
                return None

            # Get next task ID
            task_id = self._group_queues[group_id].pop(0)

            # Get task
            task = self._tasks.get(task_id)
            if not task:
                return None

            # Check if task is scheduled for future
            if task.scheduled_at and task.scheduled_at > datetime.now(UTC):
                # Re-add to queue
                self._group_queues[group_id].append(task_id)
                self._group_queues[group_id].sort(
                    key=lambda tid: (
                        -self._tasks[tid].priority,
                        self._tasks[tid].enqueued_at.timestamp(),
                    )
                )
                return None

            # Check expiration (lock already held)
            if task.expires_at and task.expires_at < datetime.now(UTC):
                # Update status directly since lock is already held
                task.status = TaskStatus.EXPIRED
                self._task_status[task.task_id] = TaskStatus.EXPIRED
                # Remove from queue
                if task.task_id in self._group_queues[group_id]:
                    self._group_queues[group_id].remove(task.task_id)
                # Update group state
                self._update_group_state_internal_locked(group_id, increment_queue_length=False)
                # Decrement global queue size
                self._global_queue_size = max(0, self._global_queue_size - 1)
                return None

            # Update status
            task.status = TaskStatus.PROCESSING
            self._task_status[task.task_id] = TaskStatus.PROCESSING

            # Update group state (lock already held)
            self._update_group_state_internal_locked(group_id, increment_queue_length=False)

            # Decrement global queue size
            self._global_queue_size = max(0, self._global_queue_size - 1)

            return task

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

        Args:
            algorithm: Fairness algorithm to use

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

        """
        async with self._lock:
            # Get all groups with tasks (lock already held)
            group_states: dict[str, GroupQueueState] = {}
            for group_id in self._groups:
                queue_length = len(self._group_queues.get(group_id, []))
                if queue_length > 0:
                    state = self._get_group_state_locked(group_id)
                    if state:
                        group_states[group_id] = state

            if not group_states:
                return None

            return algorithm.select_next_group(GroupStates(groups=group_states))

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

        Args:
            group_id: Group ID

        Returns:
            Queue length

        """
        async with self._lock:
            return len(self._group_queues.get(group_id, []))

    def _get_group_state_locked(self, group_id: str) -> GroupQueueState | None:
        """Get current state of a group's queue (assumes lock is already held).

        Args:
            group_id: Group ID

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

        """
        if group_id in self._group_states:
            state = self._group_states[group_id]
            # Update queue length
            state.queue_length = len(self._group_queues.get(group_id, []))
            return state

        queue_length = len(self._group_queues.get(group_id, []))
        if queue_length == 0:
            return None

        # Create default state
        state = GroupQueueState(
            group_id=group_id,
            queue_length=queue_length,
            virtual_finish_time=0.0,
            weight=1.0,
        )
        self._group_states[group_id] = state
        return state

    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

        """
        async with self._lock:
            return self._get_group_state_locked(group_id)

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

        Returns:
            List of all active group IDs

        """
        async with self._lock:
            return list(self._groups)

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

        Args:
            state: Group state to update

        """
        async with self._lock:
            self._group_states[state.group_id] = state

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

        Args:
            task_id: Task ID
            group_id: Optional group ID (unused in in-memory implementation)

        Returns:
            Task status or None if task not found

        """
        async with self._lock:
            return self._task_status.get(task_id)

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

        Args:
            task_id: Task ID
            group_id: Group ID

        Returns:
            FairTask if found, None otherwise

        """
        async with self._lock:
            return self._tasks.get(task_id)

    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

        """
        async with self._lock:
            return self._idempotency_keys.get(idempotency_key)

    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

        """
        async with self._lock:
            task = self._tasks.get(task_id)
            if task:
                return task.group_id
            return None

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

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

        """
        async with self._lock:
            self._task_status[task_id] = status
            if task_id in self._tasks:
                self._tasks[task_id].status = status
                if metadata:
                    # Update metadata fields
                    if metadata.processing_started_at:
                        self._tasks[task_id].metadata.processing_started_at = metadata.processing_started_at
                    if metadata.processing_completed_at:
                        self._tasks[task_id].metadata.processing_completed_at = metadata.processing_completed_at
                    if metadata.error_message:
                        self._tasks[task_id].metadata.error_message = metadata.error_message
                    if metadata.error_type:
                        self._tasks[task_id].metadata.error_type = metadata.error_type
                    if metadata.custom_data:
                        self._tasks[task_id].metadata.custom_data.update(metadata.custom_data)

    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

        """
        async with self._lock:
            self._task_status[task_id] = TaskStatus.COMPLETED
            if task_id in self._tasks:
                self._tasks[task_id].status = TaskStatus.COMPLETED
            # Remove from queue if still there
            if group_id in self._group_queues and task_id in self._group_queues[group_id]:
                self._group_queues[group_id].remove(task_id)
            # Decrement global queue size
            self._global_queue_size = max(0, self._global_queue_size - 1)

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

        Args:
            task_id: Task ID
            group_id: Group ID

        """
        async with self._lock:
            self._task_status[task_id] = TaskStatus.FAILED
            if task_id in self._tasks:
                self._tasks[task_id].status = TaskStatus.FAILED

    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

        """
        async with self._lock:
            self._dead_letter_queues[task.group_id].append(task.task_id)

    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.

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

        Returns:
            Tuple of (allowed, window_start, current_count)

        """
        async with self._lock:
            current_time = datetime.now(UTC)
            window_start, current_count = self._rate_limits.get(group_id, (current_time, 0))

            # Check if window has expired
            if current_time - window_start >= per_time:
                # Reset window
                window_start = current_time
                current_count = 0

            # Check if limit exceeded
            if current_count >= max_tasks:
                self._rate_limits[group_id] = (window_start, current_count)
                return (False, window_start, current_count)

            # Increment count
            current_count += 1
            self._rate_limits[group_id] = (window_start, current_count)
            return (True, 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)

        """
        async with self._lock:
            if self._global_queue_size >= max_size:
                return (False, self._global_queue_size)
            self._global_queue_size += 1
            return (True, self._global_queue_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.

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

        Returns:
            Tuple of (allowed, current_size)

        """
        async with self._lock:
            if self._global_queue_size + batch_size > max_size:
                return (False, self._global_queue_size)
            self._global_queue_size += batch_size
            return (True, self._global_queue_size)

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

        Returns:
            Current global queue size

        """
        async with self._lock:
            return self._global_queue_size

    async def decrement_global_queue_size(self) -> None:
        """Atomically decrement global queue size."""
        async with self._lock:
            self._global_queue_size = max(0, self._global_queue_size - 1)

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

        Args:
            task_id: Task ID
            group_id: Group ID
            result: Task result to store
            ttl_seconds: Optional TTL (not implemented in in-memory storage)

        """
        async with self._lock:
            self._task_results[task_id] = result

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

        Args:
            task_id: Task ID
            group_id: Optional group ID (unused)

        Returns:
            Task result or None if not found

        """
        async with self._lock:
            return self._task_results.get(task_id)

    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

        """
        async with self._lock:
            # Remove from queue
            if group_id in self._group_queues and task_id in self._group_queues[group_id]:
                self._group_queues[group_id].remove(task_id)

            # Remove task data
            self._tasks.pop(task_id, None)
            self._task_status.pop(task_id, None)
            self._task_results.pop(task_id, None)

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

            # Decrement global queue size
            self._global_queue_size = max(0, self._global_queue_size - 1)

    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

        """
        async with self._lock:
            count = 0
            tasks_to_remove: list[tuple[str, str]] = []

            for task_id, task in self._tasks.items():
                if task.expires_at and task.expires_at < before:
                    tasks_to_remove.append((task_id, task.group_id))

            for task_id, group_id in tasks_to_remove:
                await self.remove_task(task_id, group_id)
                count += 1

            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

        """
        async with self._lock:
            tasks: list[FairTask] = []

            if group_id:
                task_ids = self._dead_letter_queues.get(group_id, [])[offset : offset + limit]
            else:
                # Get from all groups
                all_task_ids: list[str] = []
                for task_ids_list in self._dead_letter_queues.values():
                    all_task_ids.extend(task_ids_list)
                task_ids = all_task_ids[offset : offset + limit]

            for task_id in task_ids:
                task = self._tasks.get(task_id)
                if task:
                    tasks.append(task)

            return tasks

    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

        """
        async with self._lock:
            results: list[tuple[str, str]] = []

            for task_id, task_status in self._task_status.items():
                if task_status == status:
                    task = self._tasks.get(task_id)
                    if task and (group_id is None or task.group_id == group_id):
                        results.append((task_id, task.group_id))

            # Apply pagination
            return results[offset : offset + limit]

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

        Returns:
            True if storage is healthy, False otherwise

        """
        return True

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

        Returns:
            ComponentHealthStatus model with health status details

        """
        async with self._lock:
            total_tasks = len(self._tasks)
            total_groups = len(self._groups)
            total_queue_length = sum(len(q) for q in self._group_queues.values())

            return ComponentHealthStatus(
                healthy=True,
                component="storage",
                details=HealthDetails(
                    custom_data={
                        "total_tasks": total_tasks,
                        "total_groups": total_groups,
                        "total_queue_length": total_queue_length,
                        "global_queue_size": self._global_queue_size,
                    }
                ),
                error=None,
                timestamp=datetime.now(UTC),
            )

    def _update_group_state_internal_locked(
        self,
        group_id: str,
        increment_queue_length: bool,  # noqa: ARG002
    ) -> None:
        """Internal method to update group state queue length (assumes lock is already held).

        Args:
            group_id: Group ID
            increment_queue_length: If True, increment queue length; if False, decrement

        """
        # Lock is already held by caller, so we can directly access state
        queue_length = len(self._group_queues.get(group_id, []))

        if group_id in self._group_states:
            state = self._group_states[group_id]
            state.queue_length = queue_length
        else:
            # Create new state if it doesn't exist
            state = GroupQueueState(
                group_id=group_id,
                queue_length=queue_length,
                virtual_finish_time=0.0,
                weight=1.0,
            )
            self._group_states[group_id] = state

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

        Args:
            group_id: Group ID
            increment_queue_length: If True, increment queue length; if False, decrement

        """
        async with self._lock:
            self._update_group_state_internal_locked(group_id, increment_queue_length)
