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

"""
Celery tasks for the vision pipeline — embedding, OCR, object detection,
face recognition.

All tasks run on the dedicated `vision` queue with limited concurrency
(memory-bound CPU inference). They read thumbnails generated by
generate_thumbnails, so they MUST run after thumbs complete.

DB access uses sync psycopg2 sessions (not asyncpg) because Celery
forks workers and asyncpg connections can't be shared across forks.
"""
import logging
from pathlib import Path

import numpy as np
from celery import shared_task
from sqlalchemy import create_engine, text as sa_text, select, delete
from sqlalchemy.orm import Session, sessionmaker
from PIL import Image

from app.models.embeddings import Embedding
from app.config import settings

logger = logging.getLogger(__name__)

VISION_READY_KEY = "mulita:vision:ready"


def _vision_worker_ready() -> bool:
    """Check whether the vision worker has finished model bootstrap."""
    try:
        import redis as _redis
        return bool(_redis.from_url(settings.redis_url).exists(VISION_READY_KEY))
    except Exception:
        return False


_sync_engine = None


def _get_sync_engine():
    """Return a module-level singleton engine (one per worker process)."""
    global _sync_engine
    if _sync_engine is None:
        sync_url = settings.database_url.replace("+asyncpg", "+psycopg2").replace("+aiosqlite", "")
        _sync_engine = create_engine(sync_url, pool_pre_ping=True, pool_size=3, max_overflow=5)
    return _sync_engine


def _get_sync_session() -> Session:
    """Create a sync DB session backed by the shared engine."""
    return sessionmaker(bind=_get_sync_engine())()


def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None:
    """Load a thumbnail as an RGB numpy array.

    Thumbnails may live at ``/data/thumbs/{photo_id}/`` (legacy) or
    ``/data/thumbs/{user_id}/{photo_id}/`` (multi-user). Try both.
    """
    thumb_base = Path("/data/thumbs")
    # Try legacy flat path first.
    thumb_path = thumb_base / photo_id / f"{size}.webp"
    if not thumb_path.exists():
        # Try user-prefixed paths: /data/thumbs/*/photo_id/size.webp
        matches = list(thumb_base.glob(f"*/{photo_id}/{size}.webp"))
        if matches:
            thumb_path = matches[0]
        else:
            logger.warning("Thumbnail not found: %s", thumb_path)
            return None
    try:
        img = Image.open(thumb_path).convert("RGB")
        img.load()  # force decode to catch corruption early
        arr = np.array(img)
        img.close()
        return arr
    except Exception as e:
        logger.warning("Corrupt or unreadable thumbnail for %s: %s", photo_id, e)
        return None


@shared_task(name='embed_photo', queue='vision', bind=True, max_retries=3)
def embed_photo(self, photo_id: str):
    """Generate CLIP embedding for a photo and store in pgvector."""
    if not settings.vision.enabled:
        return {'status': 'skipped', 'reason': 'vision disabled'}

    image = _load_thumb(photo_id, "medium")  # 640px
    if image is None:
        return {'status': 'error', 'message': 'thumbnail not found'}

    try:
        from app.services.vision.registry import registry
        embedder = registry.get_embedder()
        vector = embedder.embed_image(image)
    except Exception as exc:
        logger.exception("embed_photo failed for %s", photo_id)
        raise self.retry(exc=exc, countdown=60)

    model_name = settings.vision.embedder.name

    session = _get_sync_session()
    try:
        session.execute(
            delete(Embedding).where(
                Embedding.photo_id == photo_id,
                Embedding.model == model_name,
            )
        )
        emb = Embedding(
            photo_id=photo_id,
            model=model_name,
            vector=vector.tolist(),
        )
        session.add(emb)
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

    logger.info("[%s] Embedded photo %s with %s", self.request.id, photo_id, model_name)
    return {'status': 'success', 'photo_id': photo_id}


@shared_task(name='vision_fanout', queue='vision')
def vision_fanout(photo_id: str):
    """Dispatch all enabled vision tasks for a photo."""
    if not settings.vision.enabled:
        return {'status': 'skipped', 'reason': 'vision disabled'}

    embed_photo.delay(photo_id)

    if settings.vision.ocr.enabled:
        ocr_photo.delay(photo_id)
    if settings.vision.detector.enabled:
        detect_objects.delay(photo_id)
    if settings.vision.faces.enabled:
        extract_faces.delay(photo_id)
    if settings.vision.classifier.enabled:
        classify_content.delay(photo_id)

    return {'status': 'dispatched', 'photo_id': photo_id}


