FabledCurator/backend/app/services/ml/suggestions.py

"""The suggestion read-path: raw predictions + centroids -> alias-resolved,
threshold-filtered, category-grouped, ranked suggestions for one image.
"""

from dataclasses import dataclass, field

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import (
    ImageRecord,
    MLSettings,
    Tag,
    TagSuggestionRejection,
)
from ...models.tag import image_tag
from .aliases import AliasService
from .centroids import CentroidService
from .tagger import SURFACED_CATEGORIES


@dataclass(frozen=True)
class Suggestion:
    # canonical_tag_id is None when this is a raw Camie tag with no alias and
    # no existing Tag row — accepting it will create the tag.
    canonical_tag_id: int | None
    display_name: str
    category: str
    score: float
    source: str  # 'tagger' | 'centroid' | 'both'
    creates_new_tag: bool


@dataclass
class SuggestionList:
    by_category: dict[str, list[Suggestion]] = field(default_factory=dict)


class SuggestionService:
    def __init__(self, session: AsyncSession):
        self.session = session
        self.aliases = AliasService(session)
        self.centroids = CentroidService(session)

    async def _settings(self) -> MLSettings:
        return (
            await self.session.execute(select(MLSettings).where(MLSettings.id == 1))
        ).scalar_one()

    def _threshold_for(self, s: MLSettings, category: str) -> float:
        # 'artist' (FC-2d-vii-c) and 'copyright' (2026-06-01) retired;
        # both fall through to the 1.01 "never surfaces" default like any
        # unsurfaced category.
        return {
            "character": s.suggestion_threshold_character,
            "general": s.suggestion_threshold_general,
        }.get(category, 1.01)

    async def for_image(self, image_id: int) -> SuggestionList:
        img = await self.session.get(ImageRecord, image_id)
        if img is None:
            return SuggestionList()

        settings = await self._settings()
        predictions: dict = img.tagger_predictions or {}

        applied = set(
            (
                await self.session.execute(
                    select(image_tag.c.tag_id).where(
                        image_tag.c.image_record_id == image_id
                    )
                )
            ).scalars().all()
        )
        rejected = set(
            (
                await self.session.execute(
                    select(TagSuggestionRejection.tag_id).where(
                        TagSuggestionRejection.image_record_id == image_id
                    )
                )
            ).scalars().all()
        )

        # --- Camie predictions ---
        candidates: list[tuple[str, str, float]] = []
        for name, p in predictions.items():
            category = p.get("category", "general")
            if category not in SURFACED_CATEGORIES:
                continue
            conf = float(p.get("confidence", 0.0))
            if conf < self._threshold_for(settings, category):
                continue
            candidates.append((name, category, conf))

        alias_map = await self.aliases.resolve_many(
            [(n, c) for n, c, _ in candidates]
        )

        merged: dict[object, Suggestion] = {}

        def _merge(key, sug: Suggestion):
            existing = merged.get(key)
            if existing is None:
                merged[key] = sug
            elif sug.score > existing.score:
                merged[key] = Suggestion(
                    canonical_tag_id=existing.canonical_tag_id,
                    display_name=existing.display_name,
                    category=existing.category,
                    score=sug.score,
                    source="both"
                    if existing.source != sug.source
                    else existing.source,
                    creates_new_tag=existing.creates_new_tag,
                )

        for name, category, conf in candidates:
            canonical = alias_map.get((name, category))
            if canonical is not None:
                if canonical.id in applied or canonical.id in rejected:
                    continue
                _merge(
                    canonical.id,
                    Suggestion(
                        canonical_tag_id=canonical.id,
                        display_name=canonical.name,
                        category=category,
                        score=conf,
                        source="tagger",
                        creates_new_tag=False,
                    ),
                )
            else:
                existing_tag = (
                    await self.session.execute(
                        select(Tag).where(Tag.name == name)
                    )
                ).scalars().first()
                if existing_tag is not None:
                    if (
                        existing_tag.id in applied
                        or existing_tag.id in rejected
                    ):
                        continue
                    _merge(
                        existing_tag.id,
                        Suggestion(
                            canonical_tag_id=existing_tag.id,
                            display_name=existing_tag.name,
                            category=category,
                            score=conf,
                            source="tagger",
                            creates_new_tag=False,
                        ),
                    )
                else:
                    _merge(
                        f"raw:{name}:{category}",
                        Suggestion(
                            canonical_tag_id=None,
                            display_name=name,
                            category=category,
                            score=conf,
                            source="tagger",
                            creates_new_tag=True,
                        ),
                    )

        # --- Centroid augmentation ---
        hits = await self.centroids.find_similar_tags(image_id, limit=30)
        for hit in hits:
            if hit.similarity < settings.centroid_similarity_threshold:
                continue
            if hit.tag_id in applied or hit.tag_id in rejected:
                continue
            tag = await self.session.get(Tag, hit.tag_id)
            if tag is None:
                continue
            cat = tag.kind.value if hasattr(tag.kind, "value") else str(tag.kind)
            display_cat = cat if cat in SURFACED_CATEGORIES else "general"
            _merge(
                tag.id,
                Suggestion(
                    canonical_tag_id=tag.id,
                    display_name=tag.name,
                    category=display_cat,
                    score=hit.similarity,
                    source="centroid",
                    creates_new_tag=False,
                ),
            )

        result = SuggestionList()
        for sug in merged.values():
            result.by_category.setdefault(sug.category, []).append(sug)
        for cat in result.by_category:
            result.by_category[cat].sort(key=lambda s: s.score, reverse=True)
        return result

    async def for_selection(
        self,
        image_ids: list[int],
        threshold: float = 0.8,
        top_k: int = 10,
    ) -> dict[str, list[dict]]:
        """Consensus suggestions across image_ids. A tag is included iff it
        was suggested for (or already applied to) >= threshold fraction of
        the selection AND was acceptable on >= 1 image. Confidence is the
        mean over images where it was suggested. Aggregated by
        canonical_tag_id; creates-new (no canonical id) suggestions are
        skipped (bulk Accept applies by tag id)."""
        if not image_ids:
            return {}
        threshold = min(1.0, max(0.0, threshold))
        total = len(image_ids)

        stats: dict[int, dict] = {}
        for image_id in image_ids:
            sl = await self.for_image(image_id)
            for category, items in sl.by_category.items():
                for s in items:
                    if s.canonical_tag_id is None or s.creates_new_tag:
                        continue
                    st = stats.get(s.canonical_tag_id)
                    if st is None:
                        st = {
                            "tag_id": s.canonical_tag_id,
                            "name": s.display_name,
                            "category": category,
                            "source": s.source,
                            "suggested_count": 0,
                            "sum_score": 0.0,
                        }
                        stats[s.canonical_tag_id] = st
                    st["suggested_count"] += 1
                    st["sum_score"] += s.score

        rows = (
            await self.session.execute(
                select(
                    image_tag.c.image_record_id, image_tag.c.tag_id
                ).where(image_tag.c.image_record_id.in_(image_ids))
            )
        ).all()
        applied_by_tag: dict[int, set[int]] = {}
        for iid, tid in rows:
            applied_by_tag.setdefault(tid, set()).add(iid)

        result: dict[str, list[dict]] = {}
        for st in stats.values():
            existing_count = len(applied_by_tag.get(st["tag_id"], set()))
            covered = st["suggested_count"] + existing_count
            coverage = covered / total
            if coverage < threshold or st["suggested_count"] < 1:
                continue
            result.setdefault(st["category"], []).append(
                {
                    "canonical_tag_id": st["tag_id"],
                    "name": st["name"],
                    "category": st["category"],
                    "confidence": round(
                        st["sum_score"] / st["suggested_count"], 4
                    ),
                    "coverage": round(coverage, 4),
                    "covered_count": covered,
                    "source": st["source"],
                }
            )
        for cat in result:
            result[cat].sort(key=lambda x: x["confidence"], reverse=True)
            result[cat] = result[cat][:top_k]
        return result