imagerepo/app/tasks/scan.py

# app/tasks/scan.py
"""
Directory scanning task for the import pipeline.

This task walks the import directory, detects new/changed files,
creates ImportTask records, and queues appropriate processing tasks.
"""
import os
import logging
from datetime import datetime

from app.celery_app import celery
from app import db
from app.models import ImportTask, ImportBatch

log = logging.getLogger('celery.tasks.scan')

# Supported file extensions (imported from image_importer for consistency)
ALLOWED_MEDIA_EXTS = (
    '.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp', '.tiff', '.tif',
    '.mp4', '.mov', '.webm', '.avi', '.mkv'
)
ALLOWED_ARCHIVE_EXTS = (
    '.zip', '.rar', '.7z', '.tar', '.tar.gz', '.tgz', '.tar.bz2', '.tbz2',
    '.cbr', '.cbz'
)


@celery.task(bind=True, name='app.tasks.scan.scan_directory')
def scan_directory(self, source_dir: str = '/import', dest_dir: str = '/images'):
    """
    Scan source directory for new/changed files.

    Creates ImportTask records for each file and queues them for processing.

    Strategy:
    1. Walk directory tree under each artist folder
    2. For each file, check if already in ImportTask with same path+size
    3. If new/changed, create ImportTask and queue appropriate task
    4. Handle archives specially (first volume only, queue as atomic unit)

    Args:
        source_dir: Root directory to scan (default: /import)
        dest_dir: Destination for imported files (default: /images)

    Returns:
        dict with batch_id and statistics
    """
    from app.tasks.import_file import import_media_file, import_archive
    from app.utils.image_importer import is_first_volume

    log.info(f"Starting directory scan: {source_dir}")

    # Create batch record
    batch = ImportBatch(source_directory=source_dir)
    db.session.add(batch)
    db.session.commit()

    stats = {
        'total_files': 0,
        'queued_media': 0,
        'queued_archives': 0,
        'skipped_existing': 0,
        'skipped_non_first_volume': 0,
        'skipped_non_media': 0,
    }

    try:
        # Iterate through artist directories
        for artist_dir in os.listdir(source_dir):
            artist_path = os.path.join(source_dir, artist_dir)
            if not os.path.isdir(artist_path):
                continue

            log.debug(f"Scanning artist directory: {artist_dir}")

            # Walk all files in artist directory (including subdirectories)
            for root, _, files in os.walk(artist_path):
                for filename in files:
                    full_path = os.path.join(root, filename)
                    lower = filename.lower()
                    stats['total_files'] += 1

                    # Skip sidecar JSON files (handled during import)
                    if lower.endswith('.json'):
                        continue

                    # Check if archive
                    if lower.endswith(ALLOWED_ARCHIVE_EXTS):
                        # Only process first volume of multi-part archives
                        if not is_first_volume(full_path):
                            stats['skipped_non_first_volume'] += 1
                            continue

                        # Create and queue archive task
                        task_record = _create_task_if_new(
                            source_path=full_path,
                            task_type='import_archive',
                            batch_id=batch.id,
                            context={'artist': artist_dir, 'dest_dir': dest_dir}
                        )

                        if task_record:
                            result = import_archive.delay(task_record.id)
                            task_record.celery_task_id = result.id
                            task_record.status = 'queued'
                            db.session.commit()
                            stats['queued_archives'] += 1
                            log.debug(f"Queued archive: {filename}")
                        else:
                            stats['skipped_existing'] += 1
                        continue

                    # Check if media file
                    if not lower.endswith(ALLOWED_MEDIA_EXTS):
                        stats['skipped_non_media'] += 1
                        continue

                    # Create and queue media file task
                    task_record = _create_task_if_new(
                        source_path=full_path,
                        task_type='import_image',
                        batch_id=batch.id,
                        context={'artist': artist_dir, 'dest_dir': dest_dir}
                    )

                    if task_record:
                        result = import_media_file.delay(task_record.id)
                        task_record.celery_task_id = result.id
                        task_record.status = 'queued'
                        db.session.commit()
                        stats['queued_media'] += 1
                        log.debug(f"Queued media: {filename}")
                    else:
                        stats['skipped_existing'] += 1

        # Update batch with totals
        batch.total_files = stats['total_files']
        db.session.commit()

        log.info(f"Scan complete. Queued {stats['queued_media']} media, "
                 f"{stats['queued_archives']} archives. "
                 f"Skipped {stats['skipped_existing']} existing.")

        return {
            'batch_id': batch.id,
            'stats': stats
        }

    except Exception as e:
        log.error(f"Scan failed: {e}", exc_info=True)
        batch.status = 'failed'
        db.session.commit()
        raise


