"""
Queue service for managing tasks using encrypted user databases.
Replaces the service_db approach with direct access to user databases.
"""

from datetime import datetime, timedelta, UTC
from typing import Any, Dict, List, Optional

from loguru import logger
from sqlalchemy.orm import Session

from .models import QueueStatus, TaskMetadata


class UserQueueService:
    """Manages queue operations within a user's encrypted database."""

    def __init__(self, session: Session):
        """
        Initialize with a database session.

        Args:
            session: SQLAlchemy session for the user's encrypted database
        """
        self.session = session

    def _safe_commit(self) -> None:
        """Commit the current transaction, rolling back on failure."""
        try:
            self.session.commit()
        except Exception:
            self.session.rollback()
            logger.exception("Database commit failed, transaction rolled back")
            raise

    def update_queue_status(
        self,
        active_tasks: int,
        queued_tasks: int,
        last_task_id: Optional[str] = None,
    ) -> None:
        """Update queue status for the user."""
        status = self.session.query(QueueStatus).first()

        if status:
            status.active_tasks = active_tasks
            status.queued_tasks = queued_tasks
            status.last_checked = datetime.now(UTC)
            if last_task_id:
                status.last_task_id = last_task_id
        else:
            status = QueueStatus(
                active_tasks=active_tasks,
                queued_tasks=queued_tasks,
                last_task_id=last_task_id,
            )
            self.session.add(status)

        self._safe_commit()

    def get_queue_status(self) -> Optional[Dict[str, Any]]:
        """Get queue status for the user."""
        status = self.session.query(QueueStatus).first()

        if status:
            return {
                "active_tasks": status.active_tasks,
                "queued_tasks": status.queued_tasks,
                "last_checked": status.last_checked,
                "last_task_id": status.last_task_id,
            }
        return None

    def add_task_metadata(
        self,
        task_id: str,
        task_type: str,
        priority: int = 0,
    ) -> None:
        """Add metadata for a new task."""
        task = TaskMetadata(
            task_id=task_id,
            status="queued",
            task_type=task_type,
            priority=priority,
        )
        self.session.add(task)

        # Update queue counts
        self._increment_queue_count()

        self._safe_commit()

    def update_task_status(
        self, task_id: str, status: str, error_message: Optional[str] = None
    ) -> None:
        """Update task status."""
        task = (
            self.session.query(TaskMetadata).filter_by(task_id=task_id).first()
        )

        if task:
            old_status = task.status
            task.status = status
            task.error_message = error_message

            if status == "processing" and old_status == "queued":
                task.started_at = datetime.now(UTC)
                self._update_queue_counts(-1, 1)  # -1 queued, +1 active

            elif status == "queued" and old_status == "processing":
                # Dispatch was rolled back (e.g. global capacity reject left
                # the research queued for the next tick). Revert the claim so
                # the counters stay balanced instead of leaking a slot into
                # active_tasks on every retry.
                task.started_at = None
                self._update_queue_counts(1, -1)  # +1 queued, -1 active

            elif status in ["completed", "failed"]:
                task.completed_at = datetime.now(UTC)
                self._update_queue_counts(0, -1)  # 0 queued, -1 active

            self._safe_commit()

    def get_pending_tasks(self, limit: int = 50) -> List[Dict[str, Any]]:
        """Get pending tasks for the user."""
        tasks = (
            self.session.query(TaskMetadata)
            .filter_by(status="queued")
            .order_by(TaskMetadata.priority.desc(), TaskMetadata.created_at)
            .limit(limit)
            .all()
        )

        return [
            {
                "task_id": t.task_id,
                "task_type": t.task_type,
                "created_at": t.created_at,
                "priority": t.priority,
            }
            for t in tasks
        ]

    def cleanup_old_tasks(self, days: int = 7) -> int:
        """Clean up old completed/failed tasks."""
        cutoff_date = datetime.now(UTC) - timedelta(days=days)

        deleted = (
            self.session.query(TaskMetadata)
            .filter(
                TaskMetadata.status.in_(["completed", "failed"]),
                TaskMetadata.completed_at < cutoff_date,
            )
            .delete()
        )

        self._safe_commit()
        return deleted

    def _get_or_create_status(self) -> QueueStatus:
        """Get existing queue status or create a new one with zero counts."""
        status = self.session.query(QueueStatus).first()
        if status is None:
            status = QueueStatus(queued_tasks=0, active_tasks=0)
            self.session.add(status)
        return status

    def _increment_queue_count(self):
        """Increment the queued task count."""
        status = self._get_or_create_status()
        status.queued_tasks += 1
        status.last_checked = datetime.now(UTC)

    def _update_queue_counts(self, queued_delta: int, active_delta: int):
        """Update queue counts by deltas."""
        status = self._get_or_create_status()
        status.queued_tasks = max(0, status.queued_tasks + queued_delta)
        status.active_tasks = max(0, status.active_tasks + active_delta)
        status.last_checked = datetime.now(UTC)
