"""Tag CRUD + autocomplete + image-tag association."""

import logging
import time
from collections.abc import Sequence
from dataclasses import dataclass

from sqlalchemy import and_, case, exists, func, select, text, update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import (
    HeadMetric,
    Tag,
    TagHead,
    TagKind,
    TagPositiveConfirmation,
    image_tag,
)
from .db_helpers import get_or_create
from .tag_query import fandom_join_alias, tag_columns

log = logging.getLogger(__name__)


def normalize_tag_name(name: str) -> str:
    """Canonical tag form (#701): collapse whitespace + capitalize the first
    letter of each word, PRESERVING the rest of the word.

    Per-word capitalize (NOT str.title(), which mangles apostrophes:
    don't → Don'T). The word tail is left untouched so acronyms survive
    (DC stays DC, NSFW stays NSFW) — operator-revised 2026-06-06: protecting
    acronyms matters more than folding ALL-CAPS input. This MATCHES
    ml/tag_name._title_word, so a tag suggested by the Camie tagger keeps the
    exact casing the suggestion UI showed when it round-trips through the
    create endpoint on Accept.
    """
    return " ".join(
        w[:1].upper() + w[1:] for w in (name or "").split()
    )


class TagValidationError(ValueError):
    """Raised when tag construction breaks the kind/fandom rules."""


class TagMergeConflict(TagValidationError):
    """Rename hit a same-(kind, fandom) name collision. Carries the
    colliding target and the merge preview the 409 needs."""

    def __init__(
        self,
        message: str,
        *,
        target_id: int,
        target_name: str,
        source_image_count: int,
        will_alias: bool,
    ):
        super().__init__(message)
        self.target_id = target_id
        self.target_name = target_name
        self.source_image_count = source_image_count
        self.will_alias = will_alias


@dataclass(frozen=True)
class TagAutocompleteHit:
    id: int
    name: str
    kind: str
    fandom_id: int | None
    fandom_name: str | None
    image_count: int


@dataclass(frozen=True)
class MergeResult:
    target_id: int
    target_name: str
    target_kind: str
    merged_count: int
    alias_created: bool
    source_deleted: bool


@dataclass(frozen=True)
class MergePreview:
    """Non-mutating projection of what merge(source→target) would do (#8).
    Shares the apply's predicates (rule 93) so the preview can't drift."""
    source_id: int
    source_name: str
    target_id: int
    target_name: str
    compatible: bool                # same kind + fandom — would apply succeed?
    images_moving: int              # source links that move (image lacks target)
    images_already_on_target: int   # source links dropped (image has target)
    source_total: int               # images carrying source
    series_pages: int               # series pages repointed
    will_alias: bool                # source name kept as a protective alias
    sample_thumbnails: list[str]    # a few of the moving images


