imagerepo/app/tasks/maintenance.py

"""Maintenance-queue Celery tasks.

Currently holds the blocklist cleanup task that retroactively removes
blocklisted tag names from the library. The earlier sync_character_fandoms
task was removed on 2026-04-21 as part of the bare-name refactor.
"""
import logging
import re
from datetime import datetime, timezone

from celery import shared_task
from sqlalchemy import text

from app import db
from app.models import ImageRecord, ImportTask, SuggestionFeedback, Tag
from app.services.integrity import verify_path

log = logging.getLogger(__name__)

_FANDOM_SUFFIX_RE = re.compile(r'^(.+?) \(([^()]+)\)$')


@shared_task(
    name='app.tasks.maintenance.sync_character_fandoms_to_images',
    soft_time_limit=300,
    time_limit=600,
)
def _heal_malformed_character_names() -> dict:
    """Find characters whose name still contains a '(Fandom)' suffix and
    whose fandom_id is NULL (typically created by a pre-fix suggestion
    accept that didn't parse the WD14 suffix). Parse the suffix, ensure
    the fandom tag exists, and either promote the malformed row to
    (name=bare, fandom_id=<fandom>) or merge it into an existing canonical
    character with the same (bare_name, fandom_id) pair.

    Same auto-merge semantics as migration j26042101 Phase 2. Uses raw
    SQL throughout to avoid ORM autoflush ordering quirks between the
    fandom INSERT and the subsequent tag UPDATE inside a single worker
    task session.
    """
    rows = db.session.execute(text("""
        SELECT id, name FROM tag
        WHERE kind = 'character' AND fandom_id IS NULL AND name LIKE '% (%)%'
    """)).fetchall()

    healed = 0
    merged = 0
    skipped = 0

    for row in rows:
        tag_id = row[0]
        old_name = row[1]
        m = _FANDOM_SUFFIX_RE.match(old_name)
        if not m:
            skipped += 1
            continue
        bare_name = m.group(1).strip()
        fandom_name = m.group(2).strip()
        if not bare_name or not fandom_name:
            skipped += 1
            continue

        try:
            # Find-or-create the fandom via raw SQL, commit so the FK from
            # the subsequent character UPDATE can resolve.
            fandom_row = db.session.execute(
                text("SELECT id FROM tag WHERE kind='fandom' AND name=:n"),
                {'n': fandom_name},
            ).fetchone()
            if fandom_row is None:
                fandom_id = db.session.execute(
                    text("INSERT INTO tag (kind, name) VALUES ('fandom', :n) RETURNING id"),
                    {'n': fandom_name},
                ).scalar_one()
            else:
                fandom_id = fandom_row[0]
            db.session.commit()

            # Collision check against the post-refactor partial index.
            canonical_row = db.session.execute(
                text("""
                    SELECT id FROM tag
                    WHERE kind='character' AND name=:n AND fandom_id=:f AND id<>:me
                """),
                {'n': bare_name, 'f': fandom_id, 'me': tag_id},
            ).fetchone()

            if canonical_row is not None:
                canonical_id = canonical_row[0]
                # Reassign image_tags onto the canonical, delete malformed's
                # own image_tags rows.
                db.session.execute(text("""
                    INSERT INTO image_tags (image_id, tag_id)
                    SELECT DISTINCT image_id, :to_id FROM image_tags WHERE tag_id = :from_id
                    ON CONFLICT DO NOTHING
                """), {'from_id': tag_id, 'to_id': canonical_id})
                db.session.execute(
                    text("DELETE FROM image_tags WHERE tag_id = :tid"),
                    {'tid': tag_id},
                )
                _cascade_ref_embedding_rename(old_name, bare_name)
                db.session.execute(
                    text("DELETE FROM tag WHERE id = :tid"),
                    {'tid': tag_id},
                )
                merged += 1
            else:
                _cascade_ref_embedding_rename(old_name, bare_name)
                db.session.execute(
                    text("UPDATE tag SET name=:n, fandom_id=:f WHERE id=:tid"),
                    {'n': bare_name, 'f': fandom_id, 'tid': tag_id},
                )
                healed += 1
            db.session.commit()
        except Exception as e:
            db.session.rollback()
            log.warning(
                "heal_malformed_character_names: tag %s (%r) failed: %s",
                tag_id, old_name, e,
            )
            skipped += 1

    return {'malformed_scanned': len(rows), 'healed': healed, 'merged': merged, 'skipped': skipped}


