imagerepo/app/tasks/import_file.py

# app/tasks/import_file.py
"""
File import tasks for images, videos, and archives.

These tasks handle the actual import of individual files,
including filtering, duplicate detection, and storage.
"""
import os
import shutil
import tempfile
import logging
from datetime import datetime

import imagehash

from app.celery_app import celery
from app import db
from app.models import ImportTask, ImageRecord, Tag, ArchiveRecord

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


@celery.task(bind=True, name='app.tasks.import_file.import_media_file',
             max_retries=3, default_retry_delay=60)
def import_media_file(self, task_id: int):
    """
    Import a single media file (image or video).

    Steps:
    1. Load ImportTask record
    2. Apply filters (dimensions, transparency, single-color)
    3. Check for duplicates (hash and pHash)
    4. Copy to content-addressed storage
    5. Create ImageRecord
    6. Queue thumbnail task
    7. Queue sidecar metadata task
    8. Update ImportTask status

    Args:
        task_id: ID of the ImportTask record

    Returns:
        dict with status and result info
    """
    from app.tasks.thumbnail import generate_thumbnail_task
    from app.tasks.sidecar import apply_sidecar_metadata
    from app.utils.image_importer import (
        extract_metadata, calculate_hash, calculate_perceptual_hash,
        is_mostly_transparent, is_mostly_single_color, find_similar_image,
        build_hashed_dest_path, load_import_settings, supersede_image,
        generate_thumbnail, generate_video_thumbnail_mirrored
    )
    from app.utils.metadata_enrichment import find_sidecar_json

    task = ImportTask.query.get(task_id)
    if not task:
        log.error(f"Task {task_id} not found")
        return {'error': 'Task not found'}

    task.status = 'processing'
    task.started_at = datetime.utcnow()
    db.session.commit()

    try:
        src_path = task.source_path
        context = task.context or {}
        artist = context.get('artist', 'unknown')
        dest_dir = context.get('dest_dir', '/images')

        log.info(f"Processing: {src_path}")

        # Load settings
        settings = load_import_settings()

        # Validate file exists
        if not os.path.exists(src_path):
            return _fail_task(task, 'Source file not found')

        filename = os.path.basename(src_path)
        lower = filename.lower()

        # Extract metadata first for filtering
        metadata = extract_metadata(src_path)
        if not metadata.get('width') or not metadata.get('height'):
            return _skip_task(task, 'Missing dimension data')

        # Apply dimension filter
        min_width = settings.get('min_width', 0)
        min_height = settings.get('min_height', 0)
        if min_width > 0 and metadata['width'] < min_width:
            return _skip_task(task, f"Too small: {metadata['width']}x{metadata['height']}")
        if min_height > 0 and metadata['height'] < min_height:
            return _skip_task(task, f"Too small: {metadata['width']}x{metadata['height']}")

        # Apply transparency filter
        if settings.get('skip_transparent') and lower.endswith(('.png', '.gif', '.webp')):
            threshold = settings.get('transparency_threshold', 0.9)
            if is_mostly_transparent(src_path, threshold):
                return _skip_task(task, 'Mostly transparent')

        # Apply single-color filter
        if settings.get('skip_single_color', True) and lower.endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp', '.bmp')):
            if is_mostly_single_color(src_path,
                                      settings.get('single_color_threshold', 0.95),
                                      settings.get('single_color_tolerance', 30)):
                return _skip_task(task, 'Mostly single color')

        # Compute content hash
        file_hash = calculate_hash(src_path)
        task.file_hash = file_hash

        # Check for exact duplicate by hash
        existing = ImageRecord.query.filter_by(hash=file_hash).first()
        if existing:
            # Still try to enrich with sidecar metadata
            sidecar_path = find_sidecar_json(src_path)
            if sidecar_path:
                apply_sidecar_metadata.delay(existing.id, sidecar_path)
            return _skip_task(task, 'Duplicate by hash', result_image_id=existing.id)

        # Compute pHash for similarity check
        phash = calculate_perceptual_hash(src_path)
        phash_threshold = settings.get('phash_threshold', 10)

        if phash and phash_threshold > 0:
            # Load existing pHashes for comparison
            existing_phashes = [
                (imagehash.hex_to_hash(img.perceptual_hash), img.width, img.height, img.id)
                for img in ImageRecord.query.filter(ImageRecord.perceptual_hash.isnot(None)).all()
            ]

            relationship, similar_id = find_similar_image(
                phash, metadata['width'], metadata['height'],
                existing_phashes, threshold=phash_threshold
            )

            if relationship == 'larger_exists':
                return _skip_task(task, 'Similar larger image exists')

            # Handle supersede case - new image is larger than existing
            if relationship == 'smaller_exists' and settings.get('supersede_smaller', True) and similar_id:
                old_record = ImageRecord.query.get(similar_id)
                if old_record:
                    return _supersede_existing(
                        task, old_record, src_path, file_hash, phash,
                        metadata, artist, dest_dir, filename
                    )

        # Normal import - copy to destination
        target_dir = os.path.join(dest_dir, artist)
        os.makedirs(target_dir, exist_ok=True)
        dest_path = build_hashed_dest_path(target_dir, filename, file_hash)
        shutil.copy2(src_path, dest_path)

        # Generate thumbnail
        try:
            if dest_path.lower().endswith(('.mp4', '.mov', '.webm', '.avi', '.mkv')):
                thumb_path = generate_video_thumbnail_mirrored(dest_path)
            else:
                thumb_path = generate_thumbnail(dest_path)
        except Exception as e:
            log.warning(f"Thumbnail generation failed for {filename}: {e}")
            thumb_path = None

        # Create ImageRecord
        record = ImageRecord(
            filename=filename,
            filepath=dest_path,
            thumb_path=thumb_path,
            hash=file_hash,
            perceptual_hash=str(phash) if phash else None,
            file_size=metadata['file_size'],
            width=metadata['width'],
            height=metadata['height'],
            format=metadata['format'],
            camera_model=metadata.get('camera_model'),
            taken_at=metadata.get('taken_at'),
            imported_at=datetime.utcnow()
        )

        # Add artist tag
        _add_artist_tag(record, artist)

        db.session.add(record)
        db.session.commit()

        # Complete task
        task.status = 'complete'
        task.result_image_id = record.id
        task.completed_at = datetime.utcnow()
        db.session.commit()

        # Queue sidecar metadata enrichment
        sidecar_path = find_sidecar_json(src_path)
        if sidecar_path:
            apply_sidecar_metadata.delay(record.id, sidecar_path)

        # Update batch stats
        if task.batch_id:
            from app.tasks.scan import update_batch_stats
            update_batch_stats.delay(task.batch_id)

        log.info(f"Imported: {filename} -> {dest_path}")

        return {
            'status': 'imported',
            'image_id': record.id,
            'dest_path': dest_path
        }

    except Exception as e:
        db.session.rollback()
        log.error(f"Import failed for task {task_id}: {e}", exc_info=True)

        task.retry_count += 1
        if task.retry_count < 3:
            task.status = 'queued'
            db.session.commit()
            raise self.retry(exc=e)
        else:
            return _fail_task(task, str(e))