class TagService:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def find_or_create(
        self, name: str, kind: TagKind, fandom_id: int | None = None
    ) -> Tag:
        """Upserts (name, kind, fandom_id) into the tag table.

        Validates kind/fandom rules before touching the DB:
        - fandom_id may only be set when kind == TagKind.character
        - if fandom_id is set, the referenced tag must exist and have
          kind == TagKind.fandom
        """
        # NOTE: case is NOT normalized here — find_or_create is the shared path
        # the ML tagger / allowlist also use, and Title-Casing the booru
        # vocabulary breaks allowlist matching. Display-casing (Title Case) is
        # applied at the user-entry layer (api/tags create_tag) only (#701).
        name = name.strip()
        if not name:
            raise TagValidationError("Tag name cannot be empty")

        if fandom_id is not None and kind != TagKind.character:
            raise TagValidationError(
                "fandom_id is only valid for character tags"
            )

        if fandom_id is not None:
            fandom = await self.session.get(Tag, fandom_id)
            if fandom is None or fandom.kind != TagKind.fandom:
                raise TagValidationError(
                    f"fandom_id {fandom_id} does not reference a fandom tag"
                )

        # Audit 2026-06-02: race-safe upsert via savepoint +
        # IntegrityError recovery. The partial uniqueness index on
        # (name, kind, COALESCE(fandom_id, -1)) catches concurrent
        # inserts; without the savepoint the outer transaction would
        # poison and the calling request crashes. Mirrors
        # importer._get_or_create.
        # Case-insensitive match (#701): a normalized (Title Case) input must
        # find an existing differently-cased tag instead of forking a duplicate.
        stmt = (
            select(Tag)
            .where(func.lower(Tag.name) == name.lower())
            .where(Tag.kind == kind)
            .where(
                Tag.fandom_id.is_(None) if fandom_id is None else Tag.fandom_id == fandom_id
            )
            .order_by(Tag.id)
            .limit(1)
        )
        async def _create() -> Tag:
            new_tag = Tag(name=name, kind=kind, fandom_id=fandom_id)
            self.session.add(new_tag)
            await self.session.flush()
            return new_tag

        tag, _created = await get_or_create(self.session, stmt, _create)
        return tag

    async def autocomplete(
        self,
        query: str,
        kind: TagKind | None = None,
        limit: int = 20,
    ) -> list[TagAutocompleteHit]:
        """ILIKE search ranked exact > prefix > substring, then by image count.

        Returns at most `limit` hits.
        """
        q = (query or "").strip()
        if not q:
            return []

        match_rank = case(
            (func.lower(Tag.name) == q.lower(), 0),
            (func.lower(Tag.name).startswith(q.lower()), 1),
            else_=2,
        )

        fandom_alias = fandom_join_alias()
        image_count = (
            select(func.count(image_tag.c.image_record_id))
            .where(image_tag.c.tag_id == Tag.id)
            .correlate(Tag.__table__)
            .scalar_subquery()
        )

        stmt = (
            select(
                *tag_columns(fandom_alias),
                image_count.label("image_count"),
            )
            .select_from(Tag.__table__.outerjoin(
                fandom_alias, Tag.fandom_id == fandom_alias.c.id
            ))
            .where(func.lower(Tag.name).like(f"%{q.lower()}%"))
            .order_by(match_rank.asc(), image_count.desc(), Tag.name.asc())
            .limit(limit)
        )
        if kind is not None:
            stmt = stmt.where(Tag.kind == kind)

        rows = (await self.session.execute(stmt)).all()
        return [
            TagAutocompleteHit(
                id=r.id,
                name=r.name,
                kind=r.kind.value if hasattr(r.kind, "value") else r.kind,
                fandom_id=r.fandom_id,
                fandom_name=r.fandom_name,
                image_count=r.image_count or 0,
            )
            for r in rows
        ]

    async def add_to_image(self, image_id: int, tag_id: int, source: str = "manual") -> None:
        """Idempotent: re-adding an existing tag does nothing."""
        # A genuinely-new MANUAL add of a tag that already has a head is an
        # UNDER-FIRE signal — the auto-system should have caught it (#114 obs).
        is_new = source == "manual" and (
            await self.session.execute(
                select(image_tag.c.tag_id).where(
                    and_(
                        image_tag.c.image_record_id == image_id,
                        image_tag.c.tag_id == tag_id,
                    )
                )
            )
        ).first() is None
        stmt = pg_insert(image_tag).values(
            image_record_id=image_id, tag_id=tag_id, source=source
        )
        stmt = stmt.on_conflict_do_nothing(
            index_elements=["image_record_id", "tag_id"]
        )
        await self.session.execute(stmt)
        if is_new:
            await self._note_under_fire(tag_id)

    async def remove_from_image(self, image_id: int, tag_id: int) -> None:
        # Removing an auto-applied (source='head_auto') tag is a MISFIRE — read
        # the source BEFORE deleting, since it's lost with the row (#114 obs).
        src = (
            await self.session.execute(
                select(image_tag.c.source).where(
                    and_(
                        image_tag.c.image_record_id == image_id,
                        image_tag.c.tag_id == tag_id,
                    )
                )
            )
        ).scalar_one_or_none()
        await self.session.execute(
            image_tag.delete().where(
                and_(
                    image_tag.c.image_record_id == image_id,
                    image_tag.c.tag_id == tag_id,
                )
            )
        )
        if src == "head_auto":
            await self._bump_metric(tag_id, "n_misfires")

    async def _note_under_fire(self, tag_id: int) -> None:
        """Count an under-fire only when the tag actually has a head."""
        has_head = (
            await self.session.execute(
                select(TagHead.tag_id).where(TagHead.tag_id == tag_id)
            )
        ).first() is not None
        if has_head:
            await self._bump_metric(tag_id, "n_underfires")

    async def _bump_metric(self, tag_id: int, column: str) -> None:
        """Increment a HeadMetric counter (upsert), keyed by tag so it survives
        head retrain/prune."""
        col = HeadMetric.__table__.c[column]
        await self.session.execute(
            pg_insert(HeadMetric)
            .values(tag_id=tag_id, **{column: 1})
            .on_conflict_do_update(
                index_elements=["tag_id"],
                set_={column: col + 1, "updated_at": func.now()},
            )
        )

    async def list_for_image(self, image_id: int) -> Sequence:
        """Tags on an image, ordered (kind, name). Each row carries the fandom's
        NAME (not just fandom_id) via a self-join on Tag, so the UI can label a
        character chip with its fandom without an N+1 (mirrors the
        autocomplete/directory resolution)."""
        fandom_alias = fandom_join_alias()
        # source drives the auto-applied badge; confirmed = operator affirmed the
        # tag (positive + retraction-shielded, milestone 139).
        confirmed = exists().where(
            TagPositiveConfirmation.image_record_id == image_id,
            TagPositiveConfirmation.tag_id == Tag.id,
        ).label("confirmed")
        stmt = (
            select(*tag_columns(fandom_alias), image_tag.c.source, confirmed)
            .select_from(
                Tag.__table__
                .join(image_tag, image_tag.c.tag_id == Tag.id)
                .outerjoin(fandom_alias, Tag.fandom_id == fandom_alias.c.id)
            )
            .where(image_tag.c.image_record_id == image_id)
            .order_by(Tag.kind.asc(), Tag.name.asc())
        )
        return (await self.session.execute(stmt)).all()

    async def _keep_as_alias(self, tag_id: int) -> bool:
        """A merged-away tag's old name must survive as an alias iff the ML
        pipeline has ever applied it (manual accept or head auto-apply) — so a
        re-application or an alias remap resolves the canonical name. Purely-
        manual, ML-unknown tags are deleted outright (no DB bloat)."""
        is_machine = await self.session.scalar(
            select(
                exists().where(
                    and_(
                        image_tag.c.tag_id == tag_id,
                        image_tag.c.source.in_(
                            ("ml_auto", "ml_accepted", "head_auto", "auto")
                        ),
                    )
                )
            )
        )
        return bool(is_machine)

    async def rename(self, tag_id: int, new_name: str) -> Tag:
        """Rename a tag. Raises TagMergeConflict if the new name collides
        with an existing tag of the same (kind, fandom_id) — the API turns
        that into a 409 merge hint.
        """
        new_name = new_name.strip()
        if not new_name:
            raise TagValidationError("Tag name cannot be empty")
        tag = await self.session.get(Tag, tag_id)
        if tag is None:
            raise TagValidationError(f"Tag {tag_id} not found")
        # The training-hygiene machinery keys on system tags by row; a rename
        # would silently disarm it (milestone #128).
        if tag.is_system:
            raise TagValidationError(
                f"'{tag.name}' is a system tag and cannot be renamed"
            )

        # Case-insensitive clash (#701) — renaming onto a differently-cased tag
        # is still a merge, not a silent fork.
        clash_stmt = (
            select(Tag)
            .where(func.lower(Tag.name) == new_name.lower())
            .where(Tag.kind == tag.kind)
            .where(
                Tag.fandom_id.is_(None)
                if tag.fandom_id is None
                else Tag.fandom_id == tag.fandom_id
            )
            .where(Tag.id != tag_id)
            .order_by(Tag.id)
            .limit(1)
        )
        clash = (await self.session.execute(clash_stmt)).scalar_one_or_none()
        if clash is not None:
            source_image_count = await self.session.scalar(
                select(func.count())
                .select_from(image_tag)
                .where(image_tag.c.tag_id == tag_id)
            )
            will_alias = await self._keep_as_alias(tag_id)
            raise TagMergeConflict(
                f"A {tag.kind} tag named {new_name!r} already exists",
                target_id=clash.id,
                target_name=clash.name,
                source_image_count=int(source_image_count or 0),
                will_alias=will_alias,
            )

        tag.name = new_name
        await self.session.flush()
        return tag

    async def set_fandom(
        self, tag_id: int, fandom_id: int | None, *, merge: bool = False
    ) -> Tag:
        """Set / change / clear a character tag's fandom.

        Raises TagValidationError unless the tag is a character and fandom_id
        (when given) references a fandom tag. If the change would collide with
        an existing character of the same name in the TARGET fandom, raises
        TagMergeConflict (the API turns that into a 409 merge hint) — unless
        merge=True, in which case this tag is merged INTO that existing
        character (a deliberate cross-fandom merge) and the surviving target
        is returned. Passing fandom_id=None clears the fandom.
        """
        tag = await self.session.get(Tag, tag_id)
        if tag is None:
            raise TagValidationError(f"Tag {tag_id} not found")
        if tag.kind != TagKind.character:
            raise TagValidationError("Only character tags can have a fandom")
        if fandom_id is not None:
            fandom = await self.session.get(Tag, fandom_id)
            if fandom is None or fandom.kind != TagKind.fandom:
                raise TagValidationError(
                    f"fandom_id {fandom_id} does not reference a fandom tag"
                )
        if fandom_id == tag.fandom_id:
            return tag

        # Collision: another character with the same name already lives in the
        # target fandom. Mirrors rename's (name, kind, fandom_id) uniqueness.
        clash_stmt = (
            select(Tag)
            .where(Tag.name == tag.name)
            .where(Tag.kind == TagKind.character)
            .where(
                Tag.fandom_id.is_(None)
                if fandom_id is None
                else Tag.fandom_id == fandom_id
            )
            .where(Tag.id != tag_id)
        )
        clash = (await self.session.execute(clash_stmt)).scalar_one_or_none()
        if clash is not None:
            if not merge:
                source_image_count = await self.session.scalar(
                    select(func.count())
                    .select_from(image_tag)
                    .where(image_tag.c.tag_id == tag_id)
                )
                will_alias = await self._keep_as_alias(tag_id)
                raise TagMergeConflict(
                    f"A character named {tag.name!r} already exists in that fandom",
                    target_id=clash.id,
                    target_name=clash.name,
                    source_image_count=int(source_image_count or 0),
                    will_alias=will_alias,
                )
            await self._do_merge(tag, clash)
            return clash

        tag.fandom_id = fandom_id
        await self.session.flush()
        return tag

    async def merge_preview(
        self, source_id: int, target_id: int
    ) -> MergePreview:
        """Read-only dry-run of merge(source→target): the same counts the apply
        would produce, plus a thumbnail sample of the images that move. Raises
        TagValidationError only for missing/self-merge (so the UI can render a
        kind/fandom-mismatch warning rather than a hard error)."""
        if source_id == target_id:
            raise TagValidationError("Cannot merge a tag into itself")
        source = await self.session.get(Tag, source_id)
        target = await self.session.get(Tag, target_id)
        if source is None or target is None:
            raise TagValidationError("Tag not found")

        compatible = source.kind == target.kind and (
            (source.fandom_id or 0) == (target.fandom_id or 0)
        )

        # Mirrors _repoint_image_tags: links whose image already has target are
        # dropped (already_on_target); the rest move.
        target_images = select(image_tag.c.image_record_id).where(
            image_tag.c.tag_id == target_id
        )
        moving = await self.session.scalar(
            select(func.count())
            .select_from(image_tag)
            .where(
                image_tag.c.tag_id == source_id,
                image_tag.c.image_record_id.notin_(target_images),
            )
        )
        already = await self.session.scalar(
            select(func.count())
            .select_from(image_tag)
            .where(
                image_tag.c.tag_id == source_id,
                image_tag.c.image_record_id.in_(target_images),
            )
        )

        from ..models.series_page import SeriesPage

        series_pages = await self.session.scalar(
            select(func.count())
            .select_from(SeriesPage)
            .where(SeriesPage.series_tag_id == source_id)
        )

        will_alias = await self._keep_as_alias(source_id)
        sample = await self._merge_sample_thumbnails(source_id, target_images)

        return MergePreview(
            source_id=source_id,
            source_name=source.name,
            target_id=target_id,
            target_name=target.name,
            compatible=compatible,
            images_moving=moving or 0,
            images_already_on_target=already or 0,
            source_total=(moving or 0) + (already or 0),
            series_pages=series_pages or 0,
            will_alias=will_alias,
            sample_thumbnails=sample,
        )

    async def _merge_sample_thumbnails(
        self, source_id: int, target_images
    ) -> list[str]:
        """Up to 6 thumbnails of the images that WOULD move (carry source, not
        target) — a sanity check that the merge targets the right content."""
        from ..models import ImageRecord
        from .gallery_service import thumbnail_url

        rows = (
            await self.session.execute(
                select(
                    ImageRecord.thumbnail_path,
                    ImageRecord.sha256,
                    ImageRecord.mime,
                )
                .join(image_tag, image_tag.c.image_record_id == ImageRecord.id)
                .where(
                    image_tag.c.tag_id == source_id,
                    image_tag.c.image_record_id.notin_(target_images),
                )
                .limit(6)
            )
        ).all()
        return [thumbnail_url(tp, sha, mime) for tp, sha, mime in rows]

    async def merge(self, source_id: int, target_id: int) -> MergeResult:
        """Transactionally repoint every FK from source→target, optionally
        keep source's name as a tagger alias, delete source. Atomic: any
        exception rolls back via the caller's session."""
        if source_id == target_id:
            raise TagValidationError("Cannot merge a tag into itself")
        source = await self.session.get(Tag, source_id)
        if source is None:
            raise TagValidationError(f"Tag {source_id} not found")
        # Merge deletes the source row — the only tag-delete path — and the
        # training-hygiene machinery keys on system rows (milestone #128).
        # Merging INTO a system tag stays allowed (adopting an operator's old
        # 'wip sketch' tag into the system 'wip' is the intended move).
        if source.is_system:
            raise TagValidationError(
                f"'{source.name}' is a system tag and cannot be merged away"
            )
        target = await self.session.get(Tag, target_id)
        if target is None:
            raise TagValidationError(f"Tag {target_id} not found")
        if source.kind != target.kind or (source.fandom_id or 0) != (
            target.fandom_id or 0
        ):
            raise TagValidationError(
                "Tags must be the same kind and fandom to merge"
            )
        return await self._do_merge(source, target)

    async def _do_merge(self, source: Tag, target: Tag) -> MergeResult:
        """Repoint every FK source→target, optionally keep source's name as a
        tagger alias, delete source. NO kind/fandom validation — callers that
        need it (public merge()) validate first; set_fandom's collision
        resolution calls this directly for a deliberate CROSS-fandom merge."""
        source_id, target_id = source.id, target.id
        keep_as_alias = await self._keep_as_alias(source_id)
        source_name = source.name
        source_kind = source.kind
        target_name = target.name
        target_kind = (
            target.kind.value
            if hasattr(target.kind, "value")
            else str(target.kind)
        )

        merged_count = await self._repoint_image_tags(source_id, target_id)
        await self._repoint_rejections(source_id, target_id)
        await self._repoint_aliases(source_id, target_id)
        await self._repoint_fandom_children(
            source_id, target_id, source_kind
        )
        await self._repoint_series_pages(source_id, target_id)

        alias_created = False
        if keep_as_alias:
            alias_created = await self._create_protective_aliases(
                source_name, source_kind, target_id
            )

        await self.session.delete(source)
        await self.session.flush()

        return MergeResult(
            target_id=target_id,
            target_name=target_name,
            target_kind=target_kind,
            merged_count=merged_count,
            alias_created=alias_created,
            source_deleted=True,
        )

    async def _repoint_image_tags(self, src: int, tgt: int) -> int:
        """Drop source links whose image already links target, then move
        the rest. Returns the number actually moved."""
        await self.session.execute(
            image_tag.delete().where(
                and_(
                    image_tag.c.tag_id == src,
                    image_tag.c.image_record_id.in_(
                        select(image_tag.c.image_record_id).where(
                            image_tag.c.tag_id == tgt
                        )
                    ),
                )
            )
        )
        res = await self.session.execute(
            image_tag.update()
            .where(image_tag.c.tag_id == src)
            .values(tag_id=tgt)
        )
        return res.rowcount or 0

    async def _repoint_rejections(self, src: int, tgt: int) -> None:
        from ..models.tag_suggestion_rejection import TagSuggestionRejection

        await self.session.execute(
            text(
                "DELETE FROM tag_suggestion_rejection r "
                "WHERE r.tag_id = :src AND EXISTS ("
                "  SELECT 1 FROM tag_suggestion_rejection r2 "
                "  WHERE r2.tag_id = :tgt "
                "  AND r2.image_record_id = r.image_record_id)"
            ),
            {"src": src, "tgt": tgt},
        )
        await self.session.execute(
            update(TagSuggestionRejection)
            .where(TagSuggestionRejection.tag_id == src)
            .values(tag_id=tgt)
        )

    async def _repoint_aliases(self, src: int, tgt: int) -> None:
        from ..models.tag_alias import TagAlias

        await self.session.execute(
            update(TagAlias)
            .where(TagAlias.canonical_tag_id == src)
            .values(canonical_tag_id=tgt)
        )

    async def _repoint_series_pages(self, src: int, tgt: int) -> None:
        from ..models.series_chapter import SeriesChapter
        from ..models.series_page import SeriesPage

        # Move src's dividers onto tgt. Their anchor_page_id stays valid — the
        # anchored pages move too — and anchor is globally UNIQUE so no collision.
        await self.session.execute(
            update(SeriesChapter)
            .where(SeriesChapter.series_tag_id == src)
            .values(series_tag_id=tgt)
        )
        # Append src's pages after tgt's so the merged series stays ONE clean
        # series-global run (image_id is UNIQUE, so the two page sets are
        # disjoint). Offset src's page_number past tgt's current max.
        tgt_max = (
            await self.session.scalar(
                select(
                    func.coalesce(func.max(SeriesPage.page_number), 0)
                ).where(SeriesPage.series_tag_id == tgt)
            )
        ) or 0
        await self.session.execute(
            update(SeriesPage)
            .where(SeriesPage.series_tag_id == src)
            .values(
                series_tag_id=tgt,
                page_number=SeriesPage.page_number + tgt_max,
            )
        )

    async def _repoint_fandom_children(
        self, src: int, tgt: int, src_kind: TagKind
    ) -> None:
        if src_kind != TagKind.fandom:
            return
        await self.session.execute(
            update(Tag).where(Tag.fandom_id == src).values(fandom_id=tgt)
        )

    async def _image_assoc_counts(self, tag_ids: list[int]) -> dict[int, int]:
        """image_tag row counts keyed by tag_id, for the survivor heuristic
        (the best-connected tag in a collision group survives → fewest
        FK repoints). Tags with zero associations are absent from the map."""
        if not tag_ids:
            return {}
        rows = (
            await self.session.execute(
                select(image_tag.c.tag_id, func.count())
                .where(image_tag.c.tag_id.in_(tag_ids))
                .group_by(image_tag.c.tag_id)
            )
        ).all()
        return {tid: int(n) for tid, n in rows}

    async def _create_protective_aliases(
        self, src_name: str, src_kind: TagKind, tgt: int
    ) -> bool:
        """Alias (src_name, category) -> tgt so future tagger predictions of
        src_name resolve to the merge survivor. Idempotent — never clobbers a
        pre-existing operator alias.

        The category is the tag's kind. The tagger's tag_to_category map is
        one-to-one (a name has exactly ONE category), and a tag's kind is set
        from that category when it's created — so kind already IS the tagger's
        category for this name. This used to SELECT DISTINCT category by scanning
        every image_record's tagger_predictions JSON (an unindexed full scan that,
        on a large library, held a series_page lock for minutes inside the merge
        and blocked migration 0040 — operator-flagged 2026-06-07). That scan only
        ever rediscovered the kind, so it's gone."""
        from ..models.tag_alias import TagAlias

        category = src_kind.value if hasattr(src_kind, "value") else str(src_kind)
        res = await self.session.execute(
            pg_insert(TagAlias)
            .values(
                alias_string=src_name,
                alias_category=category,
                canonical_tag_id=tgt,
            )
            .on_conflict_do_nothing(
                index_elements=["alias_string", "alias_category"]
            )
        )
        return bool(res.rowcount)


