FabledCurator/backend/app/tasks/ml.py

"""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 (#747): sample frames at a fixed cadence (ml_settings
    video_frame_interval_seconds, capped at video_max_frames), keep a tag only if
    it appears in >= video_min_tag_frames frames and average its confidence over
    those frames (mean-pool, not max — kills one-frame noise); mean-pool the
    SigLIP embeddings. On no-frames returns status='no_frames' (not an error).
    """
    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,
                    interval=settings.video_frame_interval_seconds,
                    max_frames=settings.video_max_frames,
                )
                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 = _aggregate_video_predictions(
                    [tagger.infer(f, store_floor=settings.tagger_store_floor)
                     for f in frames],
                    min_frames=settings.video_min_tag_frames,
                )
                embedding = np.mean(
                    [embedder.infer(f) for f in frames], axis=0
                ).astype("float32")
                log.info(
                    "tag_and_embed video aggregated %d tag(s) from %d frame(s) "
                    "(min_frames=%d): %s",
                    len(preds), len(frames), settings.video_min_tag_frames, ctx,
                )
                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, *, interval: float, max_frames: int,
) -> list[Path]:
    """Sample frames at a fixed CADENCE — one every `interval` seconds — so a
    tag's frame-presence reflects real screen time regardless of video length
    (#747). The count is capped at `max_frames`: a video longer than
    interval×max_frames stretches the spacing instead of exploding the frame
    count (keeps cost bounded so a long video can't hog the single ml-worker).
    Frames are taken across the 5%–95% window (skip intro/outro black/cards) via
    per-frame fast-seek. Returns temp file paths (caller deletes); [] 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.05, duration * 0.95
    span = max(end - start, 0.0)
    # Cadence count, clamped to [1, max_frames]. int(span/interval)+1 ≈ one frame
    # per `interval` seconds across the window; the cap stretches spacing on very
    # long videos.
    n = max(1, min(int(span / interval) + 1, max(1, max_frames)))
    step = span / 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:04d}.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=30,
            )
            if dest.is_file():
                out.append(dest)
        except Exception:
            continue
    return out


def _aggregate_video_predictions(per_frame: list[dict], *, min_frames: int) -> dict:
    """Aggregate per-frame {name: TagPrediction} into one prediction set (#747).

    A tag is kept only if it appears (≥ the tagger store floor, already applied)
    in at least `min_frames` of the sampled frames — because sampling is at a
    fixed cadence, that means it was on screen for roughly min_frames×interval
    seconds, so a single-frame flicker / scene-transition artifact is dropped
    while a genuine scene-local tag in a long video survives. Confidence is the
    MEAN over the frames where the tag appears (not max — max re-inflated the
    one-frame noise this whole change exists to remove).

    `min_frames` is clamped to the number of frames actually sampled so a very
    short video (1–2 frames) still tags instead of dropping everything.
    """
    n = len(per_frame)
    if n == 0:
        return {}
    threshold = max(1, min(min_frames, n))
    agg: dict[str, dict] = {}
    for frame_preds in per_frame:
        for name, p in frame_preds.items():
            cur = agg.get(name)
            if cur is None:
                agg[name] = {"category": p.category, "sum": p.confidence, "count": 1}
            else:
                cur["sum"] += p.confidence
                cur["count"] += 1
    return {
        name: {"category": v["category"], "confidence": v["sum"] / v["count"]}
        for name, v in agg.items()
        if v["count"] >= threshold
    }


@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 ..services.ml.centroids import CentroidService
    from ._async_session import async_session_factory

    async def _run() -> bool:
        # Per-task NullPool engine bound to THIS asyncio.run loop — the shared
        # process-wide engine reuses connections across loops and raises
        # "Future attached to a different loop" on every call after the first.
        async_factory, async_engine = async_session_factory()
        try:
            async with async_factory() as session:
                svc = CentroidService(session)
                result = await svc.recompute_for_tag(tag_id)
                await session.commit()
                return result
        finally:
            await async_engine.dispose()

    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 ..services.ml.centroids import CentroidService
    from ._async_session import async_session_factory

    async def _list() -> list[int]:
        # Per-task NullPool engine bound to this loop (see recompute_centroid).
        async_factory, async_engine = async_session_factory()
        try:
            async with async_factory() as session:
                return await CentroidService(session).list_drifted()
        finally:
            await async_engine.dispose()

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


@celery.task(
    name="backend.app.tasks.ml.tag_eval_run",
    bind=True,
    # The head-vs-centroid eval (#1130) loads embeddings + fits sklearn heads
    # for several concepts — minutes, not seconds. Runs on the ml queue because
    # only that worker has numpy/scikit-learn.
    soft_time_limit=1800, time_limit=2100,
)
def tag_eval_run(self, run_id: int) -> str:
    """Compute the eval report into the persisted TagEvalRun row so it survives
    navigation (the admin card rehydrates from the row, not transient state)."""
    from datetime import UTC, datetime

    from ..models import TagEvalRun
    from ..services.ml.tag_eval import run_eval

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        run = session.get(TagEvalRun, run_id)
        if run is None:
            return "missing"
        run.last_progress_at = datetime.now(UTC)
        session.commit()
        try:
            report = run_eval(session, run.params)
        except SoftTimeLimitExceeded:
            run.status = "error"
            run.error = "timed out"
            run.finished_at = datetime.now(UTC)
            session.commit()
            raise
        except Exception as exc:
            log.exception("tag_eval_run %d failed", run_id)
            run.status = "error"
            run.error = str(exc)
            run.finished_at = datetime.now(UTC)
            session.commit()
            return "error"
        run.report = report
        run.status = "ready"
        run.finished_at = datetime.now(UTC)
        session.commit()
        return "ready"


@celery.task(
    name="backend.app.tasks.ml.train_heads",
    bind=True,
    # Trains a logistic-regression head per eligible concept over stored SigLIP
    # embeddings — minutes for a full library. Runs on the ml queue (only that
    # worker has scikit-learn). Commits per head so a kill leaves progress.
    soft_time_limit=3600, time_limit=3900,
)
def train_heads(self, run_id: int) -> str:
    """(Re)train all eligible concept heads into tag_head, tracked by the
    HeadTrainingRun row so the admin card shows live + historical status."""
    from datetime import UTC, datetime

    from ..models import HeadTrainingRun
    from ..services.ml.heads import train_all_heads

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        run = session.get(HeadTrainingRun, run_id)
        if run is None:
            return "missing"
        run.last_progress_at = datetime.now(UTC)
        session.commit()
        try:
            result = train_all_heads(session, run.params, run)
        except SoftTimeLimitExceeded:
            run.status = "error"
            run.error = "timed out"
            run.finished_at = datetime.now(UTC)
            session.commit()
            raise
        except Exception as exc:
            log.exception("train_heads %d failed", run_id)
            run.status = "error"
            run.error = str(exc)
            run.finished_at = datetime.now(UTC)
            session.commit()
            return "error"
        run.n_trained = result["n_trained"]
        run.n_skipped = result["n_skipped"]
        run.status = "ready"
        run.finished_at = datetime.now(UTC)
        session.commit()
        return "ready"


@celery.task(name="backend.app.tasks.ml.scheduled_train_heads")
def scheduled_train_heads() -> str:
    """Nightly passive retrain (#114): fold the day's accepts/rejects + any
    newly-eligible concepts into the heads without the operator clicking. Skips
    if a run is already in flight (one at a time). Creates + COMMITS the run row
    before dispatching so the ml-queue worker can always find it."""
    from datetime import UTC, datetime

    from sqlalchemy import select as sa_select

    from ..models import HeadTrainingRun

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        running = session.execute(
            sa_select(HeadTrainingRun.id).where(HeadTrainingRun.status == "running")
        ).scalar_one_or_none()
        if running is not None:
            return "already running"
        run = HeadTrainingRun(
            params={"source": "scheduled"}, status="running",
            last_progress_at=datetime.now(UTC),
        )
        session.add(run)
        session.commit()
        run_id = run.id
    train_heads.delay(run_id)
    return "dispatched"


@celery.task(
    name="backend.app.tasks.ml.apply_head_tags",
    bind=True,
    # Scores the whole library against the graduated heads and applies their
    # tags (or, dry_run, just counts). Streams embeddings in chunks; numpy only,
    # but ml queue keeps it off the API workers. Commits per chunk.
    soft_time_limit=3600, time_limit=3900,
)
def apply_head_tags(self, run_id: int) -> str:
    """Run an earned-auto-apply sweep into the persisted HeadAutoApplyRun row."""
    from datetime import UTC, datetime

    from ..models import HeadAutoApplyRun
    from ..services.ml.heads import auto_apply_sweep

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        run = session.get(HeadAutoApplyRun, run_id)
        if run is None:
            return "missing"
        run.last_progress_at = datetime.now(UTC)
        session.commit()
        try:
            result = auto_apply_sweep(session, run, run.dry_run)
        except SoftTimeLimitExceeded:
            run.status = "error"
            run.error = "timed out"
            run.finished_at = datetime.now(UTC)
            session.commit()
            raise
        except Exception as exc:
            log.exception("apply_head_tags %d failed", run_id)
            run.status = "error"
            run.error = str(exc)
            run.finished_at = datetime.now(UTC)
            session.commit()
            return "error"
        run.n_applied = result["n_applied"]
        run.report = {"concepts": result["concepts"]}
        run.status = "ready"
        run.finished_at = datetime.now(UTC)
        session.commit()
        return "ready"


@celery.task(name="backend.app.tasks.ml.scheduled_apply_head_tags")
def scheduled_apply_head_tags() -> str:
    """Daily passive auto-apply sweep (#114) — only when the master switch is on.
    Skips if a sweep is already in flight. Creates + COMMITS the run before
    dispatching so the worker always finds it."""
    from datetime import UTC, datetime

    from sqlalchemy import select as sa_select

    from ..models import HeadAutoApplyRun, MLSettings

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        enabled = session.execute(
            sa_select(MLSettings.head_auto_apply_enabled).where(MLSettings.id == 1)
        ).scalar_one_or_none()
        if not enabled:
            return "disabled"
        running = session.execute(
            sa_select(HeadAutoApplyRun.id).where(HeadAutoApplyRun.status == "running")
        ).scalar_one_or_none()
        if running is not None:
            return "already running"
        run = HeadAutoApplyRun(
            dry_run=False, params={"dry_run": False, "source": "scheduled"},
            status="running", last_progress_at=datetime.now(UTC),
        )
        session.add(run)
        session.commit()
        run_id = run.id
    apply_head_tags.delay(run_id)
    return "dispatched"


@celery.task(name="backend.app.tasks.ml.enqueue_gpu_backfill")
def enqueue_gpu_backfill(task_name: str) -> int:
    """Enqueue a gpu_job for every image that doesn't already have one for
    `task_name` (one INSERT…SELECT, so it scales to a full library). The desktop
    agent drains the queue over HTTP. Returns the number enqueued."""
    from sqlalchemy import exists, insert, literal
    from sqlalchemy import select as sa_select

    from ..models import GpuJob, ImageRecord

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        already = exists().where(
            GpuJob.image_record_id == ImageRecord.id,
            GpuJob.task == task_name,
            GpuJob.status.in_(["pending", "leased", "done"]),
        )
        sel = sa_select(
            ImageRecord.id, literal(task_name), literal("pending")
        ).where(~already)
        # RETURNING + count: result.rowcount is unreliable for INSERT…SELECT.
        rows = session.execute(
            insert(GpuJob)
            .from_select(["image_record_id", "task", "status"], sel)
            .returning(GpuJob.id)
        ).fetchall()
        session.commit()
        return len(rows)


@celery.task(name="backend.app.tasks.ml.recover_orphaned_gpu_jobs")
def recover_orphaned_gpu_jobs() -> int:
    """Reset expired GPU-job leases back to pending — recovers work orphaned by an
    agent that died mid-job (no graceful release). Short beat cadence so orphans
    get picked back up quickly + the queue counts read honestly. Returns the
    number recovered."""
    from datetime import UTC, datetime

    from sqlalchemy import update

    from ..models import GpuJob

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        now = datetime.now(UTC)
        res = session.execute(
            update(GpuJob)
            .where(GpuJob.status == "leased", GpuJob.lease_expires_at < now)
            .values(
                status="pending", lease_token=None, leased_at=None,
                lease_expires_at=None, updated_at=now,
            )
        )
        session.commit()
        return res.rowcount or 0


@celery.task(
    name="backend.app.tasks.ml.scheduled_ccip_auto_apply",
    soft_time_limit=1800, time_limit=2100,
)
def scheduled_ccip_auto_apply() -> str:
    """Auto-tag confident CCIP character matches (source='ccip_auto') so identity
    tags keep flowing without a button. No-op unless ccip_auto_apply_enabled.
    References come only from single-character images (unambiguous); a tag is
    applied where any figure's best cosine to a character's prototypes clears
    ccip_auto_apply_threshold and it isn't already applied/rejected. Reversible."""
    import numpy as np
    from sqlalchemy import func
    from sqlalchemy import select as sa_select
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    from ..models import ImageRegion, MLSettings, Tag, TagKind, TagSuggestionRejection
    from ..models.tag import image_tag

    fig = ("face", "figure")

    def _l2(m):
        n = np.linalg.norm(m, axis=1, keepdims=True)
        n[n == 0] = 1.0
        return m / n

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        s = session.get(MLSettings, 1)
        if s is None or not s.ccip_auto_apply_enabled:
            return "disabled"
        thr = float(s.ccip_auto_apply_threshold)

        single = (
            sa_select(image_tag.c.image_record_id)
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.kind == TagKind.character)
            .group_by(image_tag.c.image_record_id)
            .having(func.count() == 1)
        )
        ref_rows = session.execute(
            sa_select(image_tag.c.tag_id, ImageRegion.ccip_embedding)
            .select_from(ImageRegion)
            .join(
                image_tag,
                image_tag.c.image_record_id == ImageRegion.image_record_id,
            )
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.kind == TagKind.character)
            .where(ImageRegion.kind.in_(fig))
            .where(ImageRegion.ccip_embedding.is_not(None))
            .where(ImageRegion.image_record_id.in_(single))
        ).all()
        if not ref_rows:
            return "no-references"

        by_char: dict[int, list] = {}
        for tid, vec in ref_rows:
            by_char.setdefault(tid, []).append(vec)
        ref_tags = list(by_char)
        mats = [_l2(np.asarray(by_char[t], dtype=np.float32)) for t in ref_tags]
        allref = np.vstack(mats)                                  # (total, 768)
        seg = np.cumsum([0] + [len(m) for m in mats])[:-1]        # per-char start

        # Per character: images that already carry OR rejected the tag — skip.
        skip = {t: set() for t in ref_tags}
        for t in ref_tags:
            for (iid,) in session.execute(
                sa_select(image_tag.c.image_record_id).where(
                    image_tag.c.tag_id == t
                )
            ):
                skip[t].add(iid)
            for (iid,) in session.execute(
                sa_select(TagSuggestionRejection.image_record_id).where(
                    TagSuggestionRejection.tag_id == t
                )
            ):
                skip[t].add(iid)

        img_ids = list(session.execute(
            sa_select(ImageRegion.image_record_id)
            .where(ImageRegion.kind.in_(fig), ImageRegion.ccip_embedding.is_not(None))
            .distinct()
        ).scalars())

        applied = 0
        chunk_n = 500
        for start in range(0, len(img_ids), chunk_n):
            chunk = img_ids[start:start + chunk_n]
            rows = session.execute(
                sa_select(ImageRegion.image_record_id, ImageRegion.ccip_embedding)
                .where(
                    ImageRegion.image_record_id.in_(chunk),
                    ImageRegion.kind.in_(fig),
                    ImageRegion.ccip_embedding.is_not(None),
                )
            ).all()
            by_img: dict[int, list] = {}
            for iid, vec in rows:
                by_img.setdefault(iid, []).append(vec)
            for iid, vecs in by_img.items():
                q = _l2(np.asarray(vecs, dtype=np.float32))      # (nq, 768)
                colmax = (q @ allref.T).max(axis=0)              # (total,)
                charmax = np.maximum.reduceat(colmax, seg)       # (n_chars,)
                for ci in np.where(charmax >= thr)[0]:
                    t = ref_tags[int(ci)]
                    if iid in skip[t]:
                        continue
                    skip[t].add(iid)
                    session.execute(
                        pg_insert(image_tag)
                        .values(
                            image_record_id=iid, tag_id=t, source="ccip_auto",
                        )
                        .on_conflict_do_nothing()
                    )
                    applied += 1
            session.commit()
        return f"applied={applied}"