def _create_task_if_new(source_path: str, task_type: str, batch_id: int, context: dict) -> ImportTask:
    """
    Create ImportTask if no existing task for this file.

    Uses file path and size for change detection.
    Hash is computed later during actual processing.

    Args:
        source_path: Full path to source file
        task_type: Type of task ('import_image', 'import_archive', etc.)
        batch_id: ID of the current ImportBatch
        context: Context dict with artist, dest_dir, etc.

    Returns:
        New ImportTask record if created, None if already exists
    """
    try:
        file_size = os.path.getsize(source_path)
    except OSError as e:
        log.warning(f"Cannot access file {source_path}: {e}")
        return None

    # Check for existing pending/queued/processing task for same file
    existing_active = ImportTask.query.filter_by(
        source_path=source_path,
        task_type=task_type
    ).filter(
        ImportTask.status.in_(['pending', 'queued', 'processing'])
    ).first()

    if existing_active:
        log.debug(f"Task already active for: {source_path}")
        return None

    # Check for completed task with same file path and size
    # If size changed, we should re-import
    completed_same = ImportTask.query.filter_by(
        source_path=source_path,
        task_type=task_type,
        file_size=file_size,
        status='complete'
    ).first()

    if completed_same:
        log.debug(f"Already imported (same size): {source_path}")
        return None

    # Also check 'skipped' status - don't re-queue skipped files unless size changed
    skipped_same = ImportTask.query.filter_by(
        source_path=source_path,
        task_type=task_type,
        file_size=file_size,
        status='skipped'
    ).first()

    if skipped_same:
        log.debug(f"Previously skipped (same size): {source_path}")
        return None

    # Create new task
    task = ImportTask(
        source_path=source_path,
        task_type=task_type,
        file_size=file_size,
        batch_id=batch_id,
        context=context,
        status='pending'
    )
    db.session.add(task)
    db.session.commit()

    return task


@celery.task(name='app.tasks.scan.recover_interrupted_tasks')
def recover_interrupted_tasks():
    """
    Find and re-queue tasks that were interrupted.

    Called periodically by Celery Beat to recover from worker crashes.
    Tasks with status 'processing' or 'queued' that have no active
    Celery task are re-queued.

    Returns:
        dict with count of recovered tasks
    """
    from app.tasks.import_file import import_media_file, import_archive
    from app.tasks.thumbnail import generate_thumbnail_task

    log.info("Checking for interrupted tasks...")

    # Find tasks that are stuck in processing/queued state
    # We look for tasks that have been in this state for more than 10 minutes
    cutoff = datetime.utcnow()

    interrupted = ImportTask.query.filter(
        ImportTask.status.in_(['processing', 'queued'])
    ).all()

    recovered = 0
    for task in interrupted:
        # Check if Celery task is still active
        if task.celery_task_id:
            from app.celery_app import celery as celery_app
            result = celery_app.AsyncResult(task.celery_task_id)
            # If task is still pending/started, don't re-queue
            if result.state in ('PENDING', 'STARTED', 'RETRY'):
                continue

        log.info(f"Recovering task {task.id}: {task.source_path}")

        # Reset task state
        task.status = 'pending'
        task.celery_task_id = None
        task.started_at = None
        task.retry_count += 1

        # Re-queue based on type
        if task.task_type == 'import_image':
            result = import_media_file.delay(task.id)
        elif task.task_type == 'import_archive':
            result = import_archive.delay(task.id)
        elif task.task_type == 'thumbnail':
            # Thumbnail tasks store image_id in context
            image_id = task.context.get('image_id') if task.context else None
            if image_id:
                result = generate_thumbnail_task.delay(image_id)
            else:
                task.status = 'failed'
                task.error_message = 'Missing image_id in context'
                db.session.commit()
                continue
        else:
            log.warning(f"Unknown task type: {task.task_type}")
            continue

        task.celery_task_id = result.id
        task.status = 'queued'
        recovered += 1

    db.session.commit()

    if recovered > 0:
        log.info(f"Recovered {recovered} interrupted tasks")

    return {'recovered': recovered}