@celery.task(bind=True, name='app.tasks.import_file.import_archive',
             max_retries=2, default_retry_delay=120,
             soft_time_limit=1800, time_limit=3600)  # 30min soft, 60min hard for large archives
def import_archive(self, task_id: int):
    """
    Import an archive file.

    Steps:
    1. Check if archive already processed (by hash)
    2. Extract to temp directory
    3. Create ImportTask for each media file inside
    4. Queue import tasks for extracted files
    5. Create archive tag if files were found
    6. Clean up is deferred until child tasks complete

    Args:
        task_id: ID of the ImportTask record

    Returns:
        dict with status and counts
    """
    from app.utils.image_importer import calculate_hash, compute_next_archive_tag_name
    from pyunpack import Archive

    task = ImportTask.query.get(task_id)
    if not task:
        return {'error': 'Task not found'}

    task.status = 'processing'
    task.started_at = datetime.utcnow()
    db.session.commit()

    tmpdir = None
    try:
        archive_path = task.source_path
        context = task.context or {}
        artist = context.get('artist', 'unknown')
        dest_dir = context.get('dest_dir', '/images')

        log.info(f"Processing archive: {archive_path}")

        # Compute archive hash
        archive_hash = calculate_hash(archive_path)
        task.file_hash = archive_hash

        # Check if already processed
        existing = ArchiveRecord.query.filter_by(hash=archive_hash).first()
        if existing:
            return _skip_task(task, 'Archive already imported')

        # Get archive size
        archive_size = os.path.getsize(archive_path)

        # Extract archive
        tmpdir = tempfile.mkdtemp(prefix='extract_')
        log.info(f"Extracting to: {tmpdir}")

        try:
            Archive(archive_path).extractall(tmpdir)
        except Exception as e:
            log.error(f"Extraction failed: {e}")
            if tmpdir and os.path.exists(tmpdir):
                shutil.rmtree(tmpdir, ignore_errors=True)
            return _fail_task(task, f'Extraction failed: {e}')

        # Find media files in extracted content
        media_files = []
        for root, _, files in os.walk(tmpdir):
            for filename in files:
                lower = filename.lower()
                if lower.endswith(('.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp',
                                   '.mp4', '.mov', '.webm', '.avi', '.mkv')):
                    media_files.append(os.path.join(root, filename))

        if not media_files:
            log.info(f"No media files found in archive: {archive_path}")
            if tmpdir and os.path.exists(tmpdir):
                shutil.rmtree(tmpdir, ignore_errors=True)
            return _skip_task(task, 'No media files in archive')

        log.info(f"Found {len(media_files)} media files in archive")

        # Create archive tag
        tag_name = compute_next_archive_tag_name(artist)
        archive_tag = Tag.query.filter_by(name=tag_name).first()
        if not archive_tag:
            archive_tag = Tag(name=tag_name, kind='archive')
            db.session.add(archive_tag)
            db.session.flush()

        # Create ArchiveRecord
        archive_record = ArchiveRecord(
            filename=archive_path,
            file_size=archive_size,
            hash=archive_hash,
            artist=artist,
            tag_id=archive_tag.id
        )
        db.session.add(archive_record)
        db.session.commit()

        # Queue tasks for each media file
        queued_count = 0
        for media_path in media_files:
            child_task = ImportTask(
                source_path=media_path,
                task_type='import_image',
                batch_id=task.batch_id,
                context={
                    'artist': artist,
                    'dest_dir': dest_dir,
                    'archive_task_id': task.id,
                    'archive_tag_id': archive_tag.id,
                    'temp_dir': tmpdir
                },
                status='pending'
            )
            db.session.add(child_task)
            db.session.flush()

            result = import_media_file.delay(child_task.id)
            child_task.celery_task_id = result.id
            child_task.status = 'queued'
            queued_count += 1

        # Store child task info for cleanup tracking
        task.context = task.context or {}
        task.context['temp_dir'] = tmpdir
        task.context['archive_tag_id'] = archive_tag.id
        task.context['child_count'] = queued_count

        task.status = 'complete'
        task.completed_at = datetime.utcnow()
        db.session.commit()

        log.info(f"Archive processed: queued {queued_count} files from {archive_path}")

        # Note: temp directory cleanup happens via a separate periodic task
        # or when all child tasks complete

        return {
            'status': 'extracted',
            'queued_files': queued_count,
            'archive_tag': tag_name
        }

    except Exception as e:
        db.session.rollback()
        log.error(f"Archive import failed: {e}", exc_info=True)
        if tmpdir and os.path.exists(tmpdir):
            shutil.rmtree(tmpdir, ignore_errors=True)
        return _fail_task(task, str(e))


