"""ML Celery tasks: per-image inference, backfill discovery, centroid
recompute, allowlist auto-apply, model self-heal.

All run on the ml-worker (queue 'ml') except recompute_centroids and
apply_allowlist_tags sweeps which are 'maintenance' lane. Sync sessions
(Celery workers are sync processes), same pattern as FC-2a tasks.
"""

import logging
from pathlib import Path

from celery.exceptions import SoftTimeLimitExceeded
from sqlalchemy import delete, select
from sqlalchemy.exc import DBAPIError, OperationalError

from ..celery_app import celery
from ..models import ImagePrediction, ImageRecord, MLSettings
from ._sync_engine import sync_session_factory as _sync_session_factory

log = logging.getLogger(__name__)

IMAGES_ROOT = Path("/images")
VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv", ".webm", ".m4v", ".wmv", ".flv"}


def _is_video(path: Path) -> bool:
    return path.suffix.lower() in VIDEO_EXTS


@celery.task(
    name="backend.app.tasks.ml.tag_and_embed",
    bind=True,
    autoretry_for=(OperationalError, DBAPIError, OSError),
    retry_backoff=5,
    retry_backoff_max=60,
    retry_jitter=True,
    max_retries=3,
    # Sized for the video branch: sample 6 frames, run tagger +
    # embedder on each (≈12 GPU ops vs 2 for an image). A loaded
    # ml-worker can take 5-10 min on a long video; bumped from
    # 5min/7min on 2026-05-28 after operator-flagged image 6288 (a
    # .mp4) hit the recovery sweep at 5 min while still legitimately
    # processing. Image runs return in seconds; the bump doesn't
    # affect their UX.
    soft_time_limit=900,   # 15 min
    time_limit=1200,       # 20 min hard
)
def tag_and_embed(self, image_id: int) -> dict:
    """Run Camie + SigLIP on one image; store predictions + embedding;
    then enqueue per-image allowlist application.

    Video: sample frames between 10% and 90% of duration (VIDEO_ML_FRAMES,
    default 6). Max-pool tagger confidences across frames, mean-pool the
    SigLIP embeddings. On no-frames returns status='no_frames' (not an error).
    """
    import os
    import time

    from ..services.ml.embedder import get_embedder
    from ..services.ml.tagger import get_tagger

    # Phase + file context, so a timeout/crash names WHICH file and WHERE it
    # died instead of a bare SoftTimeLimitExceeded() (operator-flagged 2026-06-08:
    # the activity told them nothing about the file or why). `ctx` is enriched
    # once the record is loaded; both feed the worker log AND the re-raised
    # exception message (which becomes the activity's error_message).
    started = time.monotonic()
    phase = "open_session"
    ctx = f"image_id={image_id}"

    def _elapsed() -> float:
        return time.monotonic() - started

    try:
        SessionLocal = _sync_session_factory()
        with SessionLocal() as session:
            record = session.get(ImageRecord, image_id)
            if record is None:
                return {"status": "missing", "image_id": image_id}
            settings = session.execute(
                select(MLSettings).where(MLSettings.id == 1)
            ).scalar_one()

            src = Path(record.path)
            is_vid = _is_video(src)
            ctx = (
                f"image_id={image_id} path={record.path} mime={record.mime} "
                f"bytes={record.size_bytes} video={is_vid}"
            )
            log.info("tag_and_embed start: %s", ctx)
            if not src.is_file():
                log.warning("tag_and_embed file missing on disk: %s", ctx)
                return {"status": "file_missing", "image_id": image_id}

            phase = "load_models"
            tagger = get_tagger()
            embedder = get_embedder()

            if is_vid:
                # Layer-3 isolation: ffprobe (a separate process) validates
                # the container before we burn ~20 GPU ops sampling frames
                # from it. A corrupt video that would crash the frame
                # decoder is rejected cleanly here instead of taking down
                # the ml-worker. Operator-flagged 2026-05-28.
                phase = "video_probe"
                from ..utils import safe_probe
                vprobe = safe_probe.probe_video(src)
                if not vprobe.ok:
                    log.warning(
                        "tag_and_embed bad video (%s): %s", vprobe.reason, ctx
                    )
                    return {
                        "status": "bad_video", "image_id": image_id,
                        "reason": vprobe.reason,
                    }
                phase = "video_sample_frames"
                t0 = time.monotonic()
                frames = _sample_video_frames(
                    src, int(os.environ.get("VIDEO_ML_FRAMES", "6"))
                )
                log.info(
                    "tag_and_embed sampled %d frame(s) in %.1fs: %s",
                    len(frames), time.monotonic() - t0, ctx,
                )
                if not frames:
                    return {"status": "no_frames", "image_id": image_id}
                phase = "video_infer"
                import numpy as np

                preds = _maxpool_predictions(
                    [tagger.infer(f, store_floor=settings.tagger_store_floor)
                     for f in frames]
                )
                embedding = np.mean(
                    [embedder.infer(f) for f in frames], axis=0
                ).astype("float32")
                for f in frames:
                    f.unlink(missing_ok=True)
            else:
                phase = "tag"
                t0 = time.monotonic()
                raw = tagger.infer(src, store_floor=settings.tagger_store_floor)
                log.info(
                    "tag_and_embed tagged in %.1fs (%d tags): %s",
                    time.monotonic() - t0, len(raw), ctx,
                )
                preds = {
                    name: {"category": p.category, "confidence": p.confidence}
                    for name, p in raw.items()
                }
                phase = "embed"
                t0 = time.monotonic()
                embedding = embedder.infer(src)
                log.info(
                    "tag_and_embed embedded in %.1fs: %s",
                    time.monotonic() - t0, ctx,
                )

            phase = "persist"
            record.tagger_model_version = settings.tagger_model_version
            record.siglip_embedding = embedding.tolist()
            record.siglip_model_version = settings.embedder_model_version
            session.add(record)
            # Write the normalized image_prediction rows (#768) — the sole home
            # for predictions now (image_record.tagger_predictions was dropped in
            # migration 0046). Delete-then-insert keeps a re-tag idempotent;
            # tagger_store_floor was already applied in tagger.infer, so preds is
            # the >=floor set.
            session.execute(
                delete(ImagePrediction).where(
                    ImagePrediction.image_record_id == image_id
                )
            )
            session.add_all([
                ImagePrediction(
                    image_record_id=image_id, raw_name=name,
                    category=p.get("category", "general"),
                    score=float(p.get("confidence", 0.0)),
                )
                for name, p in preds.items()
            ])
            session.commit()
    except SoftTimeLimitExceeded:
        log.error(
            "tag_and_embed TIMED OUT after %.0fs in phase=%s: %s",
            _elapsed(), phase, ctx,
        )
        # Re-raise as SoftTimeLimitExceeded (preserves the 'timeout' status in
        # the task_run signal) but WITH context, so the activity error_message
        # names the file + phase instead of being empty.
        raise SoftTimeLimitExceeded(
            f"timed out in phase={phase} after {_elapsed():.0f}s ({ctx})"
        ) from None
    except Exception:
        # OSError/DBAPIError/OperationalError are autoretried — re-raise the
        # ORIGINAL so the type is preserved; just make sure it's logged with
        # context first.
        log.exception(
            "tag_and_embed FAILED in phase=%s after %.0fs: %s",
            phase, _elapsed(), ctx,
        )
        raise

    log.info(
        "tag_and_embed ok in %.1fs (%d tags): %s", _elapsed(), len(preds), ctx
    )
    apply_allowlist_tags.delay(image_id=image_id)
    return {"status": "ok", "image_id": image_id, "tags": len(preds)}


