"""
In-memory only queue tracking models.
These models are NOT persisted to disk and reset on application restart.
Used for tracking transient queue status without storing PII in unencrypted databases.
"""

from datetime import datetime, UTC
from functools import partial
from typing import Dict, Optional

from sqlalchemy import Column, DateTime, Integer, String, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session, sessionmaker
from sqlalchemy.pool import StaticPool

from loguru import logger

Base = declarative_base()


class InMemoryQueueStatus(Base):
    """
    Track queue status for each user IN MEMORY ONLY.
    This data is transient and will be lost on application restart.
    """

    __tablename__ = "memory_queue_status"

    username = Column(String, primary_key=True)
    active_tasks = Column(Integer, default=0)
    queued_tasks = Column(Integer, default=0)
    last_checked = Column(DateTime, default=partial(datetime.now, UTC))
    last_task_id = Column(String, nullable=True)


class InMemoryTaskMetadata(Base):
    """
    Store task metadata for background processing IN MEMORY ONLY.
    This data is transient and will be lost on application restart.
    """

    __tablename__ = "memory_task_metadata"

    task_id = Column(String, primary_key=True)
    username = Column(String, nullable=False, index=True)
    status = Column(
        String, nullable=False
    )  # queued, processing, completed, failed
    task_type = Column(String, nullable=False)  # research, benchmark, etc
    created_at = Column(DateTime, default=partial(datetime.now, UTC))
    started_at = Column(DateTime, nullable=True)
    completed_at = Column(DateTime, nullable=True)
    # Note: We don't store error messages to avoid PII leakage


class InMemoryQueueTracker:
    """
    Manages in-memory only queue tracking.
    All data is transient and resets on application restart.
    """

    def __init__(self):
        """Initialize in-memory SQLite database."""
        # Create in-memory SQLite database
        # Using StaticPool to ensure thread safety with a single connection
        self.engine = create_engine(
            "sqlite:///:memory:",
            connect_args={"check_same_thread": False},
            poolclass=StaticPool,
            echo=False,
        )

        # Create tables
        Base.metadata.create_all(self.engine)

        # Create session factory
        self.SessionLocal = sessionmaker(bind=self.engine)

        logger.info("In-memory queue tracker initialized (data is transient)")

    def get_session(self) -> Session:
        """Get a new database session."""
        return self.SessionLocal()

    def update_queue_status(
        self, username: str, active_tasks: int, queued_tasks: int
    ) -> None:
        """Update queue status for a user."""
        with self.get_session() as session:
            status = (
                session.query(InMemoryQueueStatus)
                .filter_by(username=username)
                .first()
            )

            if status:
                status.active_tasks = active_tasks
                status.queued_tasks = queued_tasks
                status.last_checked = datetime.now(UTC)
            else:
                status = InMemoryQueueStatus(
                    username=username,
                    active_tasks=active_tasks,
                    queued_tasks=queued_tasks,
                )
                session.add(status)

            session.commit()

    def get_queue_status(self, username: str) -> Optional[Dict[str, any]]:
        """Get queue status for a user."""
        with self.get_session() as session:
            status = (
                session.query(InMemoryQueueStatus)
                .filter_by(username=username)
                .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(
        self,
        task_id: str,
        username: str,
        task_type: str,
    ) -> None:
        """Add a new task to tracking."""
        with self.get_session() as session:
            task = InMemoryTaskMetadata(
                task_id=task_id,
                username=username,
                status="queued",
                task_type=task_type,
            )
            session.add(task)

            # Update queue status
            status = (
                session.query(InMemoryQueueStatus)
                .filter_by(username=username)
                .first()
            )
            if status:
                status.queued_tasks += 1
                status.last_task_id = task_id
            else:
                status = InMemoryQueueStatus(
                    username=username, queued_tasks=1, last_task_id=task_id
                )
                session.add(status)

            session.commit()

    def update_task_status(self, task_id: str, status: str) -> None:
        """Update task status."""
        with self.get_session() as session:
            task = (
                session.query(InMemoryTaskMetadata)
                .filter_by(task_id=task_id)
                .first()
            )

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

                if status == "processing" and old_status == "queued":
                    task.started_at = datetime.now(UTC)
                    # Update queue counts
                    queue_status = (
                        session.query(InMemoryQueueStatus)
                        .filter_by(username=task.username)
                        .first()
                    )
                    if queue_status:
                        queue_status.queued_tasks = max(
                            0, queue_status.queued_tasks - 1
                        )
                        queue_status.active_tasks += 1

                elif status in ["completed", "failed"]:
                    task.completed_at = datetime.now(UTC)
                    # Update queue counts
                    queue_status = (
                        session.query(InMemoryQueueStatus)
                        .filter_by(username=task.username)
                        .first()
                    )
                    if queue_status:
                        queue_status.active_tasks = max(
                            0, queue_status.active_tasks - 1
                        )

                session.commit()

    def get_user_tasks(
        self, username: str, status: Optional[str] = None
    ) -> list:
        """Get tasks for a user, optionally filtered by status."""
        with self.get_session() as session:
            query = session.query(InMemoryTaskMetadata).filter_by(
                username=username
            )

            if status:
                query = query.filter_by(status=status)

            tasks = query.order_by(InMemoryTaskMetadata.created_at).all()

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

    def cleanup_completed_tasks(self, username: Optional[str] = None) -> int:
        """
        Remove completed/failed tasks from memory.
        Useful to prevent memory growth over time.
        """
        with self.get_session() as session:
            query = session.query(InMemoryTaskMetadata).filter(
                InMemoryTaskMetadata.status.in_(["completed", "failed"])
            )

            if username:
                query = query.filter_by(username=username)

            deleted = query.delete()
            session.commit()

            if deleted > 0:
                logger.debug(
                    f"Cleaned up {deleted} completed tasks from memory"
                )

            return deleted

    def reset_user_queues(self, username: str) -> None:
        """Reset all queue data for a user (useful for testing or cleanup)."""
        with self.get_session() as session:
            # Delete all tasks for user
            session.query(InMemoryTaskMetadata).filter_by(
                username=username
            ).delete()

            # Reset or delete queue status
            status = (
                session.query(InMemoryQueueStatus)
                .filter_by(username=username)
                .first()
            )
            if status:
                status.active_tasks = 0
                status.queued_tasks = 0
                status.last_task_id = None
                status.last_checked = datetime.now(UTC)

            session.commit()
            logger.info(f"Reset queue data for user: {username}")


# Global instance - this is in-memory only and will reset on restart
_memory_queue_tracker = None


def get_memory_queue_tracker() -> InMemoryQueueTracker:
    """Get global in-memory queue tracker instance."""
    global _memory_queue_tracker
    if _memory_queue_tracker is None:
        _memory_queue_tracker = InMemoryQueueTracker()
    return _memory_queue_tracker
