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

"""
Celery tasks for thumbnail generation
"""
import os
import asyncio
from pathlib import Path
import logging
from typing import Tuple, Optional
import json

from celery import shared_task
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from PIL import Image
import imageio
from pillow_heif import register_heif_opener
import ffmpeg

# Try to import optional libraries
try:
    import pyvips
    PYVIPS_AVAILABLE = True
except ImportError:
    PYVIPS_AVAILABLE = False
    print("pyvips not available, using Pillow for image processing")

try:
    import rawpy
    RAWPY_AVAILABLE = True
except ImportError:
    RAWPY_AVAILABLE = False
    print("rawpy not available, using exiftool for RAW preview extraction")

from app.database import AsyncSessionLocal
from app.models import Photo
from app.config import settings

# Register HEIF opener with Pillow
register_heif_opener()

logger = logging.getLogger(__name__)

# Thumbnail sizes configuration
THUMB_SIZES = {
    'small': settings.thumbnails.small,
    'medium': settings.thumbnails.medium,
    'large': settings.thumbnails.large
}

def get_thumb_path(photo_id: str, size: str) -> str:
    """Get the path for a thumbnail file"""
    thumb_dir = f"/data/thumbs/{photo_id}"
    os.makedirs(thumb_dir, exist_ok=True)
    return f"{thumb_dir}/{size}.{settings.thumbnails.format}"

def process_standard_image(filepath: str) -> Image.Image:
    """Process standard image formats (JPEG, PNG, etc.)"""
    return Image.open(filepath)

def process_raw_image(filepath: str) -> Image.Image:
    """Process RAW image formats"""
    if RAWPY_AVAILABLE:
        try:
            with rawpy.imread(filepath) as raw:
                # Use half_size for faster processing
                rgb = raw.postprocess(use_camera_wb=True, half_size=True)
                # Convert numpy array to PIL Image
                return Image.fromarray(rgb, 'RGB')
        except Exception as e:
            logger.error(f"Error processing RAW file {filepath}: {e}")
            # Try to extract embedded JPEG preview
            return extract_raw_preview(filepath)
    else:
        # Use exiftool to extract embedded preview
        return extract_raw_preview(filepath)

def extract_raw_preview(filepath: str) -> Optional[Image.Image]:
    """Extract embedded JPEG preview from RAW file"""
    try:
        # Use exiftool to extract preview
        import subprocess
        import tempfile
        
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
            cmd = ['exiftool', '-b', '-PreviewImage', filepath]
            result = subprocess.run(cmd, capture_output=True)
            
            if result.returncode == 0 and result.stdout:
                tmp.write(result.stdout)
                tmp.flush()
                return Image.open(tmp.name)
    except Exception as e:
        logger.error(f"Error extracting RAW preview from {filepath}: {e}")
    
    return None

def process_heic_image(filepath: str) -> Image.Image:
    """Process HEIC/HEIF image formats"""
    try:
        # Use pillow-heif to open the image
        img = Image.open(filepath)
        # Convert to RGB if needed
        if img.mode != 'RGB':
            img = img.convert('RGB')
        return img
    except Exception as e:
        logger.error(f"Error processing HEIC file {filepath}: {e}")
        raise