def _sample_video_frames(src: Path, n: int) -> list[Path]:
    """Extract n frames evenly between 10% and 90% of duration via ffmpeg.
    Returns temp file paths (caller deletes). Empty list on failure."""
    import json
    import subprocess
    import tempfile

    try:
        probe = subprocess.run(
            [
                "ffprobe", "-v", "quiet", "-print_format", "json",
                "-show_format", str(src),
            ],
            check=True, capture_output=True, timeout=30,
        )
        duration = float(json.loads(probe.stdout)["format"]["duration"])
    except Exception:
        return []
    if duration <= 0:
        return []

    start, end = duration * 0.10, duration * 0.90
    step = (end - start) / max(n - 1, 1)
    out: list[Path] = []
    tmpdir = Path(tempfile.mkdtemp(prefix="fc_vid_"))
    for i in range(n):
        ts = start + i * step
        dest = tmpdir / f"frame_{i:02d}.jpg"
        try:
            subprocess.run(
                [
                    "ffmpeg", "-ss", f"{ts:.2f}", "-i", str(src),
                    "-frames:v", "1", "-q:v", "3", "-y", str(dest),
                ],
                check=True, capture_output=True, timeout=60,
            )
            if dest.is_file():
                out.append(dest)
        except Exception:
            continue
    return out