def _cascade_ref_embedding_rename(old_name: str, new_name: str) -> None:
    """Rename tag_reference_embedding rows from old_name to new_name where
    new_name doesn't already have a row for the same model_version; delete
    leftover old_name rows after. Matches the pattern used by migration
    j26042101."""
    if old_name == new_name:
        return
    db.session.execute(text("""
        UPDATE tag_reference_embedding SET tag_name = :new_name
        WHERE tag_name = :old_name
          AND NOT EXISTS (
              SELECT 1 FROM tag_reference_embedding w
              WHERE w.tag_name = :new_name AND w.model_version = tag_reference_embedding.model_version
          )
    """), {'old_name': old_name, 'new_name': new_name})
    db.session.execute(
        text("DELETE FROM tag_reference_embedding WHERE tag_name = :old_name"),
        {'old_name': old_name},
    )


def sync_character_fandoms_to_images() -> dict:
    """Idempotent two-pass maintenance:

    Pass A — heal characters whose 'name' still embeds a '(Fandom)' suffix
             and whose fandom_id is NULL. These come from pre-fix suggestion
             accepts that didn't parse the WD14 output. Splits the suffix,
             ensures the fandom tag exists, promotes or merges the row.

    Pass B — for every character with a non-null fandom_id, attach the
             fandom tag to every image that has the character attached.
             Backfills the gap left by migration j26042101 (which populated
             tag.fandom_id from old suffixes but didn't walk image_tags).

    Additive for image_tags; destructive only for malformed character rows
    that merge into a canonical row. Safe to re-run — both passes are
    idempotent.
    """
    pass_a = _heal_malformed_character_names()

    chars = (
        Tag.query
        .filter_by(kind='character')
        .filter(Tag.fandom_id.isnot(None))
        .all()
    )
    total_links_added = 0
    characters_processed = 0
    failures = 0
    for char in chars:
        try:
            result = db.session.execute(
                text("""
                    INSERT INTO image_tags (image_id, tag_id)
                    SELECT it.image_id, :fandom_id
                    FROM image_tags it
                    WHERE it.tag_id = :char_id
                    ON CONFLICT DO NOTHING
                """),
                {'char_id': char.id, 'fandom_id': char.fandom_id},
            )
            total_links_added += result.rowcount or 0
            db.session.commit()
            characters_processed += 1
        except Exception as e:
            db.session.rollback()
            failures += 1
            log.warning(
                "sync_character_fandoms_to_images: char tag %s (%s) failed: %s",
                char.id, char.name, e,
            )
    log.info(
        "sync_character_fandoms_to_images: heal=%s, char_attach=%d processed, %d links added, %d failures",
        pass_a, characters_processed, total_links_added, failures,
    )
    return {
        'heal': pass_a,
        'characters_scanned': len(chars),
        'characters_processed': characters_processed,
        'links_added': total_links_added,
        'characters_failed': failures,
    }


@shared_task(
    name='app.tasks.maintenance.verify_media_integrity',
    soft_time_limit=60,
    time_limit=120,
)
def verify_media_integrity(image_id: int) -> dict:
    """Run structural verification on one image and persist the result.

    Idempotent — safe to re-enqueue freely; each call re-checks the file
    on disk and writes a fresh `integrity_checked_at` timestamp. Used both
    by the post-import hook (so newly-imported rows land verified) and by
    the sweep task (which re-verifies the existing library).
    """
    image = ImageRecord.query.get(image_id)
    if image is None:
        return {'image_id': image_id, 'status': 'no_record'}

    status, detail = verify_path(image.filepath)
    image.integrity_status = status
    image.integrity_checked_at = datetime.now(timezone.utc)
    db.session.commit()

    if status != 'ok':
        log.warning(
            "verify_media_integrity: image_id=%s path=%r → %s (%s)",
            image_id, image.filepath, status, detail,
        )
    return {'image_id': image_id, 'status': status, 'detail': detail}