def _fail_task(task: ImportTask, error: str, **kwargs) -> dict:
    """Mark task as failed with error message."""
    task.status = 'failed'
    task.error_message = error
    task.completed_at = datetime.utcnow()
    for key, value in kwargs.items():
        setattr(task, key, value)
    db.session.commit()

    # Update batch stats
    if task.batch_id:
        from app.tasks.scan import update_batch_stats
        update_batch_stats.delay(task.batch_id)

    log.error(f"Task {task.id} failed: {error}")
    return {'status': 'failed', 'error': error}


def _skip_task(task: ImportTask, reason: str, **kwargs) -> dict:
    """Mark task as skipped with reason."""
    task.status = 'skipped'
    task.error_message = reason
    task.completed_at = datetime.utcnow()
    for key, value in kwargs.items():
        setattr(task, key, value)
    db.session.commit()

    # Update batch stats
    if task.batch_id:
        from app.tasks.scan import update_batch_stats
        update_batch_stats.delay(task.batch_id)

    log.info(f"Task {task.id} skipped: {reason}")
    return {'status': 'skipped', 'reason': reason}


def _add_artist_tag(record: ImageRecord, artist: str):
    """Add artist tag to image record."""
    if not artist:
        return

    tag_name = f"artist:{artist}"
    tag = Tag.query.filter_by(name=tag_name).first()
    if not tag:
        tag = Tag(name=tag_name, kind='artist')
        db.session.add(tag)
        db.session.flush()

    if tag not in record.tags:
        record.tags.append(tag)