def _maxpool_predictions(per_frame: list[dict]) -> dict:
    """Aggregate per-frame {name: TagPrediction} dicts by max confidence."""
    merged: dict[str, dict] = {}
    for frame_preds in per_frame:
        for name, p in frame_preds.items():
            cur = merged.get(name)
            if cur is None or p.confidence > cur["confidence"]:
                merged[name] = {
                    "category": p.category,
                    "confidence": p.confidence,
                }
    return merged


@celery.task(name="backend.app.tasks.ml.backfill", bind=True)
def backfill(self) -> int:
    """Enqueue tag_and_embed for images missing predictions/embeddings for
    the current model versions. Keyset pagination by id ASC (restart-safe).
    """
    SessionLocal = _sync_session_factory()
    enqueued = 0
    last_id = 0
    with SessionLocal() as session:
        settings = session.execute(
            select(MLSettings).where(MLSettings.id == 1)
        ).scalar_one()
        while True:
            rows = session.execute(
                select(ImageRecord.id)
                .where(ImageRecord.id > last_id)
                .where(
                    (ImageRecord.tagger_model_version.is_(None))
                    | (
                        ImageRecord.tagger_model_version
                        != settings.tagger_model_version
                    )
                    | (ImageRecord.siglip_embedding.is_(None))
                    | (
                        ImageRecord.siglip_model_version
                        != settings.embedder_model_version
                    )
                )
                .order_by(ImageRecord.id.asc())
                .limit(500)
            ).scalars().all()
            if not rows:
                break
            for image_id in rows:
                tag_and_embed.delay(image_id)
                enqueued += 1
            last_id = rows[-1]
    return enqueued


@celery.task(
    name="backend.app.tasks.ml.apply_allowlist_tags",
    bind=True,
    # Audit 2026-06-02 — the full-sweep mode (neither tag_id nor image_id)
    # is O(images × allowlist) and legitimately runs >5 min on large
    # libraries. Cap matches the maintenance queue's recovery threshold.
    soft_time_limit=1800, time_limit=2100,
)
def apply_allowlist_tags(self, tag_id: int | None = None,
                         image_id: int | None = None) -> int:
    """Retroactively apply allowlisted tags.

    Modes:
    - tag_id only      : scan all images for this tag.
    - image_id only    : scan all allowlisted tags for this image.
    - both             : just the (image, tag) pair.
    - neither          : full sweep (daily beat).

    Skips: already-applied, rejected (tag_suggestion_rejection), or
    confidence below the tag's allowlist min_confidence. Applied with
    source='ml_auto'.
    """
    from sqlalchemy import and_
    from sqlalchemy import select as sa_select
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    from ..models import TagAllowlist, TagSuggestionRejection
    from ..models.tag import image_tag

    SessionLocal = _sync_session_factory()
    applied = 0
    with SessionLocal() as session:
        allow_rows = session.execute(
            sa_select(TagAllowlist.tag_id, TagAllowlist.min_confidence)
            if tag_id is None
            else sa_select(
                TagAllowlist.tag_id, TagAllowlist.min_confidence
            ).where(TagAllowlist.tag_id == tag_id)
        ).all()
        allow = {r[0]: r[1] for r in allow_rows}
        if not allow:
            return 0

        # Images that have any predictions (#768: from image_prediction, not
        # the old JSON column), optionally narrowed to one image.
        img_ids_query = sa_select(ImagePrediction.image_record_id).distinct()
        if image_id is not None:
            img_ids_query = img_ids_query.where(
                ImagePrediction.image_record_id == image_id
            )

        for (img_id,) in session.execute(img_ids_query).all():
            preds = _load_predictions_sync(session, img_id)
            for a_tag_id, min_conf in allow.items():
                exists = session.execute(
                    sa_select(image_tag.c.tag_id).where(
                        and_(
                            image_tag.c.image_record_id == img_id,
                            image_tag.c.tag_id == a_tag_id,
                        )
                    )
                ).scalar_one_or_none()
                if exists is not None:
                    continue
                rej = session.get(
                    TagSuggestionRejection, (img_id, a_tag_id)
                )
                if rej is not None:
                    continue
                from ..models import Tag

                tag = session.get(Tag, a_tag_id)
                if tag is None:
                    continue
                conf = _confidence_for_tag(session, tag, preds)
                if conf is None or conf < min_conf:
                    continue
                stmt = pg_insert(image_tag).values(
                    image_record_id=img_id,
                    tag_id=a_tag_id,
                    source="ml_auto",
                )
                stmt = stmt.on_conflict_do_nothing(
                    index_elements=["image_record_id", "tag_id"]
                )
                session.execute(stmt)
                applied += 1
        session.commit()
    return applied


