"""ML Celery tasks: per-image embedding, backfill discovery, head training,
model self-heal.

All run on the ml-worker (queue 'ml'). 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 select
from sqlalchemy.exc import DBAPIError, OperationalError

from ..celery_app import celery
from ..models import 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:
    """Compute + store one image's SigLIP embedding.

    Video (#747): sample frames at a fixed cadence (ml_settings
    video_frame_interval_seconds, capped at video_max_frames) and mean-pool the
    per-frame SigLIP embeddings. On no-frames returns status='no_frames' (not an
    error). (Camie tagging was retired #1189 — heads + CCIP are the tag source.)
    """
    import time

    from ..services.ml.embedder import get_embedder

    # 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"
            embedder = get_embedder(settings.embedder_model_name)

            if is_vid:
                # Layer-3 isolation: ffprobe (a separate process) validates
                # the container before we burn 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.
                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"
                frames = _sample_video_frames(
                    src,
                    interval=settings.video_frame_interval_seconds,
                    max_frames=settings.video_max_frames,
                )
                if not frames:
                    return {"status": "no_frames", "image_id": image_id}
                phase = "video_embed"
                import numpy as np

                # Mean-pool the per-frame SigLIP embeddings into one vector.
                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 = "embed"
                t0 = time.monotonic()
                embedding = embedder.infer(src)
                log.info(
                    "tag_and_embed embedded in %.1fs: %s",
                    time.monotonic() - t0, ctx,
                )

            phase = "persist"
            record.siglip_embedding = embedding.tolist()
            record.siglip_model_version = settings.embedder_model_version
            session.add(record)
            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: %s", _elapsed(), ctx)
    return {"status": "ok", "image_id": image_id}


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


@celery.task(name="backend.app.tasks.ml.backfill", bind=True)
def backfill(self) -> int:
    """Enqueue tag_and_embed (embed-only) for images with no SigLIP embedding.
    Keyset pagination by id ASC (restart-safe).

    NB: a siglip MODEL-VERSION mismatch (an operator model swap, #1190) is NOT
    re-embedded here — the CPU ml-worker can't churn the library at 384/512px;
    the GPU agent owns version re-embeds via the 'embed' job.
    """
    SessionLocal = _sync_session_factory()
    enqueued = 0
    last_id = 0
    with SessionLocal() as session:
        while True:
            rows = session.execute(
                select(ImageRecord.id)
                .where(ImageRecord.id > last_id)
                .where(ImageRecord.siglip_embedding.is_(None))
                .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.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 still needs `task_name` (one
    INSERT…SELECT, so it scales to a full library). The desktop agent drains the
    queue over HTTP. Returns the number enqueued.

    'siglip' gates on the RESULT (no concept region yet) rather than on a prior
    job, so it picks up the back-catalogue of images that were CCIP-embedded
    before concept crops existed — without re-touching their figure/CCIP regions.

    An ERRORED job is a tombstone for its (image, task): no variant re-enqueues
    it. Retry is deliberate-only (/retry_errors), which also means an errored
    back-catalogue needs one "Retry errored jobs" press after a model swap.
    Before the tombstone rule, this loop re-minted a fresh doomed job for every
    permanently-bad file each run — ~24 duplicate error rows/day per file (the
    2026-07-02 "unprocessable" flood)."""
    from sqlalchemy import exists, insert, literal, or_
    from sqlalchemy import select as sa_select

    from ..models import GpuJob, ImageRecord, ImageRegion, MLSettings
    from ..services.ml.gpu_jobs import error_dedupe_statements

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        # Prune stale tombstones first (loop-era duplicates + rows made moot by
        # a later success), so 'error' reads as one row per distinct failing
        # file and the skip-guards below see a clean picture.
        pruned = sum(
            session.execute(s).rowcount or 0 for s in error_dedupe_statements()
        )
        if pruned:
            log.info("gpu backfill: pruned %d stale/duplicate error rows", pruned)
        cur_version = session.execute(
            select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
        ).scalar_one()
        if task_name == "embed":
            # Whole-image GPU re-embed (#1190): images with no embedding, or one
            # stamped under a DIFFERENT model version (an operator model swap).
            stale = or_(
                ImageRecord.siglip_embedding.is_(None),
                ImageRecord.siglip_model_version.is_(None),
                ImageRecord.siglip_model_version != cur_version,
            )
            # 'error' blocks too — tombstone rule, see docstring.
            blocked = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == "embed",
                GpuJob.status.in_(["pending", "leased", "error"]),
            )
            sel = sa_select(
                ImageRecord.id, literal("embed"), literal("pending")
            ).where(stale).where(~blocked)
        elif task_name == "siglip":
            # Concept-crop re-embed: enqueue when there's no concept region AT THE
            # CURRENT model version — so a model swap re-triggers crops too, not
            # only the never-embedded back-catalogue.
            has_current_concept = exists().where(
                ImageRegion.image_record_id == ImageRecord.id,
                ImageRegion.kind == "concept",
                ImageRegion.embedding_version == cur_version,
            )
            # 'error' blocks too — tombstone rule, see docstring.
            blocked = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == "siglip",
                GpuJob.status.in_(["pending", "leased", "error"]),
            )
            sel = sa_select(
                ImageRecord.id, literal("siglip"), literal("pending")
            ).where(~has_current_concept).where(~blocked)
        else:
            # ANY prior row blocks — including 'error' (tombstone rule, see
            # docstring): pre-fix this branch ran HOURLY and was the loop.
            already = exists().where(
                GpuJob.image_record_id == ImageRecord.id,
                GpuJob.task == task_name,
                GpuJob.status.in_(["pending", "leased", "done", "error"]),
            )
            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) — and convert poison-loopers
    (release/expiry cycles that never reach fail()'s attempt cap) to 'error'.
    Statements are shared with GpuJobService.recover_orphaned so the sweep and
    the service can't drift. 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 ..services.ml.gpu_jobs import recover_statements

    SessionLocal = _sync_session_factory()
    with SessionLocal() as session:
        counts = {
            name: session.execute(stmt).rowcount or 0
            for name, stmt in recover_statements(datetime.now(UTC)).items()
        }
        session.commit()
        if counts["poison_expired"] or counts["poison_pending"]:
            log.warning(
                "gpu jobs poisoned -> error: %d crash-loop (expired lease), "
                "%d never-complete (pending)",
                counts["poison_expired"], counts["poison_pending"],
            )
        return counts["recovered"]


@celery.task(name="backend.app.tasks.ml.reprocess_gpu_jobs")
def reprocess_gpu_jobs(task_name: str = "ccip") -> int:
    """Reset every done/error job of `task_name` back to pending so the agent
    re-runs the WHOLE library under the CURRENT pipeline — e.g. after adding crop
    detectors (#1202), re-cropping existing images. Heavy + operator-triggered;
    the back-catalogue won't otherwise re-process (the backfills skip images that
    already have current-version regions). Returns the number reset."""
    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.task == task_name,
                GpuJob.status.in_(["done", "error"]),
            )
            .values(
                status="pending", attempts=0, 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}"