GallerySubscriber/backend/app/tasks/maintenance.py

"""System maintenance Celery tasks."""

import asyncio
import logging
import os
import shutil
from datetime import datetime, timedelta
from pathlib import Path

from sqlalchemy import select, update

from app.tasks.db import get_async_session, cleanup_engine as _cleanup_engine
from app.tasks.celery_app import celery_app
from app.models.setting import (
    Setting,
    STORAGE_STATS_SETTING,
    LIBRARY_VALIDATION_REPORT_SETTING,
)
from app.models.download import Download, DownloadStatus
from app.models.base import utcnow

logger = logging.getLogger(__name__)

from app.config import get_settings
settings = get_settings()

# Maximum time a job can be in "running" state before considered orphaned
# This should be longer than the task_time_limit (1 hour) to account for legitimate long runs
STALE_RUNNING_THRESHOLD_MINUTES = 90

# Maximum time a job can be in "queued" state before considered lost
# If a queued job hasn't been picked up in this time, the Celery task was probably lost
# (e.g., due to Redis DB mismatch or Redis restart)
STALE_QUEUED_THRESHOLD_MINUTES = 30


def format_size(size_bytes: int) -> str:
    """Format bytes as human-readable size."""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024:
            return f"{size_bytes:.1f} {unit}"
        size_bytes /= 1024
    return f"{size_bytes:.1f} PB"


def _disk_usage_stats(path: Path) -> dict:
    """Return filesystem-level usage for the volume hosting `path`.

    Uses shutil.disk_usage so we report the actual mount's free space (the
    number that matters for capacity planning), not just what our downloads
    subtree consumes.
    """
    try:
        usage = shutil.disk_usage(path)
        return {
            "disk_total": usage.total,
            "disk_used": usage.used,
            "disk_free": usage.free,
            "disk_total_formatted": format_size(usage.total),
            "disk_used_formatted": format_size(usage.used),
            "disk_free_formatted": format_size(usage.free),
            "disk_percent_used": round(usage.used / usage.total * 100, 1) if usage.total else 0,
        }
    except OSError as e:
        logger.warning(f"Could not read disk usage for {path}: {e}")
        return {}


def calculate_storage_stats() -> dict:
    """Calculate storage statistics for the downloads directory.

    This runs synchronously since it's a filesystem operation.
    """
    downloads_path = Path(settings.download_path)

    if not downloads_path.exists():
        return {
            "total_size": 0,
            "total_size_formatted": "0 B",
            "file_count": 0,
            "calculated_at": datetime.utcnow().isoformat(),
        }

    total_size = 0
    file_count = 0

    try:
        for root, dirs, files in os.walk(downloads_path):
            for f in files:
                fp = os.path.join(root, f)
                try:
                    total_size += os.path.getsize(fp)
                    file_count += 1
                except OSError:
                    pass
    except Exception as e:
        logger.error(f"Error calculating storage stats: {e}")
        return {
            "total_size": 0,
            "total_size_formatted": "Error",
            "file_count": 0,
            "error": str(e),
            "calculated_at": datetime.utcnow().isoformat(),
        }

    return {
        "total_size": total_size,
        "total_size_formatted": format_size(total_size),
        "file_count": file_count,
        "calculated_at": datetime.utcnow().isoformat(),
        **_disk_usage_stats(downloads_path),
    }


async def _update_storage_stats_async() -> dict:
    """Async implementation to calculate and store storage stats."""
    session_factory, engine = get_async_session()

    try:
        # Calculate stats (synchronous filesystem operation)
        stats = calculate_storage_stats()
        logger.info(f"Calculated storage stats: {stats['total_size_formatted']}, {stats['file_count']} files")

        # Store in database
        async with session_factory() as session:
            # Check if setting exists
            result = await session.execute(
                select(Setting).where(Setting.key == STORAGE_STATS_SETTING)
            )
            setting = result.scalar_one_or_none()

            if setting:
                setting.value = stats
            else:
                setting = Setting(key=STORAGE_STATS_SETTING, value=stats)
                session.add(setting)

            await session.commit()

        return stats

    finally:
        await _cleanup_engine(engine)


@celery_app.task(name="tasks.update_storage_stats")
def update_storage_stats() -> dict:
    """Celery task to calculate and cache storage stats.

    This task should be scheduled to run periodically (e.g., every hour).
    """
    logger.info("Starting storage stats calculation...")
    result = asyncio.run(_update_storage_stats_async())
    logger.info(f"Storage stats updated: {result}")
    return result


