"""
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.
"""
import asyncio
import logging
from pathlib import Path

import numpy as np
from celery import shared_task
from sqlalchemy import select, delete, text
from PIL import Image

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

logger = logging.getLogger(__name__)


def _load_thumb(photo_id: str, size: str = "medium") -> np.ndarray | None:
    """Load a thumbnail as an RGB numpy array."""
    thumb_path = Path(f"/data/thumbs/{photo_id}/{size}.webp")
    if not thumb_path.exists():
        logger.warning("Thumbnail not found: %s", thumb_path)
        return None
    img = Image.open(thumb_path).convert("RGB")
    return np.array(img)


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


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

    from app.services.vision.registry import registry
    embedder = registry.get_embedder()
    vector = embedder.embed_image(image)

    model_name = settings.vision.embedder.name

    async with AsyncSessionLocal() as session:
        # Upsert: delete existing then insert
        await 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)
        await session.commit()

    logger.info("Embedded photo %s with %s", 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)

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


@shared_task(name='ocr_photo', queue='vision')
def ocr_photo(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'}
    return asyncio.run(_ocr_photo_async(photo_id))


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

    from app.services.vision.registry import registry
    ocr_engine = registry.get_ocr()
    results = ocr_engine.run(image)

    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

    async with AsyncSessionLocal() as session:
        # Delete existing OCR results for this photo (re-run safe)
        await 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,
            ))
        await session.commit()

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


@shared_task(name='detect_objects', queue='vision')
def detect_objects(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'}
    return asyncio.run(_detect_objects_async(photo_id))


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

    from app.services.vision.registry import registry
    detector = registry.get_detector()
    detections = detector.detect(image)

    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.tags import Tag, photo_tags

    source_name = "vision:yolov8n"

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

        for det in detections:
            # Find or create the object tag
            result = await session.execute(
                select(Tag).where(Tag.name == det.label, Tag.kind == 'object')
            )
            tag = result.scalar_one_or_none()
            if not tag:
                tag = Tag(name=det.label, kind='object', source=source_name)
                session.add(tag)
                await session.flush()  # get tag.id

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

        await session.commit()

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


@shared_task(name='extract_faces', queue='vision')
def extract_faces(photo_id: str):
    """Detect faces and store recognition embeddings. Clustering is
    handled separately by recluster_faces (periodic task)."""
    if not settings.vision.enabled or not settings.vision.faces.enabled:
        return {'status': 'skipped', 'reason': 'faces disabled'}
    return asyncio.run(_extract_faces_async(photo_id))


async def _extract_faces_async(photo_id: str):
    image = _load_thumb(photo_id, "large")  # 1280px for better face detection
    if image is None:
        return {'status': 'error', 'message': 'thumbnail not found'}

    from app.services.vision.registry import registry
    face_proc = registry.get_face_processor()
    faces = face_proc.process(image)

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

    from app.models.face_embedding import FaceEmbedding

    async with AsyncSessionLocal() as session:
        # Wipe previous face results for this photo (re-run safe)
        await 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,  # assigned by recluster_faces
            ))
        await session.commit()

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


@shared_task(name='recluster_faces', queue='vision')
def recluster_faces():
    """Run DBSCAN clustering over all face embeddings and assign/create
    Tag(kind=face_cluster) entries. Should be called periodically or
    manually after a batch of new faces is extracted."""
    if not settings.vision.enabled or not settings.vision.faces.enabled:
        return {'status': 'skipped', 'reason': 'faces disabled'}
    return asyncio.run(_recluster_faces_async())


async def _recluster_faces_async():
    from app.models.face_embedding import FaceEmbedding
    from app.models.tags import Tag
    from app.services.vision.clustering import cluster_faces

    async with AsyncSessionLocal() as session:
        # Load all face embeddings
        result = await session.execute(
            select(FaceEmbedding).order_by(FaceEmbedding.created_at)
        )
        face_rows = result.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,
        )

        # Map cluster labels → Tag(kind=face_cluster)
        cluster_tag_map: dict[int, str] = {}
        source_name = "vision:sface"

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

            if label not in cluster_tag_map:
                # Check if a cluster tag already exists for this cluster
                cluster_name = f"Person {label + 1}"
                tag_result = await session.execute(
                    select(Tag).where(
                        Tag.kind == 'face_cluster',
                        Tag.source == source_name,
                        Tag.name == cluster_name,
                    )
                )
                tag = tag_result.scalar_one_or_none()
                if not tag:
                    tag = Tag(
                        name=cluster_name,
                        kind='face_cluster',
                        source=source_name,
                        representative_photo_id=face_rows[i].photo_id,
                    )
                    session.add(tag)
                    await session.flush()
                cluster_tag_map[label] = tag.id

            face_rows[i].cluster_id = cluster_tag_map[label]

        await session.commit()

    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')
def backfill_vision(task: str | None = None, limit: int | None = None):
    """Queue vision tasks for photos that haven't been processed yet."""
    return asyncio.run(_backfill_vision_async(task, limit))


async def _backfill_vision_async(task: str | None, limit: int | None):
    model_name = settings.vision.embedder.name

    async with AsyncSessionLocal() as session:
        # Find photos without embeddings
        stmt = text("""
            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'
            ORDER BY p.created_at DESC
        """)
        if limit:
            stmt = text(str(stmt) + f" LIMIT {limit}")

        result = await session.execute(stmt, {"model": model_name})
        photo_ids = [row[0] for row in result.fetchall()]

    count = 0
    for pid in photo_ids:
        if task == 'embed' or task is None:
            embed_photo.delay(pid)
        if task == 'ocr' or task is None:
            ocr_photo.delay(pid)
        if task == 'detect' or task is None:
            detect_objects.delay(pid)
        if task == 'faces' or task is None:
            extract_faces.delay(pid)
        count += 1

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