@shared_task(
    name='app.tasks.maintenance.verify_unverified_images',
    soft_time_limit=300,
    time_limit=600,
)
def verify_unverified_images(only_unknown: bool = True) -> dict:
    """Sweep every ImageRecord and enqueue per-image verify tasks.

    Skips paths that are currently the target of an in-flight import — same
    contract `deep_scan_directory` uses, so a half-written file mid-import
    doesn't get false-flagged. By default only revisits rows whose status is
    still 'unknown' (post-migration default + brand-new rows that haven't
    been picked up by the import-time hook yet); pass only_unknown=False to
    force a full re-verify across the library.
    """
    active_paths = {
        row[0] for row in db.session.execute(text("""
            SELECT source_path FROM import_task
            WHERE status IN ('pending', 'queued', 'processing')
        """)).fetchall()
    }

    q = ImageRecord.query.filter(ImageRecord.filepath.isnot(None))
    if only_unknown:
        q = q.filter(ImageRecord.integrity_status == 'unknown')

    enqueued = 0
    skipped_active = 0
    last_id = 0
    BATCH = 500
    while True:
        rows = (
            q.filter(ImageRecord.id > last_id)
            .order_by(ImageRecord.id.asc())
            .limit(BATCH)
            .all()
        )
        if not rows:
            break
        for r in rows:
            if r.filepath in active_paths:
                skipped_active += 1
                continue
            verify_media_integrity.delay(r.id)
            enqueued += 1
        last_id = rows[-1].id

    log.info(
        "verify_unverified_images: enqueued=%d skipped_active=%d only_unknown=%s",
        enqueued, skipped_active, only_unknown,
    )
    return {
        'enqueued': enqueued,
        'skipped_active': skipped_active,
        'only_unknown': only_unknown,
    }


@shared_task(
    name='app.tasks.maintenance.sweep_blocklisted_tag_from_images',
    soft_time_limit=60,
    time_limit=120,
)
def sweep_blocklisted_tag_from_images(name: str) -> dict:
    """Remove a blocklisted tag name from every image that has it attached,
    then delete the Tag row itself.

    Scope is limited to kind='user' tags — that's the kind WD14's general
    category gets materialized as when accepted, which is the vast majority
    of blocklist hits. Character / fandom / artist / series / post / archive
    tags sharing the same name are left alone: those are deliberate, curated
    entities, and removing them silently because of a blocklist text match
    would be destructive.

    Returns a summary dict so the Celery result is introspectable.
    """
    tag = Tag.query.filter_by(kind='user', name=name).first()
    if tag is None:
        log.info("sweep_blocklisted_tag_from_images: no kind='user' tag named %r; nothing to do", name)
        return {'name': name, 'tag_found': False, 'image_tags_deleted': 0, 'tag_deleted': False}

    result = db.session.execute(
        text("DELETE FROM image_tags WHERE tag_id = :tid"),
        {'tid': tag.id},
    )
    rowcount = result.rowcount or 0
    db.session.delete(tag)
    db.session.commit()

    log.info(
        "sweep_blocklisted_tag_from_images: removed %r (tag_id=%d) from %d images",
        name, tag.id, rowcount,
    )
    return {
        'name': name,
        'tag_found': True,
        'image_tags_deleted': rowcount,
        'tag_deleted': True,
    }