def scan_library_for_corruption(
    root: Path,
    max_suspect_paths: int = 500,
) -> dict:
    """Walk the library and report files that fail magic-byte validation.

    The user-facing case for this is "I just deployed the validator. How many
    pre-existing files in my library are silently truncated?" — running this
    once gives that number plus a list of paths to act on. We report only;
    the live download path quarantines, but pre-existing files are left in
    place because they may be the only copy and the user should decide.

    Skips the `_quarantine` subtree (those files are already known-bad).
    `suspect_paths` is capped at `max_suspect_paths` so the persisted report
    can't grow unbounded; `suspect_count` always reflects the true total.
    """
    # Local import keeps this module's import cost low (filesystem walk
    # tasks shouldn't drag the validator into every Celery worker process
    # unconditionally).
    from app.services.file_validator import is_validatable, validate_file

    started_at = utcnow().isoformat()
    suspect_paths: list[dict] = []
    suspect_count = 0
    scanned = 0
    by_format: dict[str, int] = {}

    if not root.exists():
        return {
            "started_at": started_at,
            "completed_at": utcnow().isoformat(),
            "scanned": 0,
            "suspect_count": 0,
            "by_format": {},
            "suspect_paths": [],
            "root": str(root),
            "truncated": False,
        }

    quarantine_root = root / "_quarantine"

    for dirpath, dirnames, filenames in os.walk(root):
        # Don't re-scan the quarantine directory.
        dirnames[:] = [d for d in dirnames if Path(dirpath, d) != quarantine_root]

        for name in filenames:
            p = Path(dirpath, name)
            if not is_validatable(p):
                continue
            scanned += 1
            try:
                result = validate_file(p)
            except Exception as e:
                logger.warning(f"Validator raised on {p}: {e}")
                continue
            if result.ok:
                continue
            suspect_count += 1
            fmt = result.format or "unknown"
            by_format[fmt] = by_format.get(fmt, 0) + 1
            if len(suspect_paths) < max_suspect_paths:
                suspect_paths.append(
                    {
                        "path": str(p),
                        "format": result.format,
                        "reason": result.reason,
                        "size": result.size,
                    }
                )

    return {
        "started_at": started_at,
        "completed_at": utcnow().isoformat(),
        "scanned": scanned,
        "suspect_count": suspect_count,
        "by_format": by_format,
        "suspect_paths": suspect_paths,
        "root": str(root),
        "truncated": suspect_count > len(suspect_paths),
    }


async def _validate_library_async() -> dict:
    """Async wrapper that runs the sweep and persists the report."""
    session_factory, engine = get_async_session()
    try:
        report = scan_library_for_corruption(Path(settings.download_path))
        async with session_factory() as session:
            result = await session.execute(
                select(Setting).where(Setting.key == LIBRARY_VALIDATION_REPORT_SETTING)
            )
            setting = result.scalar_one_or_none()
            if setting:
                setting.value = report
            else:
                setting = Setting(key=LIBRARY_VALIDATION_REPORT_SETTING, value=report)
                session.add(setting)
            await session.commit()
        return report
    finally:
        await _cleanup_engine(engine)


@celery_app.task(name="tasks.validate_library")
def validate_library() -> dict:
    """Celery task: walk the library and report truncated files.

    On-demand (UI button or manual trigger). Not in the periodic schedule —
    a full filesystem walk on a large NAS-mounted library is expensive and
    the user should choose when to pay for it.
    """
    logger.info("Starting library validation sweep...")
    report = asyncio.run(_validate_library_async())
    logger.info(
        f"Library validation: scanned {report['scanned']}, "
        f"suspect={report['suspect_count']} ({report['by_format']})"
    )
    return report


