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

# Sizes the worker writes to /data/thumbs. Empty set since Phase 4 —
# the API serves all sizes via Nextcloud's /core/preview proxy.
# generate_thumbnails still runs the decode-and-pHash side-effect
# (perceptual dedup is mule-only and needs the original-resolution
# pixels) but no longer touches the disk.
WORKER_THUMB_SIZES: set[str] = set()

def get_thumb_path(photo_id: str, size: str, user_id: str = None) -> str:
    """Get the path for a thumbnail file.

    When user_id is provided, thumbnails are stored under a user-specific
    subdirectory to enforce isolation between users.
    """
    if user_id:
        thumb_dir = f"/data/thumbs/{user_id}/{photo_id}"
    else:
        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.warning(f"rawpy failed for {filepath}: {e}; trying embedded preview")
            preview = extract_raw_preview(filepath)
            if preview is not None:
                return preview
            # iPhone "Apple ProRAW" / Linear DNG has no embedded preview and
            # LibRaw rejects it as not-a-RAW. It IS a TIFF container with a
            # developed RGB image inside, so PIL opens it directly.
            try:
                logger.warning(f"embedded preview missing for {filepath}; trying PIL TIFF fallback")
                return Image.open(filepath)
            except Exception as e2:
                logger.error(f"PIL fallback also failed for {filepath}: {e2}")
                raise
    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.

    Tries pillow-heif first (fast, native). Falls back to ffmpeg for
    files that libheif rejects — e.g. iPhone photos with too many
    auxiliary image references (depth maps, gain maps).
    """
    try:
        img = Image.open(filepath)
        if img.mode != 'RGB':
            img = img.convert('RGB')
        return img
    except Exception as e:
        logger.warning(f"pillow-heif failed for {filepath}: {e} — trying vips")

    # vips fallback: handles tiled Apple HEIC files (bursts, HDR gain
    # maps, depth maps) that pillow-heif/libheif rejects due to too many
    # auxiliary image references.
    import subprocess, tempfile
    try:
        with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as tmp:
            tmp_path = tmp.name
        result = subprocess.run(
            ['vips', 'heifload', filepath, tmp_path],
            capture_output=True, timeout=60, stdin=subprocess.DEVNULL,
        )
        if result.returncode == 0:
            img = Image.open(tmp_path).convert('RGB')
            os.unlink(tmp_path)
            return img
        logger.error(f"vips HEIC decode failed for {filepath}: {result.stderr.decode()[-200:]}")
        os.unlink(tmp_path)
    except Exception as e2:
        logger.error(f"vips fallback failed for {filepath}: {e2}")
    raise RuntimeError(f"Cannot decode HEIC: {filepath}")

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 (AttributeError, KeyError, TypeError):
        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.

    Works on a copy so the caller's image is never mutated — this is
    critical because the thumbnail loop iterates multiple sizes and
    in-place shrinking would degrade later (larger) sizes.
    """
    img = image.copy()
    img.thumbnail((size, size), Image.Resampling.LANCZOS)

    img.save(
        output_path,
        'WEBP',
        quality=settings.thumbnails.quality,
        method=4  # Balance between speed and compression
    )
    img.close()

@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:
        # Declared up front so the except block below can safely check it
        # even if the initial SELECT raises (e.g. asyncpg transport error).
        photo: Optional[Photo] = None
        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)

            # Fallback: some files wear a RAW/HEIC extension but are actually
            # plain JPEGs — e.g. iPhones that write ProRAW-style .DNG for
            # images where no RAW sensor data was captured, or re-exports
            # that kept the original suffix. Pillow can open them directly,
            # so before giving up, try reading the file as a standard image.
            if not image and photo.media_type in ('raw', 'heic'):
                try:
                    image = process_standard_image(photo.filepath)
                    if image is not None:
                        logger.info(
                            f"{photo.filepath}: {photo.media_type} decode failed "
                            f"but file opens as a standard image — using fallback"
                        )
                except Exception as e:
                    logger.debug(
                        f"Standard-image fallback failed for {photo.filepath}: {e}"
                    )

            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 only the sizes the worker still owns on disk
            # (see WORKER_THUMB_SIZES above). The API serves the rest
            # via Nextcloud's preview endpoint.
            for size_name, size_value in THUMB_SIZES.items():
                if size_name not in WORKER_THUMB_SIZES:
                    continue
                thumb_path = get_thumb_path(photo_id, size_name, photo.user_id)
                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 the session is in a bad state (e.g.
            # the original failure was a transport error) rollback first so
            # the status write has a clean transaction to commit into.
            try:
                await session.rollback()
            except Exception:
                pass

            if photo is not None:
                try:
                    photo.processing_status = 'failed'
                    photo.processing_error = str(e)
                    await session.commit()
                except Exception:
                    logger.exception(
                        f"Could not mark photo {photo_id} as failed"
                    )

            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.

    Queue order matters on first-boot and recovery runs: we dispatch
    newest-first (by EXIF taken_at, fallback added_at) so the user's
    most recent photos become fully-indexed before the 2012 archive even
    starts. Picking up the library in pipeline order means the grid,
    timeline and All Photos view populate top-down instead of the worker
    chewing through random insertion-order rows while the UI still
    shows grey placeholders.
    """
    async with AsyncSessionLocal() as session:
        # Get all photos that need thumbnails, newest first.
        result = await session.execute(
            select(Photo)
            .where(Photo.processing_status.in_(['pending', 'failed']))
            .order_by(
                Photo.taken_at.desc().nullslast(),
                Photo.added_at.desc().nullslast(),
            )
        )
        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:
            # Newest-first so the recent end of the library gets phashes
            # (and therefore duplicate detection) ahead of the archive.
            result = await session.execute(
                select(Photo)
                .where(Photo.phash.is_(None))
                .where(Photo.processing_status == 'completed')
                .order_by(
                    Photo.taken_at.desc().nullslast(),
                    Photo.added_at.desc().nullslast(),
                )
                .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',
    # Full regroup scales with O(N²) on phash plus one pgvector query per
    # embedded photo. On a 16k-photo library that's comfortably past the
    # default 5-minute soft limit — bump to 2h / 2h30m. (Passing None here
    # does NOT disable limits; Celery falls back to the worker default
    # of 300s/600s. An explicit number overrides.)
    soft_time_limit=7200,
    time_limit=9000,
)
def regroup_duplicates_task():
    """Full recompute of duplicate groups (pHash + CLIP similarity).

    Used by the Settings → Re-detect duplicates button."""
    from app.services.duplicates import regroup_duplicates
    return asyncio.run(regroup_duplicates())


@shared_task(
    name='incremental_regroup_duplicates',
    # O(new × N); still cheaper than a full regroup but can easily exceed
    # the 5-minute default after a big batch import. Same caveat as
    # regroup_duplicates above — None would just re-inherit the worker
    # default, so we pass explicit values.
    soft_time_limit=3600,
    time_limit=4200,
)
def incremental_regroup_duplicates_task(since_iso: str | None = None):
    """Incremental duplicate detection for newly added photos.

    Compares only photos added after `since_iso` against the full library
    using CLIP vector similarity (O(new × log N) via HNSW) plus pHash.
    Default post-scan path — much faster than a full regroup."""
    from app.services.duplicates import incremental_regroup
    from datetime import datetime, timezone
    since = None
    if since_iso:
        since = datetime.fromisoformat(since_iso)
    return asyncio.run(incremental_regroup(since=since))