@shared_task(
    name='app.tasks.maintenance.apply_auto_accept_predictions',
    soft_time_limit=1800,
    time_limit=2700,
)
def apply_auto_accept_predictions(batch_size: int = 100) -> dict:
    """Walk every image and apply general-category WD14 predictions whose
    confidence is at or above the auto_accept_general_threshold config.

    Mirrors the side effects of the modal's per-image auto-accept flow
    (`get_image_suggestions`): for each candidate, find-or-create a
    kind='user' Tag, attach it to the image, and log a SuggestionFeedback
    row with decision='accepted'. Idempotent — re-running just skips
    images whose existing tags + per-image rejections already cover the
    candidates.

    Designed for one-shot library-wide application after a backfill
    completes, so the user doesn't have to open every modal to land
    high-confidence WD14 tags. Skips integrity-flagged rows (the suggestion
    service already returns empty for those) and amortizes the config
    fetch + Tag-table scan across the loop.
    """
    # Lazy-import the suggestion service so the maintenance worker doesn't
    # pay its overhead on unrelated tasks.
    from app.services.tag_suggestions import (
        _DEFAULTS, _config, _existing_tag_names, get_suggestions,
    )
    from app.ml.wd14 import MODEL_VERSION as WD14_VER

    cfg = _config()
    existing_all = _existing_tag_names()

    # Fast-path: if no WD14 prediction at/above the threshold exists for
    # an image-tag combo that isn't already attached and isn't user-rejected,
    # there's nothing to do. Daily-scheduled runs will hit this branch on
    # most days once the library has settled, so the full walk only runs
    # after fresh predictions land.
    try:
        threshold = float(cfg.get('auto_accept_general_threshold',
                                  _DEFAULTS['auto_accept_general_threshold']))
    except (TypeError, ValueError):
        threshold = float(_DEFAULTS['auto_accept_general_threshold'])

    pending_exists = db.session.execute(text("""
        SELECT 1
        FROM image_tag_prediction p
        WHERE p.confidence >= :thr
          AND p.tag_category = 'general'
          AND p.model_version = :wd14_ver
          AND NOT EXISTS (
            SELECT 1 FROM image_tags it
            JOIN tag t ON t.id = it.tag_id
            WHERE it.image_id = p.image_id
              AND t.name = p.tag_name
              AND t.kind = 'user'
          )
          AND NOT EXISTS (
            SELECT 1 FROM suggestion_feedback sf
            WHERE sf.image_id = p.image_id
              AND sf.tag_name = p.tag_name
              AND sf.decision = 'rejected'
          )
        LIMIT 1
    """), {'thr': threshold, 'wd14_ver': WD14_VER}).first()

    if pending_exists is None:
        log.info(
            "apply_auto_accept_predictions: no candidates above threshold=%.3f; skipping walk",
            threshold,
        )
        return {
            'scanned': 0,
            'images_with_applies': 0,
            'tags_applied': 0,
            'skipped_no_candidates': True,
        }

    scanned = 0
    images_with_applies = 0
    tags_applied = 0
    last_id = 0

    while True:
        rows = (
            ImageRecord.query
            .filter(ImageRecord.id > last_id)
            .filter(ImageRecord.integrity_status.in_(('ok', 'unknown')))
            .order_by(ImageRecord.id.asc())
            .limit(batch_size)
            .all()
        )
        if not rows:
            break

        for img in rows:
            scanned += 1
            grouped = get_suggestions(img.id, cfg=cfg, existing_all=existing_all)
            candidates = grouped.get('auto_accept_candidates') or []
            if not candidates:
                continue

            applied_for_image = 0
            for cand in candidates:
                name = cand['name']
                tag = Tag.query.filter_by(kind='user', name=name).first()
                if tag is None:
                    tag = Tag(kind='user', name=name)
                    db.session.add(tag)
                    db.session.flush()
                if tag not in img.tags:
                    img.tags.append(tag)
                    applied_for_image += 1
                db.session.add(SuggestionFeedback(
                    image_id=img.id,
                    tag_name=name,
                    suggestion_source=cand.get('source', 'wd14'),
                    confidence=cand.get('confidence', 0.0),
                    decision='accepted',
                ))

            if applied_for_image:
                images_with_applies += 1
                tags_applied += applied_for_image

        # Commit once per batch — keeps transactions short and lets the
        # next batch see freshly-created Tag rows.
        db.session.commit()
        last_id = rows[-1].id
        log.info(
            "apply_auto_accept_predictions: progress scanned=%d "
            "images_with_applies=%d tags_applied=%d last_id=%d",
            scanned, images_with_applies, tags_applied, last_id,
        )

    log.info(
        "apply_auto_accept_predictions: done scanned=%d "
        "images_with_applies=%d tags_applied=%d",
        scanned, images_with_applies, tags_applied,
    )
    return {
        'scanned': scanned,
        'images_with_applies': images_with_applies,
        'tags_applied': tags_applied,
    }