@shared_task(name='ocr_photo', queue='vision', bind=True, max_retries=3)
def ocr_photo(self, photo_id: str):
    """Run OCR on a photo and store text regions."""
    if not settings.vision.enabled or not settings.vision.ocr.enabled:
        return {'status': 'skipped', 'reason': 'OCR disabled'}

    image = _load_thumb(photo_id, "large")  # 1280px for better OCR accuracy
    if image is None:
        return {'status': 'error', 'message': 'thumbnail not found'}

    try:
        from app.services.vision.registry import registry
        ocr_engine = registry.get_ocr()
        results = ocr_engine.run(image)
    except Exception as exc:
        logger.exception("ocr_photo failed for %s", photo_id)
        raise self.retry(exc=exc, countdown=60)

    if not results:
        logger.info("No OCR text found for photo %s", photo_id)
        return {'status': 'success', 'photo_id': photo_id, 'regions': 0}

    from app.models.ocr_text import OCRText

    session = _get_sync_session()
    try:
        session.execute(delete(OCRText).where(OCRText.photo_id == photo_id))
        for r in results:
            session.add(OCRText(
                photo_id=photo_id,
                text=r.text,
                language=r.language,
                confidence=r.confidence,
                bbox=r.bbox,
            ))
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

    logger.info("[%s] OCR: %d text regions for photo %s", self.request.id, len(results), photo_id)
    return {'status': 'success', 'photo_id': photo_id, 'regions': len(results)}


@shared_task(name='detect_objects', queue='vision', bind=True, max_retries=3)
def detect_objects(self, photo_id: str):
    """Detect objects in a photo, create Tag(kind=object) rows, and
    link via photo_tags with confidence/bbox/source."""
    if not settings.vision.enabled or not settings.vision.detector.enabled:
        return {'status': 'skipped', 'reason': 'detection disabled'}

    image = _load_thumb(photo_id, "medium")  # 640px
    if image is None:
        return {'status': 'error', 'message': 'thumbnail not found'}

    try:
        from app.services.vision.registry import registry
        detector = registry.get_detector()
        detections = detector.detect(image)
    except Exception as exc:
        logger.exception("detect_objects failed for %s", photo_id)
        raise self.retry(exc=exc, countdown=60)

    if not detections:
        logger.info("No objects detected for photo %s", photo_id)
        return {'status': 'success', 'photo_id': photo_id, 'objects': 0}

    from app.models import Photo
    from app.models.tags import Tag, photo_tags

    source_name = "vision:yolov8n"

    session = _get_sync_session()
    try:
        # Get the photo's user_id so tags inherit ownership.
        photo = session.execute(
            select(Photo).where(Photo.id == photo_id)
        ).scalar_one_or_none()
        owner_id = photo.user_id if photo else None

        # Wipe previous detection results for this photo from this model
        session.execute(
            delete(photo_tags).where(
                photo_tags.c.photo_id == photo_id,
                photo_tags.c.source == source_name,
            )
        )

        # Group detections by label, keep highest confidence per label
        best_per_label: dict[str, tuple[float, list]] = {}
        for det in detections:
            if det.label not in best_per_label or det.confidence > best_per_label[det.label][0]:
                best_per_label[det.label] = (det.confidence, det.bbox)

        for label, (confidence, bbox) in best_per_label.items():
            # Find or create the object tag (scoped to user)
            tag = session.execute(
                select(Tag).where(Tag.name == label, Tag.kind == 'object', Tag.user_id == owner_id)
            ).scalar_one_or_none()

            if not tag:
                tag = Tag(name=label, kind='object', source=source_name, user_id=owner_id)
                session.add(tag)
                session.flush()  # get tag.id

            # Insert photo_tags association with ML metadata
            session.execute(
                photo_tags.insert().values(
                    photo_id=photo_id,
                    tag_id=tag.id,
                    confidence=confidence,
                    bbox=bbox,
                    source=source_name,
                )
            )

        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

    labels = [d.label for d in detections]
    logger.info("[%s] Detected %d objects in photo %s: %s", self.request.id, len(detections), photo_id, labels)
    return {'status': 'success', 'photo_id': photo_id, 'objects': len(detections)}


