"""Allowlist semantics: accepting a suggestion adds the canonical tag to
image_tag AND to tag_allowlist; per-image removal/dismiss writes a rejection.
"""

from collections.abc import Sequence
from dataclasses import dataclass

from sqlalchemy import and_, delete, distinct, func, or_, select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import (
    ImagePrediction,
    MLSettings,
    Tag,
    TagAlias,
    TagAllowlist,
    TagSuggestionRejection,
)
from ...models.tag import image_tag
from .aliases import AliasService


@dataclass(frozen=True)
class AllowlistRow:
    tag_id: int
    tag_name: str
    tag_kind: str
    min_confidence: float
    applied_count: int      # image_tag rows currently carrying this tag
    coverage_count: int     # images a sweep WOULD cover at min_confidence


class AllowlistService:
    def __init__(self, session: AsyncSession):
        self.session = session
        self.aliases = AliasService(session)

    async def _apply_image_tag(self, image_id: int, tag_id: int, source: str):
        stmt = 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)

    async def _add_to_allowlist(self, tag_id: int) -> bool:
        """Returns True if newly added (caller should kick off retro-apply)."""
        exists = await self.session.get(TagAllowlist, tag_id)
        if exists is not None:
            return False
        self.session.add(TagAllowlist(tag_id=tag_id))
        await self.session.flush()
        return True

    async def _clear_rejection(self, image_id: int, tag_id: int):
        await self.session.execute(
            delete(TagSuggestionRejection)
            .where(TagSuggestionRejection.image_record_id == image_id)
            .where(TagSuggestionRejection.tag_id == tag_id)
        )

    async def accept(self, image_id: int, tag_id: int) -> bool:
        """Accept a suggestion. Returns True if the tag was newly added to
        the allowlist (the API layer enqueues apply_allowlist_tags then)."""
        await self._apply_image_tag(image_id, tag_id, source="ml_accepted")
        await self._clear_rejection(image_id, tag_id)
        return await self._add_to_allowlist(tag_id)

    async def add_alias_and_accept(
        self,
        image_id: int,
        alias_string: str,
        alias_category: str,
        canonical_tag_id: int,
    ) -> bool:
        await self.aliases.create(
            alias_string, alias_category, canonical_tag_id
        )
        return await self.accept(image_id, canonical_tag_id)

    async def dismiss(self, image_id: int, tag_id: int) -> None:
        stmt = insert(TagSuggestionRejection).values(
            image_record_id=image_id, tag_id=tag_id
        )
        stmt = stmt.on_conflict_do_nothing(
            index_elements=["image_record_id", "tag_id"]
        )
        await self.session.execute(stmt)

    async def reject_applied_tag(self, image_id: int, tag_id: int) -> None:
        """Operator removed an applied tag from an image. Remove the
        image_tag row AND record a rejection so the allowlist won't
        re-apply it on the next maintenance sweep."""
        await self.session.execute(
            image_tag.delete()
            .where(image_tag.c.image_record_id == image_id)
            .where(image_tag.c.tag_id == tag_id)
        )
        await self.dismiss(image_id, tag_id)

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

    async def update_threshold(
        self, tag_id: int, min_confidence: float
    ) -> None:
        row = await self.session.get(TagAllowlist, tag_id)
        if row is not None:
            # An allowlist tag can't auto-apply more permissively than the
            # ingest store floor — predictions below tagger_store_floor aren't
            # stored, so a lower min_confidence would behave identically to the
            # floor. Clamp so the stored threshold matches actual behavior
            # (#764).
            floor = await self._store_floor()
            row.min_confidence = max(min_confidence, floor)

    async def remove(self, tag_id: int) -> None:
        await self.session.execute(
            delete(TagAllowlist).where(TagAllowlist.tag_id == tag_id)
        )

    async def _coverage_match(self, tag: Tag):
        """The predicate over image_prediction rows that resolve to `tag`,
        mirroring tasks.ml._confidence_for_tag's resolution: a prediction whose
        raw_name equals the tag name (any category), OR an alias maps
        (raw_name, category) -> this tag. Returns a SQLAlchemy boolean clause.
        """
        alias_rows = (
            await self.session.execute(
                select(TagAlias.alias_string, TagAlias.alias_category).where(
                    TagAlias.canonical_tag_id == tag.id
                )
            )
        ).all()
        name_clause = ImagePrediction.raw_name == tag.name
        alias_clauses = [
            and_(
                ImagePrediction.raw_name == a,
                ImagePrediction.category == c,
            )
            for a, c in alias_rows
        ]
        return or_(name_clause, *alias_clauses) if alias_clauses else name_clause

    async def coverage(self, tag_id: int, threshold: float) -> int:
        """How many distinct images a sweep WOULD cover for this tag at
        `threshold`: images with a resolving prediction scoring >= threshold.
        The gross candidate pool (NOT minus already-applied/rejected) — it's
        the tuning signal for "lower the threshold and ~N more images qualify".
        """
        tag = await self.session.get(Tag, tag_id)
        if tag is None:
            return 0
        match = await self._coverage_match(tag)
        stmt = select(
            func.count(distinct(ImagePrediction.image_record_id))
        ).where(ImagePrediction.score >= threshold, match)
        return (await self.session.execute(stmt)).scalar_one()

    async def list_all(self) -> Sequence[AllowlistRow]:
        stmt = (
            select(
                TagAllowlist.tag_id,
                Tag.name,
                Tag.kind,
                TagAllowlist.min_confidence,
            )
            .join(Tag, Tag.id == TagAllowlist.tag_id)
            .order_by(Tag.name.asc())
        )
        rows = (await self.session.execute(stmt)).all()
        tag_ids = [r[0] for r in rows]

        # Applied counts in ONE grouped query (vs N per-row counts).
        applied: dict[int, int] = {}
        if tag_ids:
            applied = dict(
                (
                    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()
            )

        result = []
        for r in rows:
            # Coverage is per-tag (alias set differs); allowlist is small.
            cov = await self.coverage(r[0], r[3])
            result.append(
                AllowlistRow(
                    tag_id=r[0],
                    tag_name=r[1],
                    tag_kind=r[2].value if hasattr(r[2], "value") else str(r[2]),
                    min_confidence=r[3],
                    applied_count=applied.get(r[0], 0),
                    coverage_count=cov,
                )
            )
        return result