def process_video_thumbnail(filepath: str) -> Image.Image:
    """Extract a still frame from a video file as a PIL Image."""
    import tempfile
    from io import BytesIO

    tmp_path: Optional[str] = None
    try:
        # Find a usable seek timestamp. Some camera MOVs only expose
        # duration at the format level, and stream 0 isn't always the
        # video stream — search explicitly and fall back to the format
        # duration, then to t=0 if neither is available.
        probe = ffmpeg.probe(filepath)
        duration: Optional[float] = None
        for stream_info in probe.get('streams', []):
            if stream_info.get('codec_type') != 'video':
                continue
            raw_duration = stream_info.get('duration')
            if raw_duration is not None:
                try:
                    duration = float(raw_duration)
                    break
                except (TypeError, ValueError):
                    pass
        if duration is None:
            raw_duration = probe.get('format', {}).get('duration')
            if raw_duration is not None:
                try:
                    duration = float(raw_duration)
                except (TypeError, ValueError):
                    duration = None

        # Seek to 10% in for a representative frame; clamp very short
        # clips to t=0 so we don't seek past the end.
        timestamp = max(0.0, (duration or 0.0) * 0.1)

        # NamedTemporaryFile creates the file on disk, so we MUST tell
        # ffmpeg to overwrite it (otherwise it prompts on stdin and the
        # call hangs/fails — which is why videos were getting the gray
        # placeholder). We close the handle immediately and clean up
        # in `finally` ourselves.
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
            tmp_path = tmp.name

        stream = ffmpeg.input(filepath, ss=timestamp)
        stream = ffmpeg.output(
            stream,
            tmp_path,
            vframes=1,
            format='image2',
            vcodec='mjpeg',
        )
        ffmpeg.run(
            stream,
            capture_stdout=True,
            capture_stderr=True,
            overwrite_output=True,
        )

        # Load the frame fully into memory so we can delete the temp
        # file immediately. Pillow's `Image.open` is lazy, which would
        # otherwise leave the file dangling.
        with open(tmp_path, 'rb') as fh:
            data = fh.read()
        if not data:
            raise RuntimeError("ffmpeg produced an empty frame")
        return Image.open(BytesIO(data)).copy()
    except ffmpeg.Error as e:
        stderr = (e.stderr or b'').decode('utf-8', errors='replace')
        logger.error(
            f"ffmpeg failed extracting video thumbnail from {filepath}: {stderr}"
        )
        return create_placeholder_thumbnail('video')
    except Exception as e:
        logger.error(f"Error extracting video thumbnail from {filepath}: {e}")
        return create_placeholder_thumbnail('video')
    finally:
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.unlink(tmp_path)
            except OSError:
                pass

def create_placeholder_thumbnail(media_type: str) -> Image.Image:
    """Create a placeholder thumbnail for failed processing"""
    # Create a simple gray placeholder
    img = Image.new('RGB', (640, 480), color=(128, 128, 128))
    return img

def auto_rotate_image(image: Image.Image) -> Image.Image:
    """Auto-rotate image based on EXIF orientation"""
    try:
        # Get EXIF data
        exif = image._getexif()
        if exif:
            orientation = exif.get(274)  # Orientation tag
            
            rotation_map = {
                3: 180,
                6: 270,  # Note: PIL uses different rotation values than vips
                8: 90
            }
            
            if orientation in rotation_map:
                image = image.rotate(rotation_map[orientation], expand=True)
    except:
        pass  # No orientation data available
    
    return image

def generate_thumbnail(image: Image.Image, size: int, output_path: str):
    """Generate a thumbnail of the specified size"""
    # Maintain aspect ratio
    image.thumbnail((size, size), Image.Resampling.LANCZOS)
    
    # Save as WebP with specified quality
    image.save(
        output_path,
        'WEBP',
        quality=settings.thumbnails.quality,
        method=4  # Balance between speed and compression
    )

@shared_task(bind=True, name='generate_thumbnails')
def generate_thumbnails(self, photo_id: str):
    """Generate thumbnails for a photo"""
    return asyncio.run(_generate_thumbnails_async(photo_id, self))

