mule-image/backend/app/tasks/scan.py

"""
Celery tasks for scanning folders and indexing photos
"""
import os
import hashlib
import asyncio
from pathlib import Path
from datetime import datetime
import logging
import json
from typing import List, Dict, Optional

from celery import shared_task
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
import aiofiles
import redis

from app.database import AsyncSessionLocal
from app.models import Photo, Folder, SourceRoot
from app.config import settings
from app.tasks.thumbs import generate_thumbnails
from app.services.metadata import extract_metadata

logger = logging.getLogger(__name__)


# Redis keys read by GET /api/v1/library/scan/status. The frontend
# ScanProgress widget polls that endpoint, so anything we want to surface
# in the UI lives here.
REDIS_KEY_ACTIVE = 'scan:active'
REDIS_KEY_CURRENT_FOLDER = 'scan:current_folder'
REDIS_KEY_PROCESSED = 'scan:processed_files'
REDIS_KEY_TOTAL = 'scan:total_files'
REDIS_KEY_ERRORS = 'scan:errors'
MAX_ERROR_ENTRIES = 50  # cap the errors list so a noisy scan doesn't blow Redis


def _get_redis():
    """Connect to the broker for progress writes. Returns None on failure
    so a Redis outage doesn't prevent the scan itself from running."""
    try:
        return redis.Redis.from_url(settings.celery_broker_url)
    except Exception as e:
        logger.warning(f"Could not reach Redis for scan progress: {e}")
        return None

# Supported file extensions
PHOTO_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.tiff', '.tif', '.webp', '.bmp'}
RAW_EXTENSIONS = {'.cr2', '.cr3', '.nef', '.arw', '.raf', '.dng', '.orf', '.rw2', '.pef', '.srw'}
HEIC_EXTENSIONS = {'.heic', '.heif'}
VIDEO_EXTENSIONS = {'.mp4', '.mov', '.avi', '.mkv', '.mts', '.m2ts', '.3gp', '.wmv', '.flv'}

SUPPORTED_EXTENSIONS = PHOTO_EXTENSIONS | RAW_EXTENSIONS | HEIC_EXTENSIONS | VIDEO_EXTENSIONS

def get_media_type(filepath: str) -> str:
    """Determine media type from file extension"""
    ext = Path(filepath).suffix.lower()
    if ext in PHOTO_EXTENSIONS:
        return 'photo'
    elif ext in RAW_EXTENSIONS:
        return 'raw'
    elif ext in HEIC_EXTENSIONS:
        return 'heic'
    elif ext in VIDEO_EXTENSIONS:
        return 'video'
    return 'unknown'

async def calculate_file_hash(filepath: str) -> str:
    """Calculate SHA-256 hash of a file"""
    hash_sha256 = hashlib.sha256()
    try:
        async with aiofiles.open(filepath, 'rb') as f:
            while chunk := await f.read(8192):
                hash_sha256.update(chunk)
        return hash_sha256.hexdigest()
    except Exception as e:
        logger.error(f"Error calculating hash for {filepath}: {e}")
        return ""

@shared_task(bind=True, name='scan_folder')
def scan_folder(self, folder_path: str, source_root_id: Optional[str] = None):
    """
    Scan a folder and index all photos/videos
    """
    # Run async function in sync context
    return asyncio.run(_scan_folder_async(folder_path, source_root_id, self))