@shared_task(name='classify_content', queue='vision', bind=True, max_retries=3)
def classify_content(self, photo_id: str):
    """Classify image content type (screenshot, document, artwork, etc.)
    using CLIP zero-shot classification. Writes Tag(kind=content_type)."""
    if not settings.vision.enabled or not settings.vision.classifier.enabled:
        return {'status': 'skipped', 'reason': 'classifier disabled'}

    image = _load_thumb(photo_id, "medium")
    if image is None:
        return {'status': 'error', 'message': 'thumbnail not found'}

    try:
        from app.services.vision.registry import registry
        classifier = registry.get_classifier()
        results = classifier.classify(image)
    except Exception as exc:
        logger.exception("classify_content failed for %s", photo_id)
        raise self.retry(exc=exc, countdown=60)

    if not results:
        logger.info("No confident classification for photo %s", photo_id)
        return {'status': 'success', 'photo_id': photo_id, 'content_type': None}

    from app.models import Photo
    from app.models.tags import Tag, photo_tags

    source_name = "vision:clip_classifier"
    best = results[0]

    session = _get_sync_session()
    try:
        # Get the photo's user_id so tags inherit ownership.
        photo = session.execute(
            select(Photo).where(Photo.id == photo_id)
        ).scalar_one_or_none()
        owner_id = photo.user_id if photo else None

        # Wipe previous classification for this photo
        session.execute(
            delete(photo_tags).where(
                photo_tags.c.photo_id == photo_id,
                photo_tags.c.source == source_name,
            )
        )

        # Find or create content_type tag (scoped to user)
        tag = session.execute(
            select(Tag).where(Tag.name == best.label, Tag.kind == 'content_type', Tag.user_id == owner_id)
        ).scalar_one_or_none()

        if not tag:
            tag = Tag(name=best.label, kind='content_type', source=source_name, user_id=owner_id)
            session.add(tag)
            session.flush()

        session.execute(
            photo_tags.insert().values(
                photo_id=photo_id,
                tag_id=tag.id,
                confidence=best.confidence,
                source=source_name,
            )
        )
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

    logger.info("[%s] Classified photo %s as '%s' (%.2f)", self.request.id, photo_id, best.label, best.confidence)
    return {'status': 'success', 'photo_id': photo_id, 'content_type': best.label}


def _load_original(photo_id: str) -> np.ndarray | None:
    """Load the original photo file as an RGB numpy array, resized to
    max 1280px on the longest edge for face detection."""
    from sqlalchemy import create_engine, select as sa_select, text as sa_text
    from app.models import Photo

    session = _get_sync_session()
    try:
        photo = session.execute(
            sa_select(Photo).where(Photo.id == photo_id)
        ).scalar_one_or_none()
        if not photo or not photo.filepath:
            return None
        filepath = photo.filepath
    finally:
        session.close()

    if not Path(filepath).exists():
        logger.warning("Original file not found: %s", filepath)
        return None

    try:
        img = Image.open(filepath).convert("RGB")
        # Cap at 4000px on longest edge to avoid OOM, but keep as large
        # as possible for face detection accuracy
        max_dim = 4000
        w, h = img.size
        if max(w, h) > max_dim:
            scale = max_dim / max(w, h)
            resized = img.resize((int(w * scale), int(h * scale)), Image.BICUBIC)
            img.close()
            img = resized
        arr = np.array(img)
        img.close()
        return arr
    except Exception as e:
        logger.warning("Failed to load original %s: %s", filepath, e)
        return None


@shared_task(name='extract_faces', queue='vision', bind=True, max_retries=3)
def extract_faces(self, photo_id: str):
    """Detect faces and store recognition embeddings using InsightFace
    (RetinaFace + ArcFace). No YOLO workaround needed — RetinaFace has
    strong human-vs-non-human precision on its own."""
    if not settings.vision.enabled or not settings.vision.faces.enabled:
        return {'status': 'skipped', 'reason': 'faces disabled'}

    image = _load_original(photo_id)
    if image is None:
        image = _load_thumb(photo_id, "large")
    if image is None:
        return {'status': 'error', 'message': 'no image available'}

    try:
        from app.services.vision.registry import registry
        face_proc = registry.get_face_processor()
        faces = face_proc.process(image)
    except Exception as exc:
        logger.exception("extract_faces failed for %s", photo_id)
        raise self.retry(exc=exc, countdown=60)

    if not faces:
        logger.info("No faces detected for photo %s", photo_id)

    return _save_faces(photo_id, faces)


