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.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()