async def _scan_folder_async(folder_path: str, source_root_id: Optional[str], task):
    """Async implementation of folder scanning. Writes progress to Redis so
    GET /api/v1/library/scan/status can surface it to the frontend
    ScanProgress widget."""
    logger.info(f"Starting scan of folder: {folder_path}")

    r = _get_redis()

    def progress_set(key: str, value) -> None:
        if r is None:
            return
        try:
            r.set(key, str(value))
        except Exception as e:
            logger.debug(f"scan progress set failed: {e}")

    def progress_push_error(message: str) -> None:
        if r is None:
            return
        try:
            r.lpush(REDIS_KEY_ERRORS, message)
            r.ltrim(REDIS_KEY_ERRORS, 0, MAX_ERROR_ENTRIES - 1)
        except Exception as e:
            logger.debug(f"scan progress push_error failed: {e}")

    # Mark scan active immediately so the UI starts polling fast.
    progress_set(REDIS_KEY_ACTIVE, 'true')
    progress_set(REDIS_KEY_CURRENT_FOLDER, folder_path)

    async with AsyncSessionLocal() as session:
        try:
            # Get or create source root
            if not source_root_id:
                source_root = await get_or_create_source_root(session, folder_path)
                source_root_id = source_root.id

            # Pre-walk to compute the total file count upfront. Without this
            # the progress bar would jump every time a new subfolder is
            # encountered because the running total kept growing.
            total_files = 0
            for _root, _dirs, files in os.walk(folder_path):
                total_files += sum(
                    1 for f in files if Path(f).suffix.lower() in SUPPORTED_EXTENSIONS
                )
            progress_set(REDIS_KEY_TOTAL, total_files)
            progress_set(REDIS_KEY_PROCESSED, 0)

            processed_files = 0
            errors = []

            for root, dirs, files in os.walk(folder_path):
                # Get or create folder entry
                folder = await get_or_create_folder(session, root, source_root_id)
                progress_set(REDIS_KEY_CURRENT_FOLDER, root)

                # Filter supported files
                supported_files = [f for f in files if Path(f).suffix.lower() in SUPPORTED_EXTENSIONS]

                # Process files in batches
                batch_size = settings.scanner.batch_size
                for i in range(0, len(supported_files), batch_size):
                    batch = supported_files[i:i + batch_size]
                    # Defer task dispatch until AFTER commit so workers don't
                    # query for rows that aren't visible to other sessions yet.
                    pending_dispatch: list[str] = []

                    for filename in batch:
                        filepath = os.path.join(root, filename)

                        try:
                            # Check if file already exists in database
                            existing = await session.execute(
                                select(Photo).where(Photo.filepath == filepath)
                            )
                            if existing.scalar_one_or_none():
                                logger.debug(f"File already indexed: {filepath}")
                                processed_files += 1
                                progress_set(REDIS_KEY_PROCESSED, processed_files)
                                continue

                            # Get file stats
                            stat = os.stat(filepath)

                            # Calculate file hash for duplicate detection
                            file_hash = await calculate_file_hash(filepath)

                            # Check for duplicate by hash. We only care
                            # whether *any* other photo shares this hash, so
                            # use a count rather than scalar_one_or_none()
                            # which raises "Multiple rows were found" the
                            # moment the library has 2+ copies of the same
                            # file (i.e. exactly the case we're trying to
                            # flag).
                            is_dup = False
                            if file_hash:
                                dup_count = (await session.execute(
                                    select(func.count(Photo.id)).where(
                                        Photo.file_hash == file_hash
                                    )
                                )).scalar() or 0
                                is_dup = dup_count > 0

                            # Create photo entry
                            photo = Photo(
                                filepath=filepath,
                                filename=filename,
                                folder_id=folder.id,
                                file_hash=file_hash,
                                media_type=get_media_type(filepath),
                                original_format=Path(filepath).suffix.upper()[1:],
                                file_size=stat.st_size,
                                taken_at=datetime.fromtimestamp(stat.st_mtime),
                                taken_at_source='filesystem',
                                is_duplicate=is_dup,
                                processing_status='pending'
                            )

                            session.add(photo)
                            await session.flush()  # Assign defaults / FK ids

                            # Queue dispatch happens after the batch commit
                            # below; otherwise the worker can race the writer
                            # and see "Photo not found".
                            pending_dispatch.append(photo.id)

                            processed_files += 1
                            progress_set(REDIS_KEY_PROCESSED, processed_files)

                            # Celery internal progress (used by celery tooling)
                            if processed_files % 10 == 0:
                                task.update_state(
                                    state='PROGRESS',
                                    meta={
                                        'current': processed_files,
                                        'total': total_files,
                                        'folder': root,
                                    }
                                )

                        except Exception as e:
                            logger.error(f"Error processing file {filepath}: {e}")
                            errors.append({'file': filepath, 'error': str(e)})
                            progress_push_error(f"{filepath}: {e}")
                            continue

                    # Commit batch, then queue worker tasks. Dispatch order
                    # matters: commit first so workers can find the rows.
                    await session.commit()

                    for photo_id in pending_dispatch:
                        generate_thumbnails.delay(photo_id)
                        extract_metadata.delay(photo_id)

            # Update folder scan timestamp
            folder.last_scanned = datetime.utcnow()
            folder.photo_count = processed_files
            await session.commit()

            logger.info(f"Scan complete. Processed {processed_files}/{total_files} files. Errors: {len(errors)}")

            return {
                'status': 'completed',
                'processed': processed_files,
                'total': total_files,
                'errors': errors,
            }

        except Exception as e:
            logger.error(f"Scan failed: {e}")
            progress_push_error(f"scan failed: {e}")
            await session.rollback()
            raise
        finally:
            # Always mark inactive on the way out so a crashed scan doesn't
            # leave the UI thinking we're still scanning.
            progress_set(REDIS_KEY_ACTIVE, 'false')

def _normalize_path(path: str) -> str:
    """Canonicalise a filesystem path so we don't get duplicate DB rows for
    the same physical directory due to trailing slashes, redundant separators,
    or `.` segments. Symlinks are NOT resolved (we want to keep mount paths
    intact for cross-machine portability)."""
    return os.path.normpath(path)


async def get_or_create_source_root(session: AsyncSession, path: str) -> SourceRoot:
    """Get or create a source root entry, matching by normalized path."""
    from sqlalchemy import select

    norm = _normalize_path(path)
    result = await session.execute(
        select(SourceRoot).where(SourceRoot.path == norm)
    )
    source_root = result.scalar_one_or_none()

    if not source_root:
        source_root = SourceRoot(
            name=Path(norm).name,
            path=norm,
        )
        session.add(source_root)
        await session.flush()

    return source_root