def _save_faces(photo_id: str, faces) -> dict:
    from app.models.face_embedding import FaceEmbedding

    session = _get_sync_session()
    try:
        session.execute(delete(FaceEmbedding).where(FaceEmbedding.photo_id == photo_id))
        for face in faces:
            session.add(FaceEmbedding(
                photo_id=photo_id,
                bbox=face.bbox,
                vector=face.embedding.tolist(),
                quality=face.quality,
                cluster_id=None,
            ))
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

    if faces:
        logger.info("Extracted %d verified face(s) from photo %s", len(faces), photo_id)
        _schedule_recluster_debounced()
    return {'status': 'success', 'photo_id': photo_id, 'faces': len(faces)}


RECLUSTER_DEBOUNCE_KEY = "mule:recluster_faces:pending"
RECLUSTER_DELAY = 120  # seconds after last face extraction


def _schedule_recluster_debounced():
    """Schedule a recluster_faces run, debounced so rapid-fire face
    extractions don't spawn hundreds of redundant cluster jobs."""
    try:
        import redis as _redis
        r = _redis.from_url(settings.redis_url)
        already_pending = r.set(RECLUSTER_DEBOUNCE_KEY, "1",
                                ex=RECLUSTER_DELAY, nx=True)
        if already_pending:
            recluster_faces.apply_async(countdown=RECLUSTER_DELAY)
            logger.info("Scheduled debounced recluster_faces in %ds", RECLUSTER_DELAY)
    except Exception as e:
        logger.debug("recluster debounce check failed: %s", e)


@shared_task(name='recluster_faces', queue='vision', bind=True, max_retries=10)
def recluster_faces(self):
    """Run DBSCAN clustering over all face embeddings and assign/create
    Tag(kind=face_cluster) entries."""
    # Clear debounce key so new face extractions can schedule another round.
    try:
        import redis as _redis
        _redis.from_url(settings.redis_url).delete(RECLUSTER_DEBOUNCE_KEY)
    except Exception:
        pass

    if not _vision_worker_ready():
        logger.info("Vision worker not ready yet — retrying in 30s")
        raise self.retry(countdown=30)

    if not settings.vision.enabled or not settings.vision.faces.enabled:
        return {'status': 'skipped', 'reason': 'faces disabled'}

    from app.models import Photo
    from app.models.face_embedding import FaceEmbedding
    from app.models.tags import Tag, photo_tags
    from app.services.vision.clustering import cluster_faces

    source_name = "vision:sface"

    session = _get_sync_session()
    try:
        face_rows = session.execute(
            select(FaceEmbedding).order_by(FaceEmbedding.created_at)
        ).scalars().all()

        if len(face_rows) < 2:
            logger.info("Not enough faces for clustering (%d)", len(face_rows))
            return {'status': 'success', 'clusters': 0}

        embeddings = np.array([f.vector for f in face_rows], dtype=np.float32)
        labels = cluster_faces(embeddings, eps=settings.vision.faces.cluster_eps)

        # Clean up old face_cluster tags and their photo_tags
        old_cluster_tags = session.execute(
            select(Tag).where(Tag.kind == 'face_cluster', Tag.source == source_name)
        ).scalars().all()
        for old_tag in old_cluster_tags:
            session.execute(
                delete(photo_tags).where(
                    photo_tags.c.tag_id == old_tag.id,
                    photo_tags.c.source == source_name,
                )
            )
            session.delete(old_tag)
        session.flush()

        # Build new clusters
        cluster_tag_map: dict[int, str] = {}
        # Track which photos belong to which cluster
        cluster_photos: dict[int, set[str]] = {}

        for i, label in enumerate(labels):
            if label == -1:
                face_rows[i].cluster_id = None
                continue

            if label not in cluster_photos:
                cluster_photos[label] = set()
            cluster_photos[label].add(face_rows[i].photo_id)

            if label not in cluster_tag_map:
                cluster_name = f"Person {label + 1}"
                # Inherit user_id from the representative photo.
                rep_photo = session.execute(
                    select(Photo.user_id).where(Photo.id == face_rows[i].photo_id)
                ).scalar_one_or_none()
                tag = Tag(
                    name=cluster_name,
                    kind='face_cluster',
                    source=source_name,
                    representative_photo_id=face_rows[i].photo_id,
                    user_id=rep_photo,
                )
                session.add(tag)
                session.flush()
                cluster_tag_map[label] = tag.id

            face_rows[i].cluster_id = cluster_tag_map[label]

        # Write photo_tags associations so the tag count and tag_ids
        # filter work for face clusters
        for label, photo_ids in cluster_photos.items():
            tag_id = cluster_tag_map[label]
            for pid in photo_ids:
                session.execute(
                    photo_tags.insert().values(
                        photo_id=pid,
                        tag_id=tag_id,
                        source=source_name,
                    )
                )

        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

    n_clusters = len(cluster_tag_map)
    logger.info("Face clustering: %d clusters from %d faces", n_clusters, len(face_rows))
    return {'status': 'success', 'clusters': n_clusters, 'faces': len(face_rows)}