async def _generate_thumbnails_async(photo_id: str, task):
    """Async implementation of thumbnail generation"""
    async with AsyncSessionLocal() as session:
        try:
            # Get photo from database
            result = await session.execute(
                select(Photo).where(Photo.id == photo_id)
            )
            photo = result.scalar_one_or_none()
            
            if not photo:
                logger.error(f"Photo not found: {photo_id}")
                return {'status': 'error', 'message': 'Photo not found'}
            
            # Check if file exists
            if not os.path.exists(photo.filepath):
                logger.error(f"File not found: {photo.filepath}")
                photo.processing_status = 'failed'
                photo.processing_error = 'File not found'
                await session.commit()
                return {'status': 'error', 'message': 'File not found'}
            
            # Update processing status
            photo.processing_status = 'processing'
            await session.commit()
            
            # Load and process the image based on type
            image = None
            
            if photo.media_type == 'photo':
                image = process_standard_image(photo.filepath)
            elif photo.media_type == 'raw':
                image = process_raw_image(photo.filepath)
            elif photo.media_type == 'heic':
                image = process_heic_image(photo.filepath)
            elif photo.media_type == 'video':
                image = process_video_thumbnail(photo.filepath)
            else:
                logger.error(f"Unsupported media type: {photo.media_type}")
                image = create_placeholder_thumbnail(photo.media_type)
            
            if not image:
                raise Exception("Failed to process image")
            
            # Auto-rotate based on EXIF
            image = auto_rotate_image(image)

            # Store original dimensions
            photo.width = image.width
            photo.height = image.height

            # Perceptual hash from the original-resolution decoded frame.
            # pHash is robust to resize/recompression but the thumbnail
            # loop below mutates `image` in place, so this MUST run before
            # the loop sees it. Failures are non-fatal — phash is a
            # nice-to-have, not a blocker for thumbnail generation.
            try:
                import imagehash
                photo.phash = str(imagehash.phash(image))  # 16-char hex
            except Exception as e:
                logger.warning(f"phash failed for {photo_id}: {e}")
                photo.phash = None

            # Generate thumbnails for each size
            for size_name, size_value in THUMB_SIZES.items():
                thumb_path = get_thumb_path(photo_id, size_name)
                generate_thumbnail(image, size_value, thumb_path)
                
                # Update database with thumbnail path
                setattr(photo, f'thumb_{size_name}', thumb_path)
                
                # Update progress
                task.update_state(
                    state='PROGRESS',
                    meta={'current_size': size_name, 'photo_id': photo_id}
                )
            
            # Update processing status
            photo.processing_status = 'completed'
            photo.processing_error = None
            await session.commit()
            
            logger.info(f"Thumbnails generated for photo {photo_id}")
            return {'status': 'success', 'photo_id': photo_id}
            
        except Exception as e:
            logger.error(f"Error generating thumbnails for {photo_id}: {e}")
            
            # Update error status
            if photo:
                photo.processing_status = 'failed'
                photo.processing_error = str(e)
                await session.commit()
            
            return {'status': 'error', 'message': str(e)}

@shared_task(name='regenerate_all_thumbnails')
def regenerate_all_thumbnails():
    """Regenerate thumbnails for all photos"""
    return asyncio.run(_regenerate_all_thumbnails_async())

async def _regenerate_all_thumbnails_async():
    """Async implementation of regenerating all thumbnails"""
    async with AsyncSessionLocal() as session:
        # Get all photos that need thumbnails
        result = await session.execute(
            select(Photo).where(
                Photo.processing_status.in_(['pending', 'failed'])
            )
        )
        photos = result.scalars().all()

        logger.info(f"Regenerating thumbnails for {len(photos)} photos")

        for photo in photos:
            generate_thumbnails.delay(photo.id)

        return {'status': 'queued', 'count': len(photos)}


# ── Perceptual hash backfill ────────────────────────────────────────────
#
# When phash was added post-launch, every existing photo has phash=NULL.
# This task fills them in by reading the existing thumb_large (the cheap
# option — pHash is robust to scale, and the thumb is already on local
# disk so we avoid re-decoding the original RAW/HEIC). Falls back to the
# original filepath if the thumb isn't available for some reason. Runs
# in batches to keep memory bounded and to give the user incremental
# progress visible in the worker logs.

@shared_task(name='backfill_phashes')
def backfill_phashes():
    """Compute and persist phash for every photo currently missing one."""
    return asyncio.run(_backfill_phashes_async())


async def _backfill_phashes_async():
    import imagehash
    from PIL import Image as _PILImage

    BATCH = 100
    total_done = 0
    total_failed = 0

    async with AsyncSessionLocal() as session:
        while True:
            result = await session.execute(
                select(Photo)
                .where(Photo.phash.is_(None))
                .where(Photo.processing_status == 'completed')
                .limit(BATCH)
            )
            batch = result.scalars().all()
            if not batch:
                break

            for photo in batch:
                source = photo.thumb_large or photo.filepath
                try:
                    if not source or not os.path.exists(source):
                        photo.phash = None
                        total_failed += 1
                        continue
                    with _PILImage.open(source) as im:
                        photo.phash = str(imagehash.phash(im))
                    total_done += 1
                except Exception as e:
                    logger.warning(f"phash backfill failed for {photo.id}: {e}")
                    total_failed += 1

            await session.commit()
            logger.info(
                f"Backfilled phashes: {total_done} done, {total_failed} failed"
            )

    return {
        'status': 'success',
        'computed': total_done,
        'failed': total_failed,
    }


@shared_task(name='regroup_duplicates')
def regroup_duplicates_task():
    """Celery wrapper around app.services.duplicates.regroup_duplicates.

    Importing the service inside the task body avoids a circular import
    at worker boot (the service uses AsyncSessionLocal which is also
    imported here at module top)."""
    from app.services.duplicates import regroup_duplicates
    return asyncio.run(regroup_duplicates())