# ---------------------------------------------------------------------------
# #714: retro-normalize existing tags to the #701 canonical (Title Case +
# collapsed whitespace) and merge case/whitespace-variant duplicates.
# ---------------------------------------------------------------------------

_NORMALIZE_SAMPLE_CAP = 50


def _group_existing_tags(
    rows,
) -> dict[tuple, list[tuple[int, str]]]:
    """Group (id, name, kind, fandom_id) rows by their post-normalization
    identity: (kind, COALESCE(fandom_id, -1), canonical_name). Every tag that
    would collapse to the same canonical lives in ONE group, so the canonical
    form is unique within (kind, fandom) once each group is resolved."""
    groups: dict[tuple, list[tuple[int, str]]] = {}
    for tag_id, name, kind, fandom_id in rows:
        canonical = normalize_tag_name(name)
        key = (kind, fandom_id if fandom_id is not None else -1, canonical)
        groups.setdefault(key, []).append((tag_id, name))
    return groups


def _group_needs_change(canonical: str, members: list[tuple[int, str]]) -> bool:
    """A group is already canonical iff it's a single member whose name equals
    the canonical form. Anything else (a collision, or a lone mis-cased tag)
    needs work."""
    if len(members) > 1:
        return True
    return members[0][1] != canonical


def _best_connected(tag_ids: list[int], counts: dict[int, int]) -> int:
    """The tag with the most image associations (→ fewest FK repoints when it
    survives), tie-broken to the lowest id for determinism. Module-level so the
    key closes over its parameter, not a loop variable (ruff B023)."""
    return max(tag_ids, key=lambda tid: (counts.get(tid, 0), -tid))


