import base64
import json
import sys
import numpy as np

from datetime import datetime

from enum import Enum, auto, IntEnum
from sqlalchemy import String, and_, desc, func, or_, text
from sqlalchemy.orm import aliased, load_only, selectinload
from sqlalchemy.types import LargeBinary
from sqlmodel import (
    Column,
    DateTime,
    SQLModel,
    Field,
    Relationship,
    exists,
    select,
    Session,
)
from typing import ClassVar, Optional, List, TYPE_CHECKING

from .character import Character
from .face import Face
from .picture_project import PictureProjectMember
from .picture_set import PictureSet, PictureSetMember
from .picture_stack import PictureStack
from .quality import Quality
from .tag import Tag
from .tag_prediction import TagPrediction

from pixlstash.pixl_logging import get_logger

if TYPE_CHECKING:
    from .character import Character
    from .picture_likeness import PictureLikeness
    from .project import Project


# Configure logging for the module
logger = get_logger(__name__)


# Class for sorting mechanisms (replaces Enum)
class SortMechanism:
    class Keys(Enum):
        DATE = auto()
        IMPORTED_AT = auto()
        SCORE = auto()
        CHARACTER_LIKENESS = auto()
        PICTURE_STACKS = auto()
        IMAGE_SIZE = auto()
        SMART_SCORE = auto()
        TEXT_CONTENT = auto()

    MECHANISMS = {
        Keys.DATE: {
            "field": "created_at",
            "description": "Date Created",
        },
        Keys.IMPORTED_AT: {
            "field": "imported_at",
            "description": "Date Imported",
        },
        Keys.SCORE: {
            "field": "score",
            "description": "Score",
        },
        Keys.SMART_SCORE: {
            "field": None,
            "description": "Smart Score",
        },
        Keys.IMAGE_SIZE: {
            "field": None,  # Special case, not a direct field
            "description": "Image Size",
        },
        Keys.TEXT_CONTENT: {
            "field": None,  # Special case, requires Quality join
            "description": "Text Content",
        },
        Keys.CHARACTER_LIKENESS: {
            "field": "character_likeness",
            "description": "Similarity to ...",
        },
        Keys.PICTURE_STACKS: {
            "field": "id",
            "description": "Likeness Groups ...",
        },
    }

    def __init__(self, key, descending: bool = True):
        self.key = key
        self.descending = descending

    @property
    def field(self):
        return self.MECHANISMS[self.key]["field"]

    @classmethod
    def all(cls):
        mechanisms = []
        for key, data in cls.MECHANISMS.items():
            data = {"key": str(key.name), **data}
            mechanisms.append(data)
        return mechanisms

    @classmethod
    def from_string(cls, key_string: str, descending: bool = True) -> "SortMechanism":
        # Try by name
        if key_string in cls.Keys.__members__:
            return SortMechanism(cls.Keys[key_string], descending=descending)

        raise ValueError(f"{key_string!r} is not a valid SortMechanism")


class LikenessParameter(IntEnum):
    SIZE_BIN = 0
    BRIGHTNESS = 1
    CONTRAST = 2
    EDGE_DENSITY = 3
    NOISE_LEVEL = 4
    ASPECT_RATIO = 5
    PHASH_PREFIX = 6
    DATE = 7
    COLORFULNESS = 8
    LUMINANCE_ENTROPY = 9
    DOMINANT_HUE = 10


LIKENESS_PARAMETER_SENTINEL = -1.0


def _default_likeness_parameters() -> np.ndarray:
    return np.full(
        len(LikenessParameter), LIKENESS_PARAMETER_SENTINEL, dtype=np.float32
    )


class ExportType(Enum):
    FULL = "full"
    FACE = "face"

    @classmethod
    def from_string(cls, value: str) -> "ExportType":
        d = (value or "").lower()
        for member in cls:
            if member.value == d:
                return member
        return cls.FULL