@celery.task(name='app.tasks.scan.cleanup_old_tasks')
def cleanup_old_tasks(days: int = 7):
    """
    Clean up old failed and skipped import task records.

    Called periodically by Celery Beat to prevent database bloat.
    Deletes ImportTask records with status 'failed' or 'skipped'
    that are older than the specified number of days.

    Also cleans up orphaned ImportBatch records that have no tasks.

    Args:
        days: Number of days to retain failed/skipped tasks (default: 7)

    Returns:
        dict with counts of deleted records
    """
    from datetime import timedelta

    cutoff = datetime.utcnow() - timedelta(days=days)

    log.info(f"Cleaning up failed/skipped tasks older than {days} days (before {cutoff})")

    # Delete old failed tasks
    failed_deleted = ImportTask.query.filter(
        ImportTask.status == 'failed',
        ImportTask.completed_at < cutoff
    ).delete(synchronize_session=False)

    # Delete old skipped tasks
    skipped_deleted = ImportTask.query.filter(
        ImportTask.status == 'skipped',
        ImportTask.completed_at < cutoff
    ).delete(synchronize_session=False)

    db.session.commit()

    # Clean up orphaned batches (batches with no tasks)
    orphaned_batches = ImportBatch.query.filter(
        ~ImportBatch.tasks.any()
    ).all()
    batches_deleted = len(orphaned_batches)
    for batch in orphaned_batches:
        db.session.delete(batch)

    db.session.commit()

    if failed_deleted or skipped_deleted or batches_deleted:
        log.info(f"Cleanup complete: {failed_deleted} failed tasks, "
                 f"{skipped_deleted} skipped tasks, {batches_deleted} orphaned batches deleted")

    return {
        'failed_deleted': failed_deleted,
        'skipped_deleted': skipped_deleted,
        'batches_deleted': batches_deleted
    }


@celery.task(name='app.tasks.scan.update_batch_stats')
def update_batch_stats(batch_id: int):
    """
    Update statistics for an ImportBatch based on its tasks.

    Args:
        batch_id: ID of the batch to update
    """
    batch = ImportBatch.query.get(batch_id)
    if not batch:
        return

    from sqlalchemy import func

    # Count tasks by status
    status_counts = dict(
        db.session.query(
            ImportTask.status,
            func.count(ImportTask.id)
        ).filter(
            ImportTask.batch_id == batch_id
        ).group_by(ImportTask.status).all()
    )

    batch.processed_files = sum(
        status_counts.get(s, 0)
        for s in ['complete', 'skipped', 'failed']
    )
    batch.imported_count = status_counts.get('complete', 0)
    batch.skipped_count = status_counts.get('skipped', 0)
    batch.failed_count = status_counts.get('failed', 0)

    # Check if batch is complete
    pending_or_active = sum(
        status_counts.get(s, 0)
        for s in ['pending', 'queued', 'processing']
    )

    if pending_or_active == 0 and batch.status == 'running':
        batch.status = 'complete'
        batch.completed_at = datetime.utcnow()
        log.info(f"Batch {batch_id} complete: {batch.imported_count} imported, "
                 f"{batch.skipped_count} skipped, {batch.failed_count} failed")

    db.session.commit()