async def _reset_orphaned_running_jobs_async(threshold_minutes: int = None) -> dict:
    """Reset jobs stuck in 'running' status back to 'queued'.

    Jobs can get stuck in 'running' status if the worker crashes or restarts
    while processing them. This function detects and resets these orphaned jobs.

    Args:
        threshold_minutes: Jobs running longer than this are considered orphaned.
                          If None, uses STALE_RUNNING_THRESHOLD_MINUTES.

    Returns:
        Dict with count of reset jobs
    """
    if threshold_minutes is None:
        threshold_minutes = STALE_RUNNING_THRESHOLD_MINUTES

    session_factory, engine = get_async_session()

    try:
        async with session_factory() as session:
            now = utcnow()
            cutoff = now - timedelta(minutes=threshold_minutes)

            # Find running jobs that started before the cutoff
            # OR running jobs with no started_at (shouldn't happen but handle it)
            query = select(Download).where(
                Download.status == DownloadStatus.RUNNING,
                (Download.started_at < cutoff) | (Download.started_at == None)
            )

            result = await session.execute(query)
            orphaned_jobs = result.scalars().all()

            reset_count = 0
            for job in orphaned_jobs:
                original_started_at = job.started_at
                job.status = DownloadStatus.QUEUED
                job.started_at = None
                job.error_message = f"Reset from orphaned running state (was running since {original_started_at})"
                reset_count += 1
                logger.info(f"Reset orphaned job {job.id} (URL: {job.url[:50]}...) to QUEUED")

            if reset_count > 0:
                await session.commit()
                logger.info(f"Reset {reset_count} orphaned running jobs to QUEUED")
            else:
                logger.debug("No orphaned running jobs found")

            return {"reset_count": reset_count, "threshold_minutes": threshold_minutes}

    finally:
        await _cleanup_engine(engine)


async def _reset_all_running_jobs_async() -> dict:
    """Reset ALL jobs in 'running' status back to 'queued'.

    This is useful on worker startup to clean up any jobs that were
    interrupted by a previous shutdown/crash.

    Returns:
        Dict with count of reset jobs
    """
    session_factory, engine = get_async_session()

    try:
        async with session_factory() as session:
            # Find all running jobs
            query = select(Download).where(Download.status == DownloadStatus.RUNNING)
            result = await session.execute(query)
            running_jobs = result.scalars().all()

            reset_count = 0
            for job in running_jobs:
                job.status = DownloadStatus.QUEUED
                job.started_at = None
                job.error_message = "Reset on worker startup"
                reset_count += 1
                logger.info(f"Reset running job {job.id} to QUEUED on startup")

            if reset_count > 0:
                await session.commit()
                logger.info(f"Reset {reset_count} running jobs to QUEUED on worker startup")

            return {"reset_count": reset_count}

    finally:
        await _cleanup_engine(engine)


@celery_app.task(name="tasks.reset_orphaned_jobs")
def reset_orphaned_jobs(threshold_minutes: int = None) -> dict:
    """Celery task to reset orphaned running jobs.

    Jobs that have been in 'running' status for longer than the threshold
    are considered orphaned (worker probably crashed) and reset to 'queued'.
    """
    logger.info("Checking for orphaned running jobs...")
    result = asyncio.run(_reset_orphaned_running_jobs_async(threshold_minutes))
    return result


def reset_all_running_jobs_sync() -> dict:
    """Synchronous function to reset all running jobs.

    Called on worker startup before Celery is fully initialized.
    """
    logger.info("Resetting all running jobs on worker startup...")
    result = asyncio.run(_reset_all_running_jobs_async())
    return result