async def get_or_create_folder(session: AsyncSession, path: str, source_root_id: str) -> Folder:
    """Get or create a folder entry, matching by normalized path."""
    from sqlalchemy import select

    norm = _normalize_path(path)
    result = await session.execute(
        select(Folder).where(Folder.path == norm)
    )
    folder = result.scalar_one_or_none()

    if not folder:
        parent_path = _normalize_path(str(Path(norm).parent))

        if parent_path != norm:  # Not the filesystem root
            parent_result = await session.execute(
                select(Folder).where(Folder.path == parent_path)
            )
            parent = parent_result.scalar_one_or_none()
            if parent:
                parent_id = parent.id
            else:
                # Recursively create parent
                parent = await get_or_create_folder(session, parent_path, source_root_id)
                parent_id = parent.id
        else:
            parent_id = None

        folder = Folder(
            name=Path(norm).name,
            path=norm,
            parent_id=parent_id,
            source_root_id=source_root_id,
        )
        session.add(folder)
        await session.flush()

    return folder

@shared_task(name='scan_all_source_roots')
def scan_all_source_roots():
    """Scan every active source root currently registered in the DB."""
    # Clear stale per-scan progress before queuing new work so the UI sees
    # a clean slate even if a previous run crashed mid-flight.
    r = _get_redis()
    if r is not None:
        try:
            r.delete(REDIS_KEY_ERRORS)
            r.set(REDIS_KEY_PROCESSED, 0)
            r.set(REDIS_KEY_TOTAL, 0)
        except Exception as e:
            logger.debug(f"scan_all_source_roots redis reset failed: {e}")

    return asyncio.run(_scan_all_source_roots_async())


async def _scan_all_source_roots_async():
    """Read every active SourceRoot from the DB and queue a scan_folder task
    for each. Source roots whose path no longer exists on disk are skipped
    with a warning (the cleanup service surfaces those at startup too)."""
    async with AsyncSessionLocal() as session:
        result = await session.execute(
            select(SourceRoot).where(SourceRoot.is_active == True)  # noqa: E712
        )
        source_roots = result.scalars().all()
        for sr in source_roots:
            if os.path.exists(sr.path):
                scan_folder.delay(sr.path, sr.id)
            else:
                logger.warning(f"Source root path does not exist: {sr.path}")


@shared_task(name='watch_folders')
def watch_folders():
    """
    Watch folders for changes using watchfiles. Long-running task that
    monitors filesystem events under every active source root.
    """
    from watchfiles import watch

    # Read source roots from the DB instead of the (now-removed) YAML
    # config. We need both the path and the id so we can dispatch
    # scan_folder with the source_root_id when an event fires.
    roots: list[tuple[str, str]] = []
    try:
        async def _load_roots():
            async with AsyncSessionLocal() as session:
                result = await session.execute(
                    select(SourceRoot).where(SourceRoot.is_active == True)  # noqa: E712
                )
                return [
                    (os.path.normpath(sr.path), sr.id)
                    for sr in result.scalars().all()
                    if os.path.exists(sr.path)
                ]
        roots = asyncio.run(_load_roots())
    except Exception as e:
        logger.error(f"watch_folders could not load source roots: {e}")
        return

    if not roots:
        logger.warning("No valid source roots to watch")
        return

    paths = [p for p, _ in roots]
    logger.info(f"Starting folder watcher for: {paths}")

    def find_source_root_for(path: str) -> Optional[str]:
        """Return the source_root id whose path contains `path`, or None."""
        normalized = os.path.normpath(path)
        for root_path, root_id in roots:
            if normalized == root_path or normalized.startswith(root_path + os.sep):
                return root_id
        return None

    for changes in watch(*paths):
        for change_type, filepath in changes:
            filepath = str(filepath)

            # Check if it's a supported file type
            if Path(filepath).suffix.lower() not in SUPPORTED_EXTENSIONS:
                continue

            if change_type == 'added' or change_type == 'modified':
                # Queue scan for the parent folder, with the source_root_id
                # resolved by ancestor lookup so scan_folder doesn't
                # auto-create a new SourceRoot for an arbitrary subdir.
                parent_dir = str(Path(filepath).parent)
                source_root_id = find_source_root_for(parent_dir)
                if source_root_id is None:
                    logger.debug(
                        f"watcher event for {filepath}: parent {parent_dir} "
                        f"not under any active source root, ignoring"
                    )
                    continue
                scan_folder.delay(parent_dir, source_root_id)
                logger.info(f"File {change_type}: {filepath}, queued scan for {parent_dir}")
            elif change_type == 'deleted':
                # Handle file deletion
                asyncio.run(handle_file_deletion(filepath))

async def handle_file_deletion(filepath: str):
    """Handle deletion of a file from the filesystem"""
    from sqlalchemy import select
    
    async with AsyncSessionLocal() as session:
        result = await session.execute(
            select(Photo).where(Photo.filepath == filepath)
        )
        photo = result.scalar_one_or_none()
        
        if photo:
            # Mark as missing or delete from database
            photo.is_discarded = True
            photo.discarded_at = datetime.utcnow()
            await session.commit()
            logger.info(f"Marked photo as discarded: {filepath}")