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 celery import shared_task
from sqlalchemy import text

from app import db
from app.models import Tag

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.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,
    }