"""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 func, select
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import (
    ImagePrediction,
    ImageRecord,
    MLSettings,
    Tag,
    TagSuggestionRejection,
)
from ...models.tag import image_tag
from .aliases import AliasService
from .centroids import CentroidService
from .tag_name import normalize as normalize_tag_name
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
    # raw_name = the booru model vocab key behind this suggestion. It's the key
    # an alias MUST be stored under (resolution looks up the raw key), so the
    # modal needs it to author an alias correctly. None for centroid-only hits
    # (no underlying prediction → nothing to alias).
    raw_name: str | None = None
    # via_alias = this suggestion was surfaced because an operator alias remapped
    # the raw prediction to this canonical tag. Lets the UI mark it + offer undo.
    via_alias: bool = False


@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()

    async def _load_predictions(self, image_id: int) -> dict:
        """Predictions for one image from the normalized image_prediction
        table (#768), in the {raw_name: {category, confidence}} shape the rest
        of this service consumed from the old JSON column — so all downstream
        threshold/alias/merge logic is unchanged."""
        rows = (
            await self.session.execute(
                select(
                    ImagePrediction.raw_name,
                    ImagePrediction.category,
                    ImagePrediction.score,
                ).where(ImagePrediction.image_record_id == image_id)
            )
        ).all()
        return {
            r.raw_name: {"category": r.category, "confidence": r.score}
            for r in rows
        }

    def _threshold_for(
        self, s: MLSettings, category: str, override: float | None = None,
    ) -> 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.
        # override (the typed-dropdown "show everything the model saw" mode)
        # applies to the surfaced categories only — unsurfaced ones are already
        # skipped before the threshold check, so they can't leak in.
        if override is not None:
            return override
        return {
            "character": s.suggestion_threshold_character,
            "general": s.suggestion_threshold_general,
        }.get(category, 1.01)

    async def for_image(
        self, image_id: int, *, threshold_override: float | None = None,
    ) -> SuggestionList:
        """Ranked suggestions for one image.

        threshold_override surfaces EVERY stored tagger prediction (down to the
        ingest STORE_FLOOR) regardless of the configured per-category suggestion
        thresholds — backs the tag-input dropdown's "search all of the model's
        predictions, including low-confidence ones, in the canonical formatting"
        mode (operator-asked 2026-06-09). The Suggestions panel still calls with
        no override so it stays the curated above-threshold list."""
        img = await self.session.get(ImageRecord, image_id)
        if img is None:
            return SuggestionList()

        settings = await self._settings()
        predictions: dict = await self._load_predictions(image_id)

        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 carry (raw_name, display_name, category, confidence).
        # raw_name = the booru-formatted vocab key, kept for alias_map
        # lookup since alias rows are hand-curated against raw keys.
        # display_name = normalize_tag_name(raw_name) — what the operator
        # sees AND what gets written to tag.name on Accept.
        candidates: list[tuple[str, 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, threshold_override):
                continue
            display = normalize_tag_name(name)
            if display is None:
                # emoticon / pure-punctuation vocab entry — drop entirely
                continue
            candidates.append((name, display, category, conf))

        alias_map = await self.aliases.resolve_many(
            [(raw, c) for raw, _disp, c, _conf 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,
                    # Keep the alias identity from `existing`: the tagger pass
                    # (which carries raw_name / via_alias) runs before centroid
                    # augmentation, so it's always the first writer for a key.
                    raw_name=existing.raw_name,
                    via_alias=existing.via_alias,
                )

        for raw, display, category, conf in candidates:
            canonical = alias_map.get((raw, 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,
                        raw_name=raw,
                        via_alias=True,
                    ),
                )
            else:
                # Case-insensitive match on BOTH the raw camie key AND
                # the normalized form — covers legacy underscore-named
                # Tag rows accepted before normalization shipped, AND
                # any tag the operator created with the human form.
                existing_tag = (
                    await self.session.execute(
                        select(Tag).where(
                            func.lower(Tag.name).in_(
                                [raw.lower(), display.lower()]
                            )
                        )
                    )
                ).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,
                            raw_name=raw,
                            via_alias=False,
                        ),
                    )
                else:
                    _merge(
                        f"raw:{display}:{category}",
                        Suggestion(
                            canonical_tag_id=None,
                            display_name=display,
                            category=category,
                            score=conf,
                            source="tagger",
                            creates_new_tag=True,
                            raw_name=raw,
                            via_alias=False,
                        ),
                    )

        # --- 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