async def _requeue_stale_queued_jobs_async(threshold_minutes: int = None) -> dict:
    """Re-queue Celery tasks for downloads stuck in 'queued' status.

    Downloads can get stuck in 'queued' status if the Celery task was lost
    (e.g., sent to wrong Redis DB, Redis restart, etc.). This function
    detects these stale queued jobs and re-sends the Celery tasks.

    Only one download per source is re-queued to prevent duplicate work.

    Args:
        threshold_minutes: Jobs queued longer than this are considered stale.
                          If None, uses STALE_QUEUED_THRESHOLD_MINUTES.

    Returns:
        Dict with count of re-queued jobs
    """
    if threshold_minutes is None:
        threshold_minutes = STALE_QUEUED_THRESHOLD_MINUTES

    # Import here to avoid circular imports
    from app.tasks.downloads import download_source

    session_factory, engine = get_async_session()

    try:
        async with session_factory() as session:
            now = utcnow()
            cutoff = now - timedelta(minutes=threshold_minutes)

            # Find queued jobs created before the cutoff, ordered by creation time
            # so we process older jobs first
            query = select(Download).where(
                Download.status == DownloadStatus.QUEUED,
                Download.created_at < cutoff
            ).order_by(Download.created_at.asc())

            result = await session.execute(query)
            stale_jobs = result.scalars().all()

            # Track which sources we've already queued to avoid duplicates
            queued_source_ids = set()
            requeued_count = 0
            skipped_duplicates = 0

            for job in stale_jobs:
                # Skip if we've already queued a job for this source
                if job.source_id in queued_source_ids:
                    # Mark duplicate as skipped so it won't be checked again
                    job.status = DownloadStatus.SKIPPED
                    job.error_message = "Duplicate queued job - another job for this source was re-queued"
                    skipped_duplicates += 1
                    logger.info(f"Marked duplicate job {job.id} as SKIPPED (source {job.source_id} already re-queued)")
                    continue

                try:
                    # Re-send the Celery task
                    download_source.delay(job.source_id, download_id=job.id)
                    queued_source_ids.add(job.source_id)
                    requeued_count += 1
                    logger.info(f"Re-queued stale job {job.id} (source_id: {job.source_id}, queued since: {job.created_at})")
                except Exception as e:
                    logger.error(f"Failed to re-queue job {job.id}: {e}")

            # Commit the status changes for skipped duplicates
            if skipped_duplicates > 0:
                await session.commit()

            if requeued_count > 0:
                logger.info(f"Re-queued {requeued_count} stale queued jobs (marked {skipped_duplicates} duplicates as SKIPPED)")
            else:
                logger.debug("No stale queued jobs found")

            return {"requeued_count": requeued_count, "skipped_duplicates": skipped_duplicates, "threshold_minutes": threshold_minutes}

    finally:
        await _cleanup_engine(engine)


async def _requeue_all_queued_jobs_async() -> dict:
    """Re-queue Celery tasks for ALL downloads in 'queued' status.

    This is useful on scheduler startup to recover from Redis DB mismatch
    or other issues that caused Celery tasks to be lost.

    Only one download per source is re-queued to prevent duplicate work.

    Returns:
        Dict with count of re-queued jobs
    """
    # Import here to avoid circular imports
    from app.tasks.downloads import download_source

    session_factory, engine = get_async_session()

    try:
        async with session_factory() as session:
            # Find all queued jobs, ordered by creation time (oldest first)
            query = select(Download).where(
                Download.status == DownloadStatus.QUEUED
            ).order_by(Download.created_at.asc())
            result = await session.execute(query)
            queued_jobs = result.scalars().all()

            # Track which sources we've already queued to avoid duplicates
            queued_source_ids = set()
            requeued_count = 0
            skipped_duplicates = 0

            for job in queued_jobs:
                # Skip if we've already queued a job for this source
                if job.source_id in queued_source_ids:
                    # Mark duplicate as skipped so it won't be checked again
                    job.status = DownloadStatus.SKIPPED
                    job.error_message = "Duplicate queued job - another job for this source was re-queued"
                    skipped_duplicates += 1
                    logger.info(f"Marked duplicate job {job.id} as SKIPPED (source {job.source_id} already re-queued)")
                    continue

                try:
                    download_source.delay(job.source_id, download_id=job.id)
                    queued_source_ids.add(job.source_id)
                    requeued_count += 1
                    logger.info(f"Re-queued job {job.id} (source_id: {job.source_id}) on startup")
                except Exception as e:
                    logger.error(f"Failed to re-queue job {job.id}: {e}")

            # Commit the status changes for skipped duplicates
            if skipped_duplicates > 0:
                await session.commit()

            if requeued_count > 0:
                logger.info(f"Re-queued {requeued_count} queued jobs on startup (marked {skipped_duplicates} duplicates as SKIPPED)")

            return {"requeued_count": requeued_count, "skipped_duplicates": skipped_duplicates}

    finally:
        await _cleanup_engine(engine)


@celery_app.task(name="tasks.requeue_stale_jobs")
def requeue_stale_jobs(threshold_minutes: int = None) -> dict:
    """Celery task to re-queue stale queued jobs.

    Jobs that have been in 'queued' status for longer than the threshold
    are considered lost (Celery task missing) and re-queued.
    """
    logger.info("Checking for stale queued jobs...")
    result = asyncio.run(_requeue_stale_queued_jobs_async(threshold_minutes))
    return result


def requeue_all_queued_jobs_sync() -> dict:
    """Synchronous function to re-queue all queued jobs.

    Called on scheduler startup to recover lost Celery tasks.
    """
    logger.info("Re-queuing all queued jobs on scheduler startup...")
    result = asyncio.run(_requeue_all_queued_jobs_async())
    return result