"""
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
import pyvips
import rawpy
import imageio
from PIL import Image
from pillow_heif import register_heif_opener
import ffmpeg

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) -> pyvips.Image:
    """Process standard image formats (JPEG, PNG, etc.)"""
    return pyvips.Image.new_from_file(filepath, access='sequential')

def process_raw_image(filepath: str) -> pyvips.Image:
    """Process RAW image formats"""
    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 pyvips image
            return pyvips.Image.new_from_array(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)

def extract_raw_preview(filepath: str) -> Optional[pyvips.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 pyvips.Image.new_from_file(tmp.name, access='sequential')
    except Exception as e:
        logger.error(f"Error extracting RAW preview from {filepath}: {e}")
    
    return None

def process_heic_image(filepath: str) -> pyvips.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')
        # Save to temp file and load with pyvips
        import tempfile
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
            img.save(tmp.name, 'JPEG')
            return pyvips.Image.new_from_file(tmp.name, access='sequential')
    except Exception as e:
        logger.error(f"Error processing HEIC file {filepath}: {e}")
        raise

def process_video_thumbnail(filepath: str) -> pyvips.Image:
    """Extract thumbnail from video file"""
    try:
        # Get video duration
        probe = ffmpeg.probe(filepath)
        duration = float(probe['streams'][0]['duration'])
        
        # Extract frame at 10% of duration
        timestamp = duration * 0.1
        
        # Extract frame using ffmpeg
        import tempfile
        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
            stream = ffmpeg.input(filepath, ss=timestamp)
            stream = ffmpeg.output(stream, tmp.name, vframes=1, format='image2', vcodec='mjpeg')
            ffmpeg.run(stream, capture_stdout=True, capture_stderr=True)
            
            return pyvips.Image.new_from_file(tmp.name, access='sequential')
    except Exception as e:
        logger.error(f"Error extracting video thumbnail from {filepath}: {e}")
        # Create a placeholder thumbnail
        return create_placeholder_thumbnail('video')

def create_placeholder_thumbnail(media_type: str) -> pyvips.Image:
    """Create a placeholder thumbnail for failed processing"""
    # Create a simple gray placeholder
    placeholder = pyvips.Image.black(640, 480)
    placeholder = placeholder + [128, 128, 128]  # Make it gray
    return placeholder

def auto_rotate_image(image: pyvips.Image) -> pyvips.Image:
    """Auto-rotate image based on EXIF orientation"""
    try:
        orientation = image.get('orientation')
        
        rotation_map = {
            3: 180,
            6: 90,
            8: 270
        }
        
        if orientation in rotation_map:
            image = image.rot(rotation_map[orientation])
    except:
        pass  # No orientation data available
    
    return image

def generate_thumbnail(image: pyvips.Image, size: int, output_path: str):
    """Generate a thumbnail of the specified size"""
    # Calculate scale to fit within size (longest edge)
    width = image.width
    height = image.height
    
    if width > height:
        scale = size / width
    else:
        scale = size / height
    
    # Only downscale, never upscale
    if scale < 1:
        image = image.resize(scale)
    
    # Save as WebP with specified quality
    image.webpsave(
        output_path,
        Q=settings.thumbnails.quality,
        effort=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
            
            # 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)}