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, timezone
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
from app.services.date_guess import has_date_warning

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

    PROGRESS_TTL = 3600  # 1 hour — auto-expire if scan crashes

    def progress_set(key: str, value) -> None:
        if r is None:
            return
        try:
            r.set(key, str(value), ex=PROGRESS_TTL)
        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
            else:
                source_root = (await session.execute(
                    select(SourceRoot).where(SourceRoot.id == source_root_id)
                )).scalar_one_or_none()

            # Inherit user_id from the source root's owner
            owner_user_id = source_root.user_id if source_root else None

            # Per-scan memoization cache for "is this folder's effective
            # is_hidden true?" Populated on first lookup by walking the
            # parent_id chain up to the source root. Keyed by folder_id
            # so repeated photos in the same folder pay only one lookup.
            hidden_folder_cache: dict[str, bool] = {}

            async def is_folder_effectively_hidden(folder_row: Folder) -> bool:
                if folder_row.id in hidden_folder_cache:
                    return hidden_folder_cache[folder_row.id]
                # Walk parents. If the current folder is hidden, short-
                # circuit. Otherwise climb until we hit a root (no
                # parent_id) or a cached ancestor.
                if folder_row.is_hidden:
                    hidden_folder_cache[folder_row.id] = True
                    return True
                parent_id = folder_row.parent_id
                while parent_id is not None:
                    if parent_id in hidden_folder_cache:
                        hidden_folder_cache[folder_row.id] = hidden_folder_cache[parent_id]
                        return hidden_folder_cache[folder_row.id]
                    parent = (
                        await session.execute(
                            select(Folder).where(Folder.id == parent_id)
                        )
                    ).scalar_one_or_none()
                    if parent is None:
                        break
                    if parent.is_hidden:
                        hidden_folder_cache[folder_row.id] = True
                        return True
                    parent_id = parent.parent_id
                hidden_folder_cache[folder_row.id] = False
                return False

            # 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, owner_user_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)
                            )
                            existing_photo = existing.scalar_one_or_none()
                            if existing_photo is not None:
                                # Resurrect a previously-discarded row if the
                                # file is back on disk. WebDAV DELETE +
                                # re-upload, trashbin restore via PUT-overwrite,
                                # and any "I removed it then put it back" flow
                                # all land here. Re-queue extract_metadata in
                                # case the bytes changed (different EXIF, new
                                # nextcloud_fileid).
                                if existing_photo.is_discarded:
                                    existing_photo.is_discarded = False
                                    existing_photo.discarded_at = None
                                    await session.commit()
                                    logger.info(
                                        f"Resurrected discarded photo on rescan: {filepath}"
                                    )
                                    extract_metadata.delay(existing_photo.id)
                                else:
                                    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

                            # Inherit the effective-hidden flag from the
                            # folder's ancestry. If any ancestor folder
                            # has is_hidden=true, the new photo is
                            # immediately marked hidden so it never
                            # briefly appears in cross-cutting views
                            # between scan and the next manual recompute.
                            effective_hidden = await is_folder_effectively_hidden(folder)

                            # Create photo entry
                            mtime_dt = datetime.fromtimestamp(stat.st_mtime)
                            photo = Photo(
                                filepath=filepath,
                                filename=filename,
                                folder_id=folder.id,
                                user_id=owner_user_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=mtime_dt,
                                taken_at_source='filesystem',
                                # First-pass flag based on the filesystem mtime;
                                # metadata.extract_metadata re-runs this once
                                # EXIF has been parsed so a real DateTimeOriginal
                                # can clear the warning.
                                has_date_warning=has_date_warning(filepath, mtime_dt),
                                is_duplicate=is_dup,
                                is_hidden=effective_hidden,
                                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, user_id: str = None
) -> 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, user_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,
            user_id=user_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).

    After dispatching the scans, queue a delayed `regroup_duplicates`
    pass so duplicate clusters are recomputed once the new photos have
    finished thumbnailing (and therefore picked up phashes). The
    countdown is a best-effort hint — on a big library the user can
    still hit Settings → Re-detect duplicates to force a fresh pass.
    """
    from app.tasks.thumbs import incremental_regroup_duplicates_task
    from app.tasks.vision import backfill_vision

    async with AsyncSessionLocal() as session:
        result = await session.execute(
            select(SourceRoot).where(SourceRoot.is_active == True)  # noqa: E712
        )
        source_roots = result.scalars().all()
        dispatched = 0
        for sr in source_roots:
            if os.path.exists(sr.path):
                scan_folder.delay(sr.path, sr.id)
                dispatched += 1
            else:
                logger.warning(f"Source root path does not exist: {sr.path}")

    if dispatched > 0:
        # 60s gives the thumbs worker a window to compute phashes for
        # the new photos before regrouping. The task is idempotent, so
        # firing too early just means the next manual run picks up the
        # late arrivals — no corrupted state.
        try:
            # Use incremental mode: only compare newly added photos
            # against the full library via CLIP HNSW + pHash.
            # O(new × log N) instead of O(N²).
            scan_start = datetime.now(timezone.utc).isoformat()
            incremental_regroup_duplicates_task.apply_async(
                kwargs={'since_iso': scan_start},
                countdown=60,
            )
        except Exception as e:
            logger.warning(f"Could not queue post-scan regroup: {e}")

        # 90s lets thumbnails finish so photos reach processing_status
        # 'completed', which backfill_vision uses as its filter.
        try:
            backfill_vision.apply_async(countdown=90)
        except Exception as e:
            logger.warning(f"Could not queue post-scan vision backfill: {e}")

        # NOTE: we used to auto-queue `backfill_gps` here so photos
        # scanned before the GPS-extraction fix would eventually get
        # their coordinates populated. That fix shipped a long time
        # ago, so on every modern restart it just re-ran
        # extract_metadata for every photo that legitimately has no
        # GPS in EXIF (screenshots, indoor shots, scans) — tens of
        # thousands of pointless tasks that saturated worker-light
        # for ~30 min after each deploy. Trigger manually via
        # POST /api/v1/library/backfill-gps if you ever need it
        # again (e.g. another extractor-logic fix lands).


@shared_task(name='watch_folders', bind=True)
def watch_folders(self):
    """Retired: file events now arrive via NC webhook_listeners.

    Kept as a no-op task so any in-flight queue items (a leftover
    apply_async from a restart before this commit, or an admin-button
    trigger) don't crash workers. Will be removed entirely once the
    queue drains.
    """
    logger.info(
        "watch_folders task is retired; file events come from NC "
        "webhook_listeners. No-op."
    )
    return {'status': 'retired'}

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


async def handle_directory_deletion(dirpath: str) -> int:
    """Mark every Photo under `dirpath` as discarded — used when Nextcloud
    fires a NodeDeletedEvent on a folder. NC emits ONE event for the
    folder itself (not one per child file), so without this we'd never
    see the children disappear except via the 30-min reconcile sweep.

    Returns the number of photos affected. Matches by `filepath LIKE
    dirpath + '/%'` (the trailing slash is important — we don't want
    `/photos/foo` to also match `/photos/foobar.jpg`).
    """
    from sqlalchemy import update

    prefix = dirpath.rstrip("/") + "/"
    async with AsyncSessionLocal() as session:
        result = await session.execute(
            update(Photo)
            .where(
                Photo.filepath.like(prefix + "%"),
                Photo.is_discarded.is_(False),
            )
            .values(is_discarded=True, discarded_at=datetime.utcnow())
        )
        await session.commit()
        n = result.rowcount or 0
    if n:
        logger.info(f"Marked {n} photos as discarded under {dirpath}")
    return n


async def handle_directory_rename(old_dirpath: str, new_dirpath: str) -> dict:
    """Reflect a Nextcloud-side folder rename in mule's DB.

    NC emits a single NodeRenamedEvent on the directory — children
    don't get their own events. We mirror the same prefix-rewrite the
    PATCH /folders/{id} endpoint does inline, so heaps, tags, ratings,
    and other state keyed on Photo.id survive intact.

    Same-source-root case (the common one): prefix-rewrite filepath /
    path on photos, folders, source_roots in one transaction.

    Cross-source-root case (folder moved between two registered roots,
    e.g. Photos/x → Memories/x): discard the old subtree and rely on
    the scan_folder dispatched by a NodeWritten/NodeCreated event (or
    the 30-min reconcile sweep) to add fresh Photo rows under the new
    root. Mirrors the "different boundary, different identity" model
    mule has elsewhere.

    Idempotent: re-running with the same args is a no-op because no
    row matches `LIKE old_prefix||'/%'` after the first pass. That
    makes the feedback loop (mule PATCH → WebDAV MOVE → NC webhook →
    handler) safe.
    """
    from sqlalchemy import or_, text, update
    from app.models.folders import SourceRoot

    old_prefix = old_dirpath.rstrip("/")
    new_prefix = new_dirpath.rstrip("/")
    if not old_prefix or not new_prefix or old_prefix == new_prefix:
        return {"status": "noop"}

    async def _source_root_id_for(session, path: str) -> Optional[str]:
        """Find the active SourceRoot whose path is a prefix of `path`."""
        roots = (await session.execute(
            select(SourceRoot).where(SourceRoot.is_active == True)  # noqa: E712
        )).scalars().all()
        norm = os.path.normpath(path)
        for sr in roots:
            root = os.path.normpath(sr.path)
            if norm == root or norm.startswith(root + os.sep):
                return sr.id
        return None

    async with AsyncSessionLocal() as session:
        old_root_id = await _source_root_id_for(session, old_prefix)
        new_root_id = await _source_root_id_for(session, new_prefix)

        # Cross-root rename: discard old subtree; let webhook-dispatched
        # scan_folder add fresh rows under the new root.
        if old_root_id and new_root_id and old_root_id != new_root_id:
            result = await session.execute(
                update(Photo)
                .where(
                    Photo.filepath.like(old_prefix + "/%"),
                    Photo.is_discarded.is_(False),
                )
                .values(is_discarded=True, discarded_at=datetime.utcnow())
            )
            await session.commit()
            n = result.rowcount or 0
            logger.info(
                f"Cross-root rename {old_prefix} -> {new_prefix}: "
                f"discarded {n} photos in old root"
            )
            return {"status": "cross_root", "discarded": n}

        # Same-root: iterate the matching rows in Python and rewrite
        # the prefix attribute-side. We tried a single UPDATE … SET …
        # SUBSTRING(... FROM LENGTH(:old)+1) raw-SQL approach but
        # asyncpg miscategorises the LENGTH() result and rejects it
        # as "$2: int (expected str)". The PATCH /folders/{id}
        # endpoint already loops in Python for the same reason — match
        # its pattern. Folder renames are rare and typically span ≤1k
        # photos, so per-row UPDATEs are fine.
        old_pat = old_prefix + "/%"
        photos = (await session.execute(
            select(Photo).where(Photo.filepath.like(old_pat))
        )).scalars().all()
        for p in photos:
            p.filepath = new_prefix + p.filepath[len(old_prefix):]

        folders = (await session.execute(
            select(Folder).where(
                or_(
                    Folder.path == old_prefix,
                    Folder.path.like(old_pat),
                )
            )
        )).scalars().all()
        for f in folders:
            f.path = new_prefix if f.path == old_prefix else \
                new_prefix + f.path[len(old_prefix):]

        source_roots = (await session.execute(
            select(SourceRoot).where(SourceRoot.path == old_prefix)
        )).scalars().all()
        for sr in source_roots:
            sr.path = new_prefix
        await session.commit()
        return {
            "status": "renamed",
            "photos": len(photos),
            "folders": len(folders),
            "source_roots": len(source_roots),
        }


@shared_task(name='backfill_gps')
def backfill_gps():
    """Re-run metadata extraction on every non-discarded photo that is
    missing latitude/longitude. Used both as a one-shot kick-off after the
    GPS columns are added on an existing install (see app/database.py) and
    as a manual trigger from POST /api/v1/library/backfill-gps. Each
    extract_metadata call is itself a Celery task, so this just enqueues —
    it does not block on extraction completing."""
    return asyncio.run(_backfill_gps_async())


async def _backfill_gps_async():
    async with AsyncSessionLocal() as session:
        # Newest-first so the most recent photos get their GPS + EXIF
        # written before the worker climbs back through the archive.
        result = await session.execute(
            select(Photo.id)
            .where(
                Photo.latitude.is_(None),
                Photo.is_discarded.is_(False),
            )
            .order_by(
                Photo.taken_at.desc().nullslast(),
                Photo.added_at.desc().nullslast(),
            )
        )
        photo_ids = [row[0] for row in result.all()]

    for pid in photo_ids:
        extract_metadata.delay(pid)

    logger.info(f"backfill_gps: queued extract_metadata for {len(photo_ids)} photos")
    return {'queued': len(photo_ids)}


@shared_task(name='discard_missing_photos_beat')
def discard_missing_photos_beat():
    """Periodic catch-up for filesystem deletions the watcher missed
    (e.g. while the worker was restarting). Walks every active source
    root that is currently mounted and present, and soft-discards any
    Photo whose file is gone. Hard-deletion stays manual via
    POST /api/v1/library/maintenance/prune-missing.

    Wired to a 30-minute beat schedule in app/tasks/celery.py.
    """
    from app.services.cleanup import discard_missing_photos
    return asyncio.run(discard_missing_photos())