def _add_archive_tag(record: ImageRecord, archive_tag_id: int):
    """Add archive tag to image record."""
    if not archive_tag_id:
        return

    tag = Tag.query.get(archive_tag_id)
    if tag and tag not in record.tags:
        record.tags.append(tag)


def _supersede_existing(task: ImportTask, old_record: ImageRecord, src_path: str,
                        file_hash: str, phash, metadata: dict, artist: str,
                        dest_dir: str, filename: str) -> dict:
    """
    Replace an existing smaller image with a larger version.

    Preserves tags and metadata from the old record.
    """
    from app.utils.image_importer import (
        build_hashed_dest_path, generate_thumbnail, generate_video_thumbnail_mirrored,
        supersede_image
    )
    from app.utils.metadata_enrichment import find_sidecar_json, enrich_from_sidecar
    from app.tasks.sidecar import apply_sidecar_metadata

    log.info(f"Superseding {old_record.filename} ({old_record.width}x{old_record.height}) "
             f"with {filename} ({metadata['width']}x{metadata['height']})")

    # Enrich metadata from sidecar if available
    sidecar_path = find_sidecar_json(src_path)
    if sidecar_path:
        enriched = enrich_from_sidecar(sidecar_path, {'taken_at': metadata.get('taken_at')})
        if enriched.get('taken_at'):
            metadata['taken_at'] = enriched['taken_at']

    # Copy new file to destination
    target_dir = os.path.join(dest_dir, artist)
    os.makedirs(target_dir, exist_ok=True)
    dest_path = build_hashed_dest_path(target_dir, filename, file_hash)
    shutil.copy2(src_path, dest_path)

    # Generate new thumbnail
    try:
        if dest_path.lower().endswith(('.mp4', '.mov', '.webm', '.avi', '.mkv')):
            thumb_path = generate_video_thumbnail_mirrored(dest_path)
        else:
            thumb_path = generate_thumbnail(dest_path)
    except Exception as e:
        log.warning(f"Thumbnail generation failed: {e}")
        thumb_path = None

    # Supersede the old record
    record = supersede_image(
        old_record, dest_path, thumb_path, file_hash, phash, metadata, filename
    )

    db.session.commit()

    # Complete task
    task.status = 'complete'
    task.result_image_id = record.id
    task.completed_at = datetime.utcnow()
    db.session.commit()

    # Apply sidecar metadata if available
    if sidecar_path:
        apply_sidecar_metadata.delay(record.id, sidecar_path)

    # Update batch stats
    if task.batch_id:
        from app.tasks.scan import update_batch_stats
        update_batch_stats.delay(task.batch_id)

    return {
        'status': 'superseded',
        'image_id': record.id,
        'dest_path': dest_path
    }