"""Tag centroids: the mean SigLIP embedding of a tag's member images.

Powers centroid-augmented suggestions (a tag whose centroid is close to an
image's embedding becomes a suggestion even if Camie didn't predict it).
"""

from dataclasses import dataclass

import numpy as np
from sqlalchemy import func, select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import (
    ImageRecord,
    MLSettings,
    Tag,
    TagKind,
    TagReferenceEmbedding,
)
from ...models.tag import image_tag
from .embedder import MODEL_VERSION as SIGLIP_VERSION

ELIGIBLE_KINDS = {
    TagKind.character,
    TagKind.fandom,
    TagKind.general,
    TagKind.series,
}


@dataclass(frozen=True)
class CentroidHit:
    tag_id: int
    similarity: float


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

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

    async def recompute_for_tag(self, tag_id: int) -> bool:
        """Recompute one tag's centroid. Returns True if a centroid was
        written, False if skipped (ineligible kind or too few members)."""
        tag = await self.session.get(Tag, tag_id)
        if tag is None or tag.kind not in ELIGIBLE_KINDS:
            return False

        min_refs = await self._min_reference_images()

        stmt = (
            select(ImageRecord.siglip_embedding)
            .join(image_tag, image_tag.c.image_record_id == ImageRecord.id)
            .where(image_tag.c.tag_id == tag_id)
            .where(ImageRecord.siglip_embedding.is_not(None))
        )
        embeddings = [
            np.array(e, dtype=np.float32)
            for e in (await self.session.execute(stmt)).scalars().all()
        ]
        if len(embeddings) < min_refs:
            return False

        centroid = np.mean(np.stack(embeddings), axis=0).astype(np.float32)

        stmt = insert(TagReferenceEmbedding).values(
            tag_id=tag_id,
            embedding=centroid.tolist(),
            reference_count=len(embeddings),
            model_version=SIGLIP_VERSION,
        )
        stmt = stmt.on_conflict_do_update(
            index_elements=["tag_id"],
            set_={
                "embedding": centroid.tolist(),
                "reference_count": len(embeddings),
                "model_version": SIGLIP_VERSION,
                "updated_at": func.now(),
            },
        )
        await self.session.execute(stmt)
        return True

    async def list_drifted(self) -> list[int]:
        """Tag ids whose centroid is stale: member count != reference_count,
        OR no centroid row, OR centroid built on a different SigLIP version.
        Only considers eligible-kind tags with embeddings present."""
        member_counts = (
            select(
                image_tag.c.tag_id.label("tag_id"),
                func.count(image_tag.c.image_record_id).label("members"),
            )
            .join(ImageRecord, ImageRecord.id == image_tag.c.image_record_id)
            .where(ImageRecord.siglip_embedding.is_not(None))
            .group_by(image_tag.c.tag_id)
            .subquery()
        )
        stmt = (
            select(Tag.id)
            .join(member_counts, member_counts.c.tag_id == Tag.id)
            .outerjoin(
                TagReferenceEmbedding,
                TagReferenceEmbedding.tag_id == Tag.id,
            )
            .where(Tag.kind.in_(ELIGIBLE_KINDS))
            .where(
                (TagReferenceEmbedding.tag_id.is_(None))
                | (
                    TagReferenceEmbedding.reference_count
                    != member_counts.c.members
                )
                | (TagReferenceEmbedding.model_version != SIGLIP_VERSION)
            )
        )
        return list((await self.session.execute(stmt)).scalars().all())

    async def find_similar_tags(
        self, image_id: int, limit: int = 20
    ) -> list[CentroidHit]:
        """Cosine similarity between an image's embedding and stored
        centroids. Returns top-`limit` by similarity DESC. pgvector's
        cosine_distance gives 1 - cosine_similarity."""
        img = await self.session.get(ImageRecord, image_id)
        if img is None or img.siglip_embedding is None:
            return []
        emb = img.siglip_embedding
        distance = TagReferenceEmbedding.embedding.cosine_distance(emb)
        stmt = (
            select(
                TagReferenceEmbedding.tag_id,
                (1 - distance).label("similarity"),
            )
            .order_by(distance.asc())
            .limit(limit)
        )
        rows = (await self.session.execute(stmt)).all()
        return [
            CentroidHit(tag_id=r.tag_id, similarity=float(r.similarity))
            for r in rows
        ]