class Picture(SQLModel, table=True):
    ExportType: ClassVar[type["ExportType"]] = ExportType
    id: int = Field(default=None, primary_key=True)
    file_path: Optional[str] = None
    description: Optional[str] = None
    format: Optional[str] = None
    width: Optional[int] = None
    height: Optional[int] = None
    size_bytes: Optional[int] = None
    size_bin_index: Optional[int] = Field(default=None, index=True)
    created_at: Optional[datetime] = Field(
        default=None, sa_column=Column("created_at", type_=DateTime, nullable=True)
    )
    imported_at: Optional[datetime] = Field(
        default=None,
        sa_column=Column("imported_at", type_=DateTime, nullable=True, index=True),
    )
    text_embedding: Optional[np.ndarray] = Field(
        sa_column=Column("text_embedding", LargeBinary, default=None, nullable=True)
    )
    image_embedding: Optional[np.ndarray] = Field(
        sa_column=Column("image_embedding", LargeBinary, default=None, nullable=True)
    )
    likeness_parameters: Optional[np.ndarray] = Field(
        default_factory=_default_likeness_parameters,
        sa_column=Column(
            "likeness_parameters", LargeBinary, default=None, nullable=True
        ),
    )
    perceptual_hash: Optional[str] = Field(
        default=None,
        sa_column=Column("perceptual_hash", String, default=None, nullable=True),
    )
    thumbnail_left: Optional[int] = Field(default=None)
    thumbnail_top: Optional[int] = Field(default=None)
    thumbnail_side: Optional[int] = Field(default=None)
    score: Optional[int] = None
    aesthetic_score: Optional[float] = None
    pixel_sha: Optional[str] = Field(default=None, index=True)
    deleted: bool = Field(default=False, index=True)
    stack_id: Optional[int] = Field(
        default=None, foreign_key="picturestack.id", index=True
    )
    stack_position: Optional[int] = Field(default=None, index=True)
    original_file_name: Optional[str] = Field(
        default=None,
        sa_column=Column("original_file_name", String, default=None, nullable=True),
    )
    comfyui_positive_prompt: Optional[str] = Field(
        default=None,
        sa_column=Column(
            "comfyui_positive_prompt", String, default=None, nullable=True
        ),
    )
    # JSON-serialised list[str] of checkpoint / UNET names used at generation time.
    # NULL means this picture has never been checked for ComfyUI metadata;
    # "[]" is the sentinel meaning checked but no models found.
    comfyui_models: Optional[str] = Field(
        default=None,
        sa_column=Column("comfyui_models", String, default=None, nullable=True),
    )
    # JSON-serialised list[str] of LoRA names used at generation time.
    # NULL means this picture has never been checked for ComfyUI metadata;
    # "[]" is the sentinel meaning checked but no LoRAs found.
    comfyui_loras: Optional[str] = Field(
        default=None,
        sa_column=Column("comfyui_loras", String, default=None, nullable=True),
    )
    project_id: Optional[int] = Field(
        default=None, foreign_key="project.id", index=True
    )
    tag_uncertainty: Optional[float] = Field(default=None, index=True)
    anomaly_tag_uncertainty: Optional[float] = Field(default=None, index=True)

    # Relationships
    quality: Optional["Quality"] = Relationship(
        back_populates="picture",
        sa_relationship_kwargs={
            # face Quality rows share the same picture_id as their parent picture,
            # so we must filter to face_id IS NULL to get only picture-level quality.
            "primaryjoin": "(Quality.picture_id == Picture.id) & (Quality.face_id == None)",
            "foreign_keys": "[Quality.picture_id]",
        },
    )
    faces: List["Face"] = Relationship(
        back_populates="picture",
        sa_relationship_kwargs={
            "overlaps": "characters",
            "cascade": "all, delete-orphan",
            "passive_deletes": True,
        },
    )
    tags: List["Tag"] = Relationship(
        back_populates="picture",
        sa_relationship_kwargs={
            "cascade": "all, delete-orphan",
            "passive_deletes": True,
            "foreign_keys": "[Tag.picture_id]",
        },
    )
    tag_predictions: List["TagPrediction"] = Relationship(
        back_populates="picture",
        sa_relationship_kwargs={
            "cascade": "all, delete-orphan",
            "passive_deletes": True,
            "foreign_keys": "[TagPrediction.picture_id]",
        },
    )
    characters: List["Character"] = Relationship(  # Many-to-many via Face
        back_populates="pictures",
        link_model=Face,
        sa_relationship_kwargs={"overlaps": "faces,picture,character"},
    )
    picture_sets: List["PictureSet"] = Relationship(
        back_populates="members", link_model=PictureSetMember
    )
    stack: Optional["PictureStack"] = Relationship(back_populates="pictures")
    projects: List["Project"] = Relationship(
        back_populates="pictures", link_model=PictureProjectMember
    )

    likeness_a: List["PictureLikeness"] = Relationship(
        back_populates="picture_a",
        sa_relationship_kwargs={
            "primaryjoin": "Picture.id==PictureLikeness.picture_id_a",
            "cascade": "all, delete-orphan",
            "passive_deletes": True,
        },
    )
    likeness_b: List["PictureLikeness"] = Relationship(
        back_populates="picture_b",
        sa_relationship_kwargs={
            "primaryjoin": "Picture.id==PictureLikeness.picture_id_b",
            "cascade": "all, delete-orphan",
            "passive_deletes": True,
        },
    )

    class Config:
        arbitrary_types_allowed = True

    def __hash__(self):
        # Use the unique id for hashing
        return hash(self.id)

    def __eq__(self, other):
        # Compare by id for equality
        if isinstance(other, Picture):
            return self.id == other.id
        return False

    def text_embedding_data(self):
        """
        Returns a structured dict for embedding: description, tags, character info,
        and ComfyUI generation metadata (positive prompt, models, LoRAs) when stored.
        """
        data = {
            "description": self.description or None,
            "tags": [
                tag.tag
                for tag in getattr(self, "tags", [])
                if getattr(tag, "tag", None)
            ],
            "characters": [],
        }
        for character in getattr(self, "characters", []):
            char_info = {
                "name": getattr(character, "name", None),
                "description": getattr(character, "description", None),
            }
            data["characters"].append(char_info)
        # comfyui_models is NULL when not yet checked; "[]" is the sentinel for
        # "checked, no models".  Only add the comfyui key when there is actually
        # something useful to embed.
        comfyui_models_list = (
            json.loads(self.comfyui_models) if self.comfyui_models else []
        )
        comfyui_loras_list = (
            json.loads(self.comfyui_loras) if self.comfyui_loras else []
        )
        if self.comfyui_positive_prompt or comfyui_models_list or comfyui_loras_list:
            data["comfyui"] = {
                "positive_prompt": self.comfyui_positive_prompt or None,
                "models": comfyui_models_list,
                "loras": comfyui_loras_list,
            }
        return data

    @classmethod
    def semantic_search(
        cls: "Picture",
        session: Session,
        query: str,
        query_words: List[str],
        text_to_embedding: callable,
        clip_text_to_embedding: callable = None,
        fuzzy_weight: float = 0.5,
        embedding_weight: float = 0.5,
        threshold: float = 0.0,
        offset: int = 0,
        limit: int = sys.maxsize,
        format: Optional[List[str]] = None,
        select_fields: Optional[List[str]] = None,
        include_deleted: bool = False,
        only_deleted: bool = False,
        include_unimported: bool = True,
        candidate_ids: Optional[List[int]] = None,
        min_score: Optional[int] = None,
        comfyui_models_filter: Optional[List[str]] = None,
        comfyui_loras_filter: Optional[List[str]] = None,
        tags_filter: Optional[List[str]] = None,
        tags_rejected_filter: Optional[List[str]] = None,
    ) -> List["Picture"]:
        """
        Hybrid semantic search: combines fuzzy tag search (levenshtein SQL function) and embedding similarity (cosine_similarity SQL function).
        Orders by combined score in SQL.
        """
        if candidate_ids is not None and not candidate_ids:
            return []
        # 1. Generate SBERT embedding for tag search (Text-to-Text)
        query_embedding = text_to_embedding(query)
        if query_embedding is None:
            logger.warning("Semantic search: Failed to generate SBERT embedding.")
            query_embedding_bytes = None
        else:
            query_embedding_bytes = query_embedding.tobytes()

        # 2. Generate CLIP embedding for visual search (Text-to-Image)
        if clip_text_to_embedding:
            clip_query_embedding = clip_text_to_embedding(query)
            clip_query_embedding_bytes = (
                clip_query_embedding.tobytes()
                if clip_query_embedding is not None
                else None
            )
        else:
            clip_query_embedding_bytes = None

        logger.debug(
            f"Performing semantic search for query='{query}' and query_words={query_words} with fuzzy_weight={fuzzy_weight}, embedding_weight={embedding_weight}"
        )

        query_str = " ".join(query_words)
        # Subquery: calculate levenshtein distance for all tags of each picture
        tag_query = select(
            Tag.picture_id,
            func.levenshtein_with_id(
                func.group_concat(Tag.tag, " "), query_str, Tag.picture_id
            ).label("min_tag_dist"),
        ).group_by(Tag.picture_id)

        if candidate_ids:
            tag_query = tag_query.where(Tag.picture_id.in_(candidate_ids))

        tag_subq = tag_query.subquery()

        # Calculate cosine similarity for both text (tags) and image (visuals) embeddings
        if query_embedding_bytes:
            text_sim = (
                func.coalesce(
                    func.cosine_similarity(cls.text_embedding, query_embedding_bytes),
                    0.0,
                )
                * 2.0
            )
        else:
            text_sim = 0.0

        if clip_query_embedding_bytes:
            # Boost logic: CLIP similarity for unrelated text-image pairs is low (0.1-0.2).
            # A good match is often 0.25-0.35.
            # Fuzzy match is 0.0 to 1.0 (usually 1.0 for matches).
            # To make CLIP comparable and impactful, we multiply by a factor (e.g., 2.5).
            # This brings 0.3 -> 0.75, which can rival a messy tag match.
            image_sim = (
                func.coalesce(
                    func.cosine_similarity(
                        cls.image_embedding, clip_query_embedding_bytes
                    ),
                    0.0,
                )
                * 2.5
            )
        else:
            image_sim = 0.0

        # Combined embedding score: average of text and image similarity to capture both explicit tags and visual concepts
        embedding_score_raw = (text_sim + image_sim) / 2.0
        embedding_score = func.sqrt(func.max(0.0, embedding_score_raw))

        raw_fuzzy_score = func.max(
            0.0, 1.0 - func.coalesce(tag_subq.c.min_tag_dist, 1.0)
        )
        fuzzy_score = func.pow(raw_fuzzy_score, 1.5)

        # Main query: join pictures with tag_subq, compute combined score
        stmt = (
            select(
                cls,
                (fuzzy_weight * fuzzy_score + embedding_weight * embedding_score).label(
                    "combined_score"
                ),
                fuzzy_score.label("fuzzy_score"),
                embedding_score.label("embedding_score"),
                tag_subq.c.min_tag_dist.label(
                    "min_tag_dist"
                ),  # Explicitly include min_tag_dist
            )
            .outerjoin(tag_subq, cls.id == tag_subq.c.picture_id)
            .order_by(desc("combined_score"))
            .offset(offset)
            .limit(limit)
        )

        if only_deleted:
            stmt = stmt.where(cls.deleted.is_(True))
        elif not include_deleted:
            stmt = stmt.where(cls.deleted.is_(False))

        if not include_unimported:
            stmt = stmt.where(cls.imported_at.is_not(None))

        # Apply select_fields logic (like in find)
        if select_fields:
            select_fields = list(set(select_fields) | {"id"})

            # Use load_only for scalar fields
            scalar_attrs = [
                getattr(cls, field)
                for field in cls.scalar_fields().intersection(select_fields)
            ]
            if scalar_attrs:
                stmt = stmt.options(load_only(*scalar_attrs))
            # Use selectinload for relationships present in select_fields
            rel_attrs = [
                getattr(cls, field)
                for field in cls.relationship_fields().intersection(select_fields)
            ]
            for rel_attr in rel_attrs:
                stmt = stmt.options(selectinload(rel_attr))

        if format:
            stmt = stmt.where(cls.format.in_(format))

        if candidate_ids:
            stmt = stmt.where(cls.id.in_(candidate_ids))

        if min_score is not None:
            stmt = stmt.where(cls.score >= min_score)

        if comfyui_models_filter:
            for i, m in enumerate(comfyui_models_filter):
                stmt = stmt.where(
                    text(
                        f"EXISTS (SELECT 1 FROM json_each(picture.comfyui_models) WHERE value = :comfyui_model_{i})"
                    ).bindparams(**{f"comfyui_model_{i}": m})
                )

        if comfyui_loras_filter:
            for i, m in enumerate(comfyui_loras_filter):
                stmt = stmt.where(
                    text(
                        f"EXISTS (SELECT 1 FROM json_each(picture.comfyui_loras) WHERE value = :comfyui_lora_{i})"
                    ).bindparams(**{f"comfyui_lora_{i}": m})
                )

        if tags_filter:
            for i, tag in enumerate(tags_filter):
                stmt = stmt.where(
                    text(
                        f"EXISTS (SELECT 1 FROM tag WHERE tag.picture_id = picture.id AND tag.tag = :tag_filter_{i})"
                    ).bindparams(**{f"tag_filter_{i}": tag})
                )

        if tags_rejected_filter:
            for i, tag in enumerate(tags_rejected_filter):
                stmt = stmt.where(
                    text(
                        f"EXISTS (SELECT 1 FROM tag_prediction WHERE tag_prediction.picture_id = picture.id AND tag_prediction.tag = :rejected_tag_filter_{i} AND tag_prediction.status = 'REJECTED')"
                    ).bindparams(**{f"rejected_tag_filter_{i}": tag})
                )

        results = session.exec(stmt).all()

        if results:
            top_rows = results[:20]
            header = f"{'rank':>4} {'id':>6} {'combined':>9} {'fuzzy':>7} {'embed':>7} {'min_tag':>8}"
            lines = [header]
            for idx, row in enumerate(top_rows, start=1):
                pic = row[0]
                combined_score = row[1] if row[1] is not None else 0.0
                fuzzy_val = row[2] if row[2] is not None else 0.0
                embed_val = row[3] if row[3] is not None else 0.0
                min_tag_val = row[4] if row[4] is not None else 0.0
                pic_id = getattr(pic, "id", None)
                lines.append(
                    f"{idx:>4} {pic_id:>6} {float(combined_score):>9.4f} {float(fuzzy_val):>7.4f} {float(embed_val):>7.4f} {float(min_tag_val):>8.4f}"
                )
            logger.info(
                "Semantic search score breakdown (top %d) query=%r:\n%s",
                len(top_rows),
                query,
                "\n".join(lines),
            )

        # Log tag contribution for each result
        for result in results:
            logger.debug(
                f"Picture ID: {result.Picture.id}, "
                f"Tag contribution: min_tag_dist={getattr(result, 'min_tag_dist', 'N/A')}, "
                f"fuzzy_score={getattr(result, 'fuzzy_score', 'N/A')}, "
                f"embedding_score={getattr(result, 'embedding_score', 'N/A')}, "
                f"combined_score={getattr(result, 'combined_score', 'N/A')}"
            )
        output = []
        for row in results:
            pic, combined_score, _, _ = (
                row[0],
                row[1],
                row[2],
                row[3],
            )
            if combined_score and combined_score >= threshold:
                output.append((pic, combined_score))
        return output

    @staticmethod
    def serialize_with_likeness(picture_and_score):
        pic, score = picture_and_score
        d = pic.to_serializable_dict()
        d["likeness_score"] = max(0.0, score)
        return d

    @classmethod
    def _get_stack_leader_ids(cls, session, only_deleted: bool = False) -> set[int]:
        """Return the set of picture IDs that are the leader of their stack.

        Uses a single window-function query (ROW_NUMBER OVER PARTITION BY stack_id)
        so cost is O(N) over stacked pictures rather than a correlated subquery
        that runs once per row. Matches the JS compareStackOrder ranking:
        lowest stack_position (NULLs last) → highest score → newest created_at → lowest id.
        """
        deleted_val = 1 if only_deleted else 0
        sql = text("""
            SELECT id FROM (
                SELECT id,
                       ROW_NUMBER() OVER (
                           PARTITION BY stack_id
                           ORDER BY COALESCE(stack_position, 999999) ASC,
                                    COALESCE(score, 0) DESC,
                                    created_at DESC,
                                    id ASC
                       ) AS rn
                FROM picture
                WHERE stack_id IS NOT NULL
                  AND deleted = :deleted
            ) WHERE rn = 1
        """)
        rows = session.execute(sql.bindparams(deleted=deleted_val)).all()
        return {row[0] for row in rows}

    @classmethod
    def find(
        cls,
        session,
        *,
        sort_mech: Optional[SortMechanism] = None,
        offset: int = 0,
        limit: int = sys.maxsize,
        select_fields: Optional[List[str]] = None,
        format: Optional[List[str]] = None,
        include_deleted: bool = False,
        only_deleted: bool = False,
        include_unimported: bool = True,
        stack_leaders_only: bool = False,
        comfyui_models_filter: Optional[List[str]] = None,
        comfyui_loras_filter: Optional[List[str]] = None,
        tags_filter: Optional[List[str]] = None,
        tags_rejected_filter: Optional[List[str]] = None,
        min_score: Optional[int] = None,
        **search,
    ) -> List["Picture"]:
        """
        Find pictures based on provided filters.
        """
        query = select(cls)

        logger.debug("Got search parameters: %s", search)
        if select_fields:
            # Always include 'id' in select_fields
            select_fields = list(set(select_fields) | {"id"})
            # Use load_only for scalar fields
            scalar_attrs = [
                getattr(cls, field)
                for field in cls.scalar_fields().intersection(select_fields)
            ]
            if scalar_attrs:
                query = query.options(load_only(*scalar_attrs))
            # Use selectinload for relationships present in select_fields
            rel_attrs = [
                getattr(cls, field)
                for field in cls.relationship_fields().intersection(select_fields)
            ]
            for rel_attr in rel_attrs:
                query = query.options(selectinload(rel_attr))

        if only_deleted:
            query = query.where(cls.deleted.is_(True))
        elif not include_deleted:
            query = query.where(cls.deleted.is_(False))

        if not include_unimported:
            query = query.where(cls.imported_at.is_not(None))

        for attr, value in search.items():
            if attr == "project_id":
                membership_query = select(PictureProjectMember.picture_id).where(
                    PictureProjectMember.picture_id == cls.id
                )
                if value is None:
                    query = query.where(~exists(membership_query))
                elif isinstance(value, list):
                    query = query.where(
                        exists(
                            membership_query.where(
                                PictureProjectMember.project_id.in_(value)
                            )
                        )
                    )
                else:
                    query = query.where(
                        exists(
                            membership_query.where(
                                PictureProjectMember.project_id == value
                            )
                        )
                    )
                continue
            if hasattr(cls, attr):
                if isinstance(value, list):
                    query = query.where(getattr(cls, attr).in_(value))
                else:
                    query = query.where(getattr(cls, attr) == value)

        if format:
            query = query.where(cls.format.in_(format))

        if min_score is not None:
            query = query.where(cls.score >= min_score)

        # Build comfyui filter conditions (applied below, aware of stack_leaders_only)
        comfyui_conditions = []
        if comfyui_models_filter:
            comfyui_conditions.extend(
                text(
                    f"EXISTS (SELECT 1 FROM json_each(picture.comfyui_models) WHERE value = :comfyui_model_{i})"
                ).bindparams(**{f"comfyui_model_{i}": m})
                for i, m in enumerate(comfyui_models_filter)
            )
        if comfyui_loras_filter:
            comfyui_conditions.extend(
                text(
                    f"EXISTS (SELECT 1 FROM json_each(picture.comfyui_loras) WHERE value = :comfyui_lora_{i})"
                ).bindparams(**{f"comfyui_lora_{i}": m})
                for i, m in enumerate(comfyui_loras_filter)
            )

        if tags_filter:
            for i, tag in enumerate(tags_filter):
                query = query.where(
                    text(
                        f"EXISTS (SELECT 1 FROM tag WHERE tag.picture_id = picture.id AND tag.tag = :tag_filter_{i})"
                    ).bindparams(**{f"tag_filter_{i}": tag})
                )

        if tags_rejected_filter:
            for i, tag in enumerate(tags_rejected_filter):
                query = query.where(
                    text(
                        f"EXISTS (SELECT 1 FROM tag_prediction WHERE tag_prediction.picture_id = picture.id AND tag_prediction.tag = :rejected_tag_filter_{i} AND tag_prediction.status = 'REJECTED')"
                    ).bindparams(**{f"rejected_tag_filter_{i}": tag})
                )

        if stack_leaders_only:
            leader_ids = cls._get_stack_leader_ids(session, only_deleted=only_deleted)
            if comfyui_conditions:
                # Find stack_ids that contain at least one picture matching the filter,
                # then show the leader of those stacks (plus any unstacked matches).
                matching_stack_subq = (
                    select(cls.stack_id)
                    .where(
                        cls.stack_id.is_not(None),
                        cls.deleted.is_(False),
                        or_(*comfyui_conditions),
                    )
                    .distinct()
                )
                query = query.where(
                    or_(
                        and_(cls.stack_id.is_(None), or_(*comfyui_conditions)),
                        and_(
                            cls.id.in_(leader_ids),
                            cls.stack_id.in_(matching_stack_subq),
                        ),
                    )
                )
            else:
                query = query.where(or_(cls.stack_id.is_(None), cls.id.in_(leader_ids)))
        elif comfyui_conditions:
            query = query.where(or_(*comfyui_conditions))

        if sort_mech:
            if sort_mech.key == SortMechanism.Keys.IMAGE_SIZE:
                # Sort by width * height
                if sort_mech.descending:
                    query = query.order_by(
                        (cls.width * cls.height).desc(), cls.id.desc()
                    )
                else:
                    query = query.order_by((cls.width * cls.height).asc(), cls.id.asc())
            elif sort_mech.key == SortMechanism.Keys.TEXT_CONTENT:
                query = query.join(
                    Quality,
                    (Quality.picture_id == cls.id) & Quality.face_id.is_(None),
                    isouter=True,
                )
                if sort_mech.descending:
                    query = query.order_by(Quality.text_score.desc(), cls.id.desc())
                else:
                    query = query.order_by(Quality.text_score.asc(), cls.id.asc())
            else:
                field_name = sort_mech.field
                field = (
                    getattr(cls, field_name, None)
                    if isinstance(field_name, str)
                    else None
                )
                if field is not None:
                    if sort_mech.descending:
                        query = query.order_by(field.desc(), cls.id.desc())
                    else:
                        query = query.order_by(field.asc(), cls.id.asc())
        if offset > 0 or limit != sys.maxsize:
            query = query.offset(offset).limit(limit)

        results = session.exec(query).all()
        return results

    @classmethod
    def metadata_fields(cls):
        """
        Return a list of simple scalar fields
        """
        return cls.scalar_fields() - cls.large_binary_fields()

    @classmethod
    def grid_fields(cls):
        """
        Return a minimal set of fields for grid listing.
        """
        return {
            "id",
            "width",
            "height",
            "format",
            "score",
            "created_at",
            "imported_at",
            "stack_id",
            "stack_position",
            "tag_uncertainty",
            "anomaly_tag_uncertainty",
        }

    @classmethod
    def scalar_fields(cls):
        """
        Return a list of simple scalar fields
        """
        return set(cls.__table__.columns.keys())

    @classmethod
    def relationship_fields(cls):
        """
        Return a list of relationship fields
        """
        return set(Picture.__mapper__.relationships.keys())

    @classmethod
    def large_binary_fields(cls):
        """
        Return a list of LargeBinary fields
        """
        return {
            field.name
            for field in cls.__table__.columns
            if isinstance(field.type, LargeBinary)
        }

    def to_serializable_dict(self):
        """
        Returns a dict suitable for JSON serialization, encoding all large binary fields as base64 if present.
        """
        d = self.model_dump()
        for field in self.large_binary_fields():
            val = d.get(field, None)
            if val is not None:
                try:
                    if isinstance(val, np.ndarray):
                        val_bytes = val.tobytes()
                    elif isinstance(val, (bytes, bytearray)):
                        val_bytes = val
                    else:
                        val_bytes = bytes(val)
                    d[field] = base64.b64encode(val_bytes).decode("utf-8")
                except Exception:
                    d[field] = None
        return d

    @classmethod
    def clear_field(cls, session, picture_ids, field_name: str):
        pictures = cls.find(session=session, id=picture_ids, select_fields=[field_name])
        for pic in pictures:
            if hasattr(pic, field_name):
                setattr(pic, field_name, None)
        session.add_all(pictures)
        session.commit()

    @classmethod
    def find_unassigned(
        cls,
        session,
        sort_mech: Optional[SortMechanism] = None,
        offset: int = 0,
        limit: int = sys.maxsize,
        format: list[str] | None = None,
        metadata_fields: list[str] | None = None,
        stack_leaders_only: bool = False,
        min_score: Optional[int] = None,
        project_id: Optional[int] = None,
        only_unassigned_project: bool = False,
    ):
        query = select(Picture)
        unassigned_conditions = cls.build_unassigned_conditions(
            enforce_stack_assignment=True,
            assignment_project_id=project_id,
            assignment_unassigned_project=only_unassigned_project,
        )
        query = query.where(
            *unassigned_conditions,
            Picture.deleted.is_(False),
        )

        project_membership_query = select(PictureProjectMember.picture_id).where(
            PictureProjectMember.picture_id == Picture.id
        )
        if only_unassigned_project:
            query = query.where(~exists(project_membership_query))
        elif project_id is not None:
            query = query.where(
                exists(
                    project_membership_query.where(
                        PictureProjectMember.project_id == project_id
                    )
                )
            )

        if format:
            query = query.where(Picture.format.in_(format))

        if min_score is not None:
            query = query.where(Picture.score >= min_score)

        select_fields = metadata_fields or cls.metadata_fields()
        if select_fields:
            select_fields = list(set(select_fields) | {"id"})
            scalar_attrs = [
                getattr(Picture, field)
                for field in Picture.scalar_fields().intersection(select_fields)
            ]
            if scalar_attrs:
                query = query.options(load_only(*scalar_attrs))
            rel_attrs = [
                getattr(Picture, field)
                for field in Picture.relationship_fields().intersection(select_fields)
            ]
            for rel_attr in rel_attrs:
                query = query.options(selectinload(rel_attr))

        if sort_mech:
            if sort_mech.key == SortMechanism.Keys.IMAGE_SIZE:
                order_expr = Picture.width * Picture.height
                query = query.order_by(
                    order_expr.desc() if sort_mech.descending else order_expr.asc(),
                    Picture.id.desc() if sort_mech.descending else Picture.id.asc(),
                )
            elif sort_mech.key == SortMechanism.Keys.TEXT_CONTENT:
                query = query.join(
                    Quality,
                    (Quality.picture_id == Picture.id) & Quality.face_id.is_(None),
                    isouter=True,
                )
                query = query.order_by(
                    Quality.text_score.desc()
                    if sort_mech.descending
                    else Quality.text_score.asc(),
                    Picture.id.desc() if sort_mech.descending else Picture.id.asc(),
                )
            else:
                field = getattr(Picture, sort_mech.field, None)
                if field is not None:
                    query = query.order_by(
                        field.desc() if sort_mech.descending else field.asc(),
                        Picture.id.desc() if sort_mech.descending else Picture.id.asc(),
                    )

        if (not stack_leaders_only) and (offset > 0 or limit != sys.maxsize):
            query = query.offset(offset).limit(limit)

        results = session.exec(query).all()

        if stack_leaders_only:
            seen_stack_ids = set()
            deduped = []
            for pic in results:
                stack_id = getattr(pic, "stack_id", None)
                if stack_id is None:
                    deduped.append(pic)
                    continue
                if stack_id in seen_stack_ids:
                    continue
                seen_stack_ids.add(stack_id)
                deduped.append(pic)

            if offset > 0 or limit != sys.maxsize:
                start = max(0, offset)
                if limit == sys.maxsize:
                    deduped = deduped[start:]
                else:
                    deduped = deduped[start : start + max(0, limit)]
            return deduped

        return results

    @classmethod
    def build_unassigned_conditions(
        cls,
        *,
        enforce_stack_assignment: bool = False,
        assignment_project_id: Optional[int] = None,
        assignment_unassigned_project: bool = False,
    ) -> list:
        """Build SQL predicates for pictures that should count as unassigned.

        When enforce_stack_assignment is enabled, any stack with at least one
        assigned member (character face or set membership) is treated as assigned,
        so all pictures in that stack are excluded from unassigned queries.

        When assignment project scope is provided, assignment is evaluated only
        against characters/sets in that same project scope.
        """
        if assignment_unassigned_project:
            project_scope = Character.project_id.is_(None)
            set_scope = PictureSet.project_id.is_(None)
        elif assignment_project_id is not None:
            project_scope = Character.project_id == assignment_project_id
            set_scope = PictureSet.project_id == assignment_project_id
        else:
            project_scope = None
            set_scope = None

        assigned_face_query = select(Face.id).where(
            Face.picture_id == cls.id,
            Face.character_id.is_not(None),
        )
        if project_scope is not None:
            assigned_face_query = assigned_face_query.join(
                Character, Character.id == Face.character_id
            ).where(project_scope)

        assigned_set_query = select(PictureSetMember.picture_id).where(
            PictureSetMember.picture_id == cls.id
        )
        if set_scope is not None:
            assigned_set_query = assigned_set_query.join(
                PictureSet,
                PictureSet.id == PictureSetMember.set_id,
            ).where(set_scope)

        conditions = [~exists(assigned_face_query), ~exists(assigned_set_query)]

        if not enforce_stack_assignment:
            return conditions

        stack_picture = aliased(cls)
        stack_assigned_face_query = (
            select(Face.id)
            .join(stack_picture, stack_picture.id == Face.picture_id)
            .where(
                stack_picture.stack_id.is_not(None),
                stack_picture.stack_id == cls.stack_id,
                Face.character_id.is_not(None),
            )
        )
        if project_scope is not None:
            stack_assigned_face_query = stack_assigned_face_query.join(
                Character, Character.id == Face.character_id
            ).where(project_scope)

        stack_assigned_set_query = (
            select(PictureSetMember.picture_id)
            .join(
                stack_picture,
                stack_picture.id == PictureSetMember.picture_id,
            )
            .where(
                stack_picture.stack_id.is_not(None),
                stack_picture.stack_id == cls.stack_id,
            )
        )
        if set_scope is not None:
            stack_assigned_set_query = stack_assigned_set_query.join(
                PictureSet,
                PictureSet.id == PictureSetMember.set_id,
            ).where(set_scope)

        stack_has_assigned_face = exists(stack_assigned_face_query)
        stack_has_set_member = exists(stack_assigned_set_query)
        conditions.extend([~stack_has_assigned_face, ~stack_has_set_member])
        return conditions

    @staticmethod
    def fetch_features(session, picture_ids):
        """
        Fetch faces for a list of picture IDs.
        Returns:
            faces_by_pic: dict[picture_id, list[Face]]
            hands_by_pic: dict (empty, hand detection removed)
            face_tags_by_face: dict (empty, face tags removed)
            hand_tags_by_hand: dict (empty, hand tags removed)
        """
        faces = session.exec(select(Face).where(Face.picture_id.in_(picture_ids))).all()

        faces_by_pic = {}
        for face in faces:
            faces_by_pic.setdefault(face.picture_id, []).append(face)

        return (
            faces_by_pic,
            {},
            {},
            {},
        )
