# 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.

Two scan modes:
- Quick scan (default): Pre-loads existing task/hash data for O(1) lookups.
  Only queues truly new files. Skips pHash comparison for speed.
- Deep scan (on-demand): Full reprocessing of all files. Re-extracts metadata,
  recalculates hashes, regenerates thumbnails. For data integrity verification.
"""
import os
import logging
from datetime import datetime
from typing import Optional

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

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'
)


def _build_task_lookup() -> dict:
    """
    Pre-load all ImportTask records into a lookup dict for O(1) checks.
    Key: (source_path, task_type)
    Value: dict with 'status', 'file_size' for quick comparison

    This replaces per-file database queries with a single upfront query.
    """
    lookup = {}
    for task in ImportTask.query.with_entities(
        ImportTask.source_path,
        ImportTask.task_type,
        ImportTask.status,
        ImportTask.file_size
    ).all():
        key = (task.source_path, task.task_type)
        # Keep the most relevant status for each path
        # Priority: active > complete > skipped > failed
        existing = lookup.get(key)
        if existing is None:
            lookup[key] = {'status': task.status, 'file_size': task.file_size}
        elif task.status in ('pending', 'queued', 'processing'):
            # Active task takes priority
            lookup[key] = {'status': task.status, 'file_size': task.file_size}
        elif existing['status'] not in ('pending', 'queued', 'processing'):
            # Replace if existing is not active and new is complete
            if task.status == 'complete':
                lookup[key] = {'status': task.status, 'file_size': task.file_size}
    return lookup


def _should_skip_file(source_path: str, file_size: int, task_type: str,
                      task_lookup: dict) -> Optional[str]:
    """
    Check if file should be skipped using pre-loaded lookup.

    Returns:
        Skip reason string if should skip, None if should process
    """
    key = (source_path, task_type)
    existing = task_lookup.get(key)

    if existing is None:
        return None  # New file, process it

    status = existing['status']

    # Skip if task is currently active
    if status in ('pending', 'queued', 'processing'):
        return 'already_active'

    # Skip if completed/skipped with same size
    if status in ('complete', 'skipped') and existing['file_size'] == file_size:
        return 'already_processed'

    # File size changed - re-process
    return None


@celery.task(bind=True, name='app.tasks.scan.scan_directory')
def scan_directory(self, source_dir: str = '/import', dest_dir: str = '/images'):
    """
    Quick scan: Efficiently detect and queue new/changed files.

    Optimizations over naive approach:
    - Pre-loads all existing task records in one query (O(1) lookup per file)
    - Batches task creation commits
    - Skips pHash comparison (relies on content hash for duplicates)

    Strategy:
    1. Load existing ImportTask lookup dict (single query)
    2. Walk directory tree under each artist folder
    3. For each file, check lookup dict (no DB queries per file)
    4. Queue new/changed files for processing
    5. Handle archives specially (first volume only)

    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 quick scan: {source_dir}")

    # Pre-load existing tasks for O(1) lookup (single query)
    log.debug("Loading existing task lookup...")
    task_lookup = _build_task_lookup()
    log.debug(f"Loaded {len(task_lookup)} existing task records")

    # 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,
    }

    # Collect tasks to create in batches
    pending_tasks = []
    BATCH_SIZE = 50  # Commit every 50 tasks

    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

                    # Get file size once (used for lookup)
                    try:
                        file_size = os.path.getsize(full_path)
                    except OSError:
                        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

                        # Check lookup dict (no DB query)
                        skip_reason = _should_skip_file(
                            full_path, file_size, 'import_archive', task_lookup
                        )
                        if skip_reason:
                            stats['skipped_existing'] += 1
                            continue

                        # Queue archive task
                        task = ImportTask(
                            source_path=full_path,
                            task_type='import_archive',
                            file_size=file_size,
                            batch_id=batch.id,
                            context={'artist': artist_dir, 'dest_dir': dest_dir},
                            status='pending'
                        )
                        pending_tasks.append(('archive', task))
                        stats['queued_archives'] += 1
                        continue

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

                    # Check lookup dict (no DB query)
                    skip_reason = _should_skip_file(
                        full_path, file_size, 'import_image', task_lookup
                    )
                    if skip_reason:
                        stats['skipped_existing'] += 1
                        continue

                    # Queue media file task
                    task = ImportTask(
                        source_path=full_path,
                        task_type='import_image',
                        file_size=file_size,
                        batch_id=batch.id,
                        context={'artist': artist_dir, 'dest_dir': dest_dir},
                        status='pending'
                    )
                    pending_tasks.append(('media', task))
                    stats['queued_media'] += 1

                    # Batch commit
                    if len(pending_tasks) >= BATCH_SIZE:
                        _flush_pending_tasks(pending_tasks)
                        pending_tasks = []

        # Flush remaining tasks
        if pending_tasks:
            _flush_pending_tasks(pending_tasks)

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

        log.info(f"Quick 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 _flush_pending_tasks(pending_tasks: list):
    """
    Batch add tasks to session, commit, and queue Celery jobs.
    """
    from app.tasks.import_file import import_media_file, import_archive

    # Add all tasks to session
    for task_type, task in pending_tasks:
        db.session.add(task)

    # Flush to get IDs
    db.session.flush()

    # Queue Celery jobs
    for task_type, task in pending_tasks:
        if task_type == 'archive':
            result = import_archive.delay(task.id)
        else:
            result = import_media_file.delay(task.id)
        task.celery_task_id = result.id
        task.status = 'queued'

    # Commit batch
    db.session.commit()


@celery.task(bind=True, name='app.tasks.scan.deep_scan_directory')
def deep_scan_directory(self, source_dir: str = '/import', dest_dir: str = '/images',
                        regenerate_thumbnails: bool = True,
                        reapply_sidecars: bool = True,
                        verify_hashes: bool = True):
    """
    Deep scan: Full reprocessing of all files for data integrity verification.

    Unlike quick scan, this:
    - Ignores existing task status (forces reprocessing)
    - Re-extracts all metadata
    - Recalculates content hashes (optional verification)
    - Regenerates thumbnails (optional)
    - Re-applies sidecar metadata (optional)
    - Runs full pHash comparison for similarity detection

    Use cases:
    - Data integrity verification after storage issues
    - Picking up metadata from new/updated sidecar files
    - Regenerating thumbnails after algorithm changes
    - Recovery from partial imports

    Args:
        source_dir: Root directory to scan (default: /import)
        dest_dir: Destination for imported files (default: /images)
        regenerate_thumbnails: Regenerate all thumbnails (default: True)
        reapply_sidecars: Re-apply metadata from sidecar files (default: True)
        verify_hashes: Verify content hashes match stored records (default: True)

    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 DEEP scan: {source_dir}")
    log.info(f"Options: thumbnails={regenerate_thumbnails}, sidecars={reapply_sidecars}, verify={verify_hashes}")

    # Create batch record with deep scan flag
    batch = ImportBatch(source_directory=source_dir)
    db.session.add(batch)
    db.session.commit()

    stats = {
        'total_files': 0,
        'queued_media': 0,
        'queued_archives': 0,
        'skipped_active': 0,  # Only skip currently processing files
        'skipped_non_first_volume': 0,
        'skipped_non_media': 0,
    }

    # For deep scan, only load active tasks to avoid double-processing
    active_tasks = set()
    for task in ImportTask.query.filter(
        ImportTask.status.in_(['pending', 'queued', 'processing'])
    ).with_entities(ImportTask.source_path, ImportTask.task_type).all():
        active_tasks.add((task.source_path, task.task_type))

    log.debug(f"Found {len(active_tasks)} currently active tasks to skip")

    pending_tasks = []
    BATCH_SIZE = 50

    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"Deep 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

                    # Get file size
                    try:
                        file_size = os.path.getsize(full_path)
                    except OSError:
                        continue

                    # Check if archive
                    if lower.endswith(ALLOWED_ARCHIVE_EXTS):
                        if not is_first_volume(full_path):
                            stats['skipped_non_first_volume'] += 1
                            continue

                        # Only skip if currently active (deep scan reprocesses everything else)
                        if (full_path, 'import_archive') in active_tasks:
                            stats['skipped_active'] += 1
                            continue

                        task = ImportTask(
                            source_path=full_path,
                            task_type='import_archive',
                            file_size=file_size,
                            batch_id=batch.id,
                            context={
                                'artist': artist_dir,
                                'dest_dir': dest_dir,
                                'deep_scan': True,
                                'regenerate_thumbnails': regenerate_thumbnails,
                                'reapply_sidecars': reapply_sidecars,
                                'verify_hashes': verify_hashes
                            },
                            status='pending'
                        )
                        pending_tasks.append(('archive', task))
                        stats['queued_archives'] += 1
                        continue

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

                    # Only skip if currently active
                    if (full_path, 'import_image') in active_tasks:
                        stats['skipped_active'] += 1
                        continue

                    task = ImportTask(
                        source_path=full_path,
                        task_type='import_image',
                        file_size=file_size,
                        batch_id=batch.id,
                        context={
                            'artist': artist_dir,
                            'dest_dir': dest_dir,
                            'deep_scan': True,
                            'regenerate_thumbnails': regenerate_thumbnails,
                            'reapply_sidecars': reapply_sidecars,
                            'verify_hashes': verify_hashes
                        },
                        status='pending'
                    )
                    pending_tasks.append(('media', task))
                    stats['queued_media'] += 1

                    # Batch commit
                    if len(pending_tasks) >= BATCH_SIZE:
                        _flush_pending_tasks(pending_tasks)
                        pending_tasks = []

        # Flush remaining tasks
        if pending_tasks:
            _flush_pending_tasks(pending_tasks)

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

        log.info(f"Deep scan complete. Queued {stats['queued_media']} media, "
                 f"{stats['queued_archives']} archives. "
                 f"Skipped {stats['skipped_active']} active.")

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

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


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


@celery.task(name='app.tasks.scan.update_system_stats')
def update_system_stats():
    """
    Compute and cache system-wide statistics for the dashboard.

    Called periodically by Celery Beat to avoid expensive queries on page load.
    Stores results in AppSettings as JSON for quick retrieval.

    Stats computed:
    - Total images and storage usage
    - Tag counts by kind
    - Pending files/archives in import folder

    Returns:
        dict with computed stats
    """
    import json
    from sqlalchemy import func
    from app.models import ImageRecord, Tag, AppSettings

    log.info("Updating system stats cache...")

    try:
        # Database stats
        total_images = ImageRecord.query.count()
        total_tags = Tag.query.count()

        # Tag kind breakdown
        tag_kinds = dict(
            db.session.query(Tag.kind, func.count(Tag.id))
            .group_by(Tag.kind).all()
        )

        # Storage stats (sum of file_size column)
        storage_result = db.session.query(func.sum(ImageRecord.file_size)).scalar()
        total_storage_bytes = storage_result or 0

        # Import folder stats
        import_dir = '/import'
        pending_files = 0
        pending_archives = 0
        if os.path.exists(import_dir):
            for artist_dir in os.listdir(import_dir):
                artist_path = os.path.join(import_dir, artist_dir)
                if not os.path.isdir(artist_path):
                    continue
                for root, _, files in os.walk(artist_path):
                    for f in files:
                        lower = f.lower()
                        if lower.endswith(ALLOWED_MEDIA_EXTS):
                            pending_files += 1
                        elif lower.endswith(ALLOWED_ARCHIVE_EXTS):
                            pending_archives += 1

        # Format storage for display
        storage_human = _format_bytes(total_storage_bytes)

        # Build stats dict
        stats = {
            'images': {
                'total': total_images,
                'storage_bytes': total_storage_bytes,
                'storage_human': storage_human
            },
            'tags': {
                'total': total_tags,
                'by_kind': tag_kinds
            },
            'import_folder': {
                'pending_files': pending_files,
                'pending_archives': pending_archives
            },
            'updated_at': datetime.utcnow().isoformat()
        }

        # Store in AppSettings
        setting = AppSettings.query.filter_by(key='system_stats_cache').first()
        if not setting:
            setting = AppSettings(key='system_stats_cache')
            db.session.add(setting)
        setting.value = json.dumps(stats)
        db.session.commit()

        log.info(f"System stats updated: {total_images} images, {total_tags} tags, "
                 f"{storage_human} storage, {pending_files + pending_archives} pending")

        return stats

    except Exception as e:
        log.error(f"Failed to update system stats: {e}", exc_info=True)
        raise


def _format_bytes(size_bytes: int) -> str:
    """Format bytes to human readable string."""
    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"