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
from sqlalchemy.ext.asyncio import AsyncSession
import aiofiles

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

# 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"""
    logger.info(f"Starting scan of 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
            
            # Walk the directory tree
            total_files = 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)
                
                # Filter supported files
                supported_files = [f for f in files if Path(f).suffix.lower() in SUPPORTED_EXTENSIONS]
                total_files += len(supported_files)
                
                # 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]
                    
                    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
                                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
                            duplicate = await session.execute(
                                select(Photo).where(Photo.file_hash == file_hash)
                            ) if file_hash else None
                            
                            # 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=bool(duplicate.scalar_one_or_none() if duplicate else False),
                                processing_status='pending'
                            )
                            
                            session.add(photo)
                            await session.flush()  # Get the photo ID
                            
                            # Queue thumbnail generation
                            generate_thumbnails.delay(photo.id)
                            
                            # Queue metadata extraction
                            extract_metadata.delay(photo.id)
                            
                            processed_files += 1
                            
                            # Update progress
                            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)})
                            continue
                    
                    # Commit batch
                    await session.commit()
            
            # 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}")
            await session.rollback()
            raise

async def get_or_create_source_root(session: AsyncSession, path: str) -> SourceRoot:
    """Get or create a source root entry"""
    from sqlalchemy import select
    
    result = await session.execute(
        select(SourceRoot).where(SourceRoot.path == path)
    )
    source_root = result.scalar_one_or_none()
    
    if not source_root:
        source_root = SourceRoot(
            name=Path(path).name,
            path=path
        )
        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"""
    from sqlalchemy import select
    
    result = await session.execute(
        select(Folder).where(Folder.path == path)
    )
    folder = result.scalar_one_or_none()
    
    if not folder:
        parent_path = str(Path(path).parent)
        parent = None
        
        if parent_path != path:  # Not root folder
            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(path).name,
            path=path,
            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 all configured source roots"""
    for source_root in settings.source_roots:
        if os.path.exists(source_root.path):
            scan_folder.delay(source_root.path)
        else:
            logger.warning(f"Source root path does not exist: {source_root.path}")

@shared_task(name='watch_folders')
def watch_folders():
    """
    Watch folders for changes using watchfiles
    This is a long-running task that monitors file system events
    """
    from watchfiles import watch
    
    paths = [sr.path for sr in settings.source_roots if os.path.exists(sr.path)]
    
    if not paths:
        logger.warning("No valid source roots to watch")
        return
    
    logger.info(f"Starting folder watcher for: {paths}")
    
    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
                parent_dir = str(Path(filepath).parent)
                scan_folder.delay(parent_dir)
                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}")