def _load_predictions_sync(session, image_id: int) -> dict:
    """Predictions for one image from image_prediction (#768), in the
    {raw_name: {category, confidence}} shape _confidence_for_tag consumes —
    keeps the allowlist resolution logic unchanged."""
    from sqlalchemy import select as sa_select

    rows = session.execute(
        sa_select(
            ImagePrediction.raw_name,
            ImagePrediction.category,
            ImagePrediction.score,
        ).where(ImagePrediction.image_record_id == image_id)
    ).all()
    return {
        r.raw_name: {"category": r.category, "confidence": r.score}
        for r in rows
    }


def _confidence_for_tag(session, tag, preds: dict) -> float | None:
    """Highest confidence among predictions that resolve to `tag` —
    either the prediction name equals the tag name, or an alias maps
    (prediction name, category) -> tag.id.
    """
    from sqlalchemy import select as sa_select

    from ..models import TagAlias

    best: float | None = None
    direct = preds.get(tag.name)
    if direct is not None:
        best = float(direct.get("confidence", 0.0))
    alias_rows = session.execute(
        sa_select(TagAlias.alias_string, TagAlias.alias_category).where(
            TagAlias.canonical_tag_id == tag.id
        )
    ).all()
    for alias_string, alias_category in alias_rows:
        p = preds.get(alias_string)
        if p is None:
            continue
        if p.get("category") != alias_category:
            continue
        c = float(p.get("confidence", 0.0))
        if best is None or c > best:
            best = c
    return best


@celery.task(name="backend.app.tasks.ml.recompute_centroid", bind=True)
def recompute_centroid(self, tag_id: int) -> bool:
    import asyncio

    from ..extensions import get_session
    from ..services.ml.centroids import CentroidService

    async def _run() -> bool:
        async with get_session() as session:
            svc = CentroidService(session)
            result = await svc.recompute_for_tag(tag_id)
            await session.commit()
            return result

    return asyncio.run(_run())


@celery.task(
    name="backend.app.tasks.ml.recompute_centroids",
    bind=True,
    # Audit 2026-06-02 — drifted-centroid rebuild over potentially
    # hundreds of tags.
    soft_time_limit=1800, time_limit=2100,
)
def recompute_centroids(self) -> int:
    """Daily: find drifted centroids, enqueue recompute_centroid for each."""
    import asyncio

    from ..extensions import get_session
    from ..services.ml.centroids import CentroidService

    async def _list() -> list[int]:
        async with get_session() as session:
            return await CentroidService(session).list_drifted()

    drifted = asyncio.run(_list())
    for tid in drifted:
        recompute_centroid.delay(tid)
    return len(drifted)