async def normalize_existing_tags(
    session: AsyncSession,
    *,
    dry_run: bool = False,
    time_budget_seconds: float | None = None,
) -> dict:
    """Convert the back-catalog to the #701 canonical tag form.

    For each (kind, fandom, canonical) group: pick a survivor, merge any
    case/whitespace-variant siblings INTO it via the tested merge path
    (TagService._do_merge — FK repoints + protective aliases), then rename the
    survivor to the canonical form. Idempotent: a group that is already a lone
    canonical tag is a no-op, so re-running is safe.

    A first run over a fresh back-catalog can touch tens of thousands of tags
    (the whole booru-derived vocabulary needs recasing) and won't finish inside
    one Celery time limit — it timed out at 40 min (operator-flagged 2026-06-07).
    `time_budget_seconds` time-boxes the live run: it stops cleanly at the budget
    and reports `partial`/`remaining` so the caller can re-enqueue and continue.
    Because it commits per group and is idempotent, the next run just picks up
    the groups still needing change.

    dry_run=True returns a projection (counts + a sample of the changes) with no
    mutations. Live runs commit per group and isolate failures per group so one
    bad group can't strand the rest.

    Returns (dry_run):
      {"groups": N, "collisions": M, "tags_to_merge": K, "tags_to_rename": R,
       "total_changes": T, "sample": [{"to", "from": [...], "kind", "merge"}]}
    Returns (live):
      {"groups_processed", "merged", "renamed", "aliases_created", "errors",
       "total_changes", "remaining", "partial", "sample": [...]}
    """
    rows = (
        await session.execute(
            select(Tag.id, Tag.name, Tag.kind, Tag.fandom_id)
            # System tags (#128) are exempt: their names are the keys the
            # hygiene machinery matches on, and canonicalization would recase
            # 'wip' → 'Wip' (this path bypasses TagService.rename's guard).
            .where(Tag.is_system.is_(False))
        )
    ).all()
    groups = _group_existing_tags(rows)

    # Deterministic sample/ordering: by kind then canonical name.
    touched = sorted(
        (
            (key, members)
            for key, members in groups.items()
            if _group_needs_change(key[2], members)
        ),
        key=lambda km: (
            km[0][0].value if hasattr(km[0][0], "value") else str(km[0][0]),
            km[0][2].lower(),
        ),
    )

    sample = [
        {
            "to": key[2],
            "from": [name for _id, name in members],
            "kind": key[0].value if hasattr(key[0], "value") else str(key[0]),
            "merge": len(members) > 1,
        }
        for key, members in touched[:_NORMALIZE_SAMPLE_CAP]
    ]

    if dry_run:
        collisions = sum(1 for key, m in touched if len(m) > 1)
        tags_to_merge = sum(len(m) - 1 for key, m in touched if len(m) > 1)
        # A group renames iff the canonical form isn't already one of its
        # members' exact names (else that member is picked as survivor → no
        # rename, the rest merge into it).
        tags_to_rename = sum(
            1
            for key, m in touched
            if key[2] not in {name for _id, name in m}
        )
        return {
            "groups": len(groups),
            "collisions": collisions,
            "tags_to_merge": tags_to_merge,
            "tags_to_rename": tags_to_rename,
            "total_changes": len(touched),
            "sample": sample,
        }

    svc = TagService(session)
    summary = {
        "groups_processed": 0,
        "merged": 0,
        "renamed": 0,
        "aliases_created": 0,
        "errors": 0,
        "total_changes": len(touched),
        "remaining": len(touched),
        "partial": False,
        "sample": sample,
    }
    start = time.monotonic()
    log.info(
        "normalize_existing_tags: %d group(s) need changes (budget=%ss)",
        len(touched), time_budget_seconds,
    )
    for done, (key, members) in enumerate(touched):
        # Time-box: stop cleanly before the Celery limit kills us mid-group and
        # strands the run as a timeout. The caller re-enqueues to finish the
        # rest (idempotent — already-canonical groups are skipped next pass).
        if (
            time_budget_seconds is not None
            and time.monotonic() - start >= time_budget_seconds
        ):
            summary["partial"] = True
            summary["remaining"] = len(touched) - done
            break
        # Heartbeat so a long run is diagnosable instead of silent — the timeout
        # operator-flagged 2026-06-07 produced zero logs because the only log was
        # per finished group and it was stuck mid-group on a lock.
        if done and done % 25 == 0:
            log.info(
                "normalize_existing_tags: %d/%d groups (%d merged, %d errors, "
                "%.0fs elapsed)",
                done, len(touched), summary["merged"], summary["errors"],
                time.monotonic() - start,
            )
        canonical = key[2]
        names_by_id = dict(members)
        # Survivor: prefer a member already named canonically (no rename, no
        # self-alias); else the best-connected (fewest FK repoints); else
        # lowest id for determinism.
        survivor_id = next(
            (tid for tid, name in members if name == canonical), None
        )
        if survivor_id is None:
            counts = await svc._image_assoc_counts(list(names_by_id))
            survivor_id = _best_connected(list(names_by_id), counts)
        loser_ids = [tid for tid in names_by_id if tid != survivor_id]
        try:
            survivor = await session.get(Tag, survivor_id)
            for loser_id in loser_ids:
                loser = await session.get(Tag, loser_id)
                if loser is None:
                    continue
                result = await svc._do_merge(loser, survivor)
                if result.alias_created:
                    summary["aliases_created"] += 1
            if survivor.name != canonical:
                survivor.name = canonical
                await session.flush()
                summary["renamed"] += 1
            await session.commit()
            summary["groups_processed"] += 1
            summary["merged"] += len(loser_ids)
        except Exception as exc:  # one bad group must not strand the rest
            await session.rollback()
            summary["errors"] += 1
            log.warning(
                "tag normalize failed for group %r: %s", canonical, exc
            )
    else:
        # Loop finished without hitting the time budget — nothing left to do.
        summary["remaining"] = 0
    return summary