@shared_task(name='backfill_vision', bind=True, max_retries=10)
def backfill_vision(self, task: str | None = None, limit: int | None = None):
    """Queue vision tasks for photos that haven't been processed yet.
    Uses a sync DB connection to avoid asyncpg conflicts in Celery."""
    if not _vision_worker_ready():
        logger.info("Vision worker not ready yet — retrying in 30s")
        raise self.retry(countdown=30)

    model_name = settings.vision.embedder.name
    ordering = "ORDER BY p.taken_at DESC NULLS LAST, p.added_at DESC NULLS LAST"
    limit_clause = " LIMIT :lim" if limit else ""
    params: dict = {"model": model_name}
    if limit:
        params["lim"] = int(limit)

    session = _get_sync_session()
    try:
        # Each query finds photos missing a specific pipeline output so
        # enabling a new processor after import still back-fills.
        embed_ids = []
        if task in ('embed', None):
            sql = f"""
                SELECT p.id FROM photos p
                LEFT JOIN embeddings e ON e.photo_id = p.id AND e.model = :model
                WHERE e.photo_id IS NULL AND p.processing_status = 'completed'
                {ordering}{limit_clause}
            """
            embed_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]

        ocr_ids = []
        if task in ('ocr', None) and settings.vision.ocr.enabled:
            sql = f"""
                SELECT p.id FROM photos p
                LEFT JOIN ocr_text o ON o.photo_id = p.id
                WHERE o.photo_id IS NULL AND p.processing_status = 'completed'
                {ordering}{limit_clause}
            """
            ocr_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]

        detect_ids = []
        if task in ('detect', None) and settings.vision.detector.enabled:
            sql = f"""
                SELECT p.id FROM photos p
                WHERE p.processing_status = 'completed'
                AND NOT EXISTS (
                    SELECT 1 FROM photo_tags pt WHERE pt.photo_id = p.id
                    AND pt.source = 'vision:yolov8n'
                )
                {ordering}{limit_clause}
            """
            detect_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]

        face_ids = []
        if task in ('faces', None) and settings.vision.faces.enabled:
            sql = f"""
                SELECT p.id FROM photos p
                LEFT JOIN face_embeddings fe ON fe.photo_id = p.id
                WHERE fe.photo_id IS NULL AND p.processing_status = 'completed'
                {ordering}{limit_clause}
            """
            face_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]

        classify_ids = []
        if task in ('classify', None) and settings.vision.classifier.enabled:
            sql = f"""
                SELECT p.id FROM photos p
                WHERE p.processing_status = 'completed'
                AND NOT EXISTS (
                    SELECT 1 FROM photo_tags pt WHERE pt.photo_id = p.id
                    AND pt.source = 'vision:clip_classifier'
                )
                {ordering}{limit_clause}
            """
            classify_ids = [r[0] for r in session.execute(sa_text(sql), params).fetchall()]
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()

    # Dispatch — deduplicate across query results.
    all_ids = set(embed_ids) | set(ocr_ids) | set(detect_ids) | set(face_ids) | set(classify_ids)
    for pid in embed_ids:
        embed_photo.delay(pid)
    for pid in ocr_ids:
        ocr_photo.delay(pid)
    for pid in detect_ids:
        detect_objects.delay(pid)
    for pid in face_ids:
        extract_faces.delay(pid)
    for pid in classify_ids:
        classify_content.delay(pid)

    logger.info("Backfill queued %d photos for vision processing", len(all_ids))
    return {'status': 'queued', 'count': len(all_ids)}