Spoolman/spoolman/database/extra_field_query.py

"""Helpers for filtering and sorting extra fields."""

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import TYPE_CHECKING

import sqlalchemy
from sqlalchemy import Alias, ColumnElement, Select
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.orm import aliased
from sqlalchemy.sql.expression import FunctionElement

from spoolman.database import models
from spoolman.database.utils import LIKE_ESCAPE, SortOrder, escape_like
from spoolman.extra_field_registry import EntityType, ExtraField, ExtraFieldType, get_extra_fields

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession
    from sqlalchemy.orm.attributes import InstrumentedAttribute


class _JsonArrayFirstElement(FunctionElement):
    """Cross-database helper: return the first element of a JSON array stored as text."""

    name = "json_array_first_element"
    inherit_cache = True


@compiles(_JsonArrayFirstElement, "postgresql")
@compiles(_JsonArrayFirstElement, "cockroachdb")
def _compile_json_array_first_pg(element: _JsonArrayFirstElement, compiler: object, **kw: object) -> str:  # type: ignore[misc]
    """PostgreSQL/CockroachDB: CAST(value AS JSON)->>0 returns the first element as TEXT."""
    (col_expr,) = element.clauses
    col_sql = compiler.process(col_expr, **kw)  # type: ignore[union-attr]
    return f"(CAST({col_sql} AS JSON)->>0)"


@compiles(_JsonArrayFirstElement)
def _compile_json_array_first_default(element: _JsonArrayFirstElement, compiler: object, **kw: object) -> str:  # type: ignore[misc]
    """SQLite/MariaDB: json_extract(value, '$[0]') returns the first element as a scalar."""
    (col_expr,) = element.clauses
    col_sql = compiler.process(col_expr, **kw)  # type: ignore[union-attr]
    return f"JSON_EXTRACT({col_sql}, '$[0]')"


class _JsonScalarText(FunctionElement):
    """Cross-database helper: decode a top-level JSON scalar and return it as unquoted TEXT.

    Used to compare against the *decoded* value rather than reconstructing the exact JSON
    serialization the client wrote. This decouples filtering from json.dumps/JSON.stringify
    encoding quirks (non-ASCII escaping, surrounding quotes, etc.). Only valid for JSON string
    scalars (text/choice/datetime) - booleans and numbers are handled with numeric casts instead,
    since scalar decoding is not consistent for them across dialects.
    """

    name = "json_scalar_text"
    inherit_cache = True


@compiles(_JsonScalarText, "postgresql")
@compiles(_JsonScalarText, "cockroachdb")
def _compile_json_scalar_text_pg(element: _JsonScalarText, compiler: object, **kw: object) -> str:  # type: ignore[misc]
    """PostgreSQL/CockroachDB: CAST(value AS JSON) #>> '{}' extracts the root scalar as TEXT."""
    (col_expr,) = element.clauses
    col_sql = compiler.process(col_expr, **kw)  # type: ignore[union-attr]
    return f"(CAST({col_sql} AS JSON) #>> '{{}}')"


@compiles(_JsonScalarText, "mysql")
def _compile_json_scalar_text_mysql(element: _JsonScalarText, compiler: object, **kw: object) -> str:  # type: ignore[misc]
    """MySQL/MariaDB: JSON_EXTRACT keeps the surrounding quotes, so JSON_UNQUOTE is required."""
    (col_expr,) = element.clauses
    col_sql = compiler.process(col_expr, **kw)  # type: ignore[union-attr]
    return f"JSON_UNQUOTE(JSON_EXTRACT({col_sql}, '$'))"


@compiles(_JsonScalarText)
def _compile_json_scalar_text_default(element: _JsonScalarText, compiler: object, **kw: object) -> str:  # type: ignore[misc]
    """SQLite: json_extract(value, '$') already returns string scalars unquoted."""
    (col_expr,) = element.clauses
    col_sql = compiler.process(col_expr, **kw)  # type: ignore[union-attr]
    return f"JSON_EXTRACT({col_sql}, '$')"


class _JsonArraySecondElement(FunctionElement):
    """Cross-database helper: return the second element of a JSON array stored as text."""

    name = "json_array_second_element"
    inherit_cache = True


@compiles(_JsonArraySecondElement, "postgresql")
@compiles(_JsonArraySecondElement, "cockroachdb")
def _compile_json_array_second_pg(element: _JsonArraySecondElement, compiler: object, **kw: object) -> str:  # type: ignore[misc]
    """PostgreSQL/CockroachDB: CAST(value AS JSON)->>1 returns the second element as TEXT."""
    (col_expr,) = element.clauses
    col_sql = compiler.process(col_expr, **kw)  # type: ignore[union-attr]
    return f"(CAST({col_sql} AS JSON)->>1)"


@compiles(_JsonArraySecondElement)
def _compile_json_array_second_default(element: _JsonArraySecondElement, compiler: object, **kw: object) -> str:  # type: ignore[misc]
    """SQLite/MariaDB: json_extract(value, '$[1]') returns the second element as a scalar."""
    (col_expr,) = element.clauses
    col_sql = compiler.process(col_expr, **kw)  # type: ignore[union-attr]
    return f"JSON_EXTRACT({col_sql}, '$[1]')"


@dataclass
class ExtraFieldJoin:
    """One extra field pulled onto a query as a plain column, ready to group or order by.

    Built by :func:`extra_field_join`, applied with :meth:`apply`. The two steps are separate
    because the value expression has to exist before the SELECT is constructed, while the join
    can only be added afterwards.
    """

    alias: Alias
    field_key: str
    id_column_name: str
    #: The DB-decoded scalar, i.e. the string the user typed rather than its JSON encoding.
    value: ColumnElement[str]

    def apply(self, stmt: Select, base_id_column: InstrumentedAttribute[int]) -> Select:
        """Outer-join the field's row onto `stmt`.

        OUTER, so entities with no row for this key survive with a NULL value — the same
        "unset" bucket a nullable built-in column produces.
        """
        return stmt.join(
            self.alias,
            sqlalchemy.and_(
                self.alias.c[self.id_column_name] == base_id_column,
                self.alias.c.key == self.field_key,
            ),
            isouter=True,
        )


def extra_field_value_text(column: ColumnElement[str]) -> ColumnElement[str]:
    """Decode a stored extra-field value to its scalar, as plain TEXT.

    Compare against this rather than a reconstructed JSON string, so matching doesn't depend
    on how the value was encoded when it was written (quoting, non-ASCII escaping).
    """
    return _JsonScalarText(column)


def extra_field_join(entity_type: EntityType, field_key: str) -> ExtraFieldJoin:
    """Prepare an extra field to be selected as a column of an entity query.

    The field table is aliased so it stays a distinct FROM element: the extra-field *filters*
    reference the un-aliased table inside correlated subqueries, and without the alias those
    subqueries would correlate to this join instead of standing on their own.
    """
    field_table = _get_field_table_for_entity(entity_type)
    id_column_name = _get_entity_id_column(field_table).key
    alias = field_table.__table__.alias()
    return ExtraFieldJoin(
        alias=alias,
        field_key=field_key,
        id_column_name=id_column_name,
        value=extra_field_value_text(alias.c.value),
    )


def _get_field_table_for_entity(entity_type: EntityType) -> type[models.Base]:
    """Map an entity type to its extra-field table."""
    if entity_type == EntityType.spool:
        return models.SpoolField
    if entity_type == EntityType.filament:
        return models.FilamentField
    if entity_type == EntityType.vendor:
        return models.VendorField
    raise ValueError(f"Unknown entity type: {entity_type}")


def _get_entity_id_column(field_table: type[models.Base]) -> InstrumentedAttribute[int]:
    """Map an extra-field table to its owning entity id column."""
    if field_table == models.SpoolField:
        return models.SpoolField.spool_id
    if field_table == models.FilamentField:
        return models.FilamentField.filament_id
    if field_table == models.VendorField:
        return models.VendorField.vendor_id
    raise ValueError(f"Unknown field table: {field_table}")


def _parse_boolean_filter(value: str) -> bool:
    """Parse a boolean filter using explicit true/false tokens only."""
    normalized = value.strip().lower()
    if normalized == "true":
        return True
    if normalized == "false":
        return False
    raise ValueError(f"Invalid boolean filter value: {value!r}")


async def apply_extra_field_filters_and_sort(
    *,
    db: AsyncSession,
    stmt: Select,
    base_obj: type[models.Base],
    entity_type: EntityType,
    extra_field_filters: dict[str, str] | None,
    sort_by: dict[str, SortOrder] | None,
) -> Select:
    """Apply extra-field filtering and sorting to a query."""
    if not extra_field_filters and not (sort_by is not None and any(field.startswith("extra.") for field in sort_by)):
        return stmt

    extra_fields = await get_extra_fields(db, entity_type)
    extra_fields_dict: dict[str, ExtraField] = {field.key: field for field in extra_fields}

    if extra_field_filters:
        field_table = _get_field_table_for_entity(entity_type)
        id_select = sqlalchemy.select(_get_entity_id_column(field_table))
        for field_key, value in extra_field_filters.items():
            field = extra_fields_dict.get(field_key)
            if field is None:
                continue
            stmt = add_where_clause_extra_field(
                stmt=stmt,
                link_column=base_obj.id,
                id_select=id_select,
                field_table=field_table,
                field_key=field_key,
                field_type=field.field_type,
                value=value,
                multi_choice=field.multi_choice if field.field_type == ExtraFieldType.choice else None,
            )

    if sort_by is not None:
        for field_name, order in sort_by.items():
            if not field_name.startswith("extra."):
                continue

            field_key = field_name[6:]
            extra_field = extra_fields_dict.get(field_key)
            if extra_field is None:
                continue

            stmt = add_order_by_extra_field(
                stmt=stmt,
                base_obj=base_obj,
                entity_type=entity_type,
                field_key=field_key,
                field_type=extra_field.field_type,
                order=order,
            )

    return stmt


async def apply_spool_related_extra_filters(
    *,
    db: AsyncSession,
    stmt: Select,
    filament_filters: dict[str, str] | None,
    vendor_filters: dict[str, str] | None,
) -> Select:
    """Filter a spool query by extra fields that live on the spool's filament or its vendor.

    Spool extra fields are handled by apply_extra_field_filters_and_sort; this handles the two
    related entities, linking through Spool.filament_id (a spool matches when its filament — or that
    filament's vendor — matches the extra-field condition).
    """
    if filament_filters:
        fields = {f.key: f for f in await get_extra_fields(db, EntityType.filament)}
        for field_key, value in filament_filters.items():
            field = fields.get(field_key)
            if field is None:
                continue
            stmt = add_where_clause_extra_field(
                stmt=stmt,
                link_column=models.Spool.filament_id,
                id_select=sqlalchemy.select(models.FilamentField.filament_id),
                field_table=models.FilamentField,
                field_key=field_key,
                field_type=field.field_type,
                value=value,
                multi_choice=field.multi_choice if field.field_type == ExtraFieldType.choice else None,
            )

    if vendor_filters:
        fields = {f.key: f for f in await get_extra_fields(db, EntityType.vendor)}
        for field_key, value in vendor_filters.items():
            field = fields.get(field_key)
            if field is None:
                continue
            # Map matching vendors back to filament ids: a spool matches when its filament's vendor
            # matches. Alias Filament so this subquery doesn't collide with the outer query's join.
            fil = aliased(models.Filament)
            id_select = sqlalchemy.select(fil.id).join(
                models.VendorField,
                models.VendorField.vendor_id == fil.vendor_id,
            )
            stmt = add_where_clause_extra_field(
                stmt=stmt,
                link_column=models.Spool.filament_id,
                id_select=id_select,
                field_table=models.VendorField,
                field_key=field_key,
                field_type=field.field_type,
                value=value,
                multi_choice=field.multi_choice if field.field_type == ExtraFieldType.choice else None,
            )

    return stmt


def add_where_clause_extra_field(  # noqa: C901, PLR0912, PLR0915
    stmt: Select,
    *,
    link_column: InstrumentedAttribute[int],
    id_select: Select,
    field_table: type[models.Base],
    field_key: str,
    field_type: ExtraFieldType,
    value: str,
    multi_choice: bool | None = None,
) -> Select:
    """Add a where clause to a select statement for an extra field.

    `link_column` is the column on the outer query to constrain (e.g. Spool.id, or Spool.filament_id
    when filtering spools by a filament/vendor extra field). `id_select` is a SELECT of ids in that
    same space (with any join to `field_table` already applied); this function narrows a copy of it
    by the field key/value conditions. This indirection is what lets one query filter by extra fields
    that live on a related entity.
    """
    conditions = []
    for value_part in value.split(","):
        # Empty-string filters follow the existing string-query API semantics.
        if len(value_part) == 0:
            empty_conditions = [
                field_table.value.is_(None),
                field_table.value == "null",
            ]
            if field_type == ExtraFieldType.boolean:
                empty_conditions.append(field_table.value == json.dumps(bool(0)))

            field_has_empty_value = id_select.where(
                sqlalchemy.and_(field_table.key == field_key, sqlalchemy.or_(*empty_conditions))
            )
            # The linked entity either stores an empty value, or has no row for this key at all.
            conditions.append(link_column.in_(field_has_empty_value))
            conditions.append(link_column.not_in(id_select.where(field_table.key == field_key)))
            continue

        exact_match = value_part.startswith('"') and value_part.endswith('"')
        parsed_value = value_part[1:-1] if exact_match else value_part

        if field_type == ExtraFieldType.text:
            # Compare against the DB-decoded scalar rather than a reconstructed JSON string, so
            # matching is independent of how the value was encoded (quote handling, non-ASCII
            # escaping). Substring search escapes LIKE wildcards to match user input literally.
            decoded = _JsonScalarText(field_table.value)
            field_condition = (
                decoded == parsed_value
                if exact_match
                else decoded.ilike(f"%{escape_like(parsed_value)}%", escape=LIKE_ESCAPE)
            )
        elif field_type == ExtraFieldType.integer:
            if ":" in parsed_value:
                min_val_str, max_val_str = parsed_value.split(":", 1)
                int_conditions = []
                try:
                    stored = sqlalchemy.cast(field_table.value, sqlalchemy.Integer)
                    if min_val_str:
                        int_conditions.append(stored >= int(min_val_str))
                    if max_val_str:
                        int_conditions.append(stored <= int(max_val_str))
                except (ValueError, TypeError) as exc:
                    raise ValueError(f"Invalid integer range filter value for '{field_key}': {parsed_value}") from exc
                if not int_conditions:
                    raise ValueError(f"Invalid integer range filter value for '{field_key}': {parsed_value}")
                field_condition = sqlalchemy.and_(*int_conditions)
            else:
                try:
                    field_condition = field_table.value == json.dumps(int(parsed_value))
                except ValueError as exc:
                    raise ValueError(f"Invalid integer filter value for '{field_key}': {parsed_value}") from exc
        elif field_type == ExtraFieldType.float:
            if ":" in parsed_value:
                min_val_str, max_val_str = parsed_value.split(":", 1)
                float_conditions = []
                try:
                    stored = sqlalchemy.cast(field_table.value, sqlalchemy.Float)
                    if min_val_str:
                        float_conditions.append(stored >= float(min_val_str))
                    if max_val_str:
                        float_conditions.append(stored <= float(max_val_str))
                except (ValueError, TypeError) as exc:
                    raise ValueError(f"Invalid float range filter value for '{field_key}': {parsed_value}") from exc
                if not float_conditions:
                    raise ValueError(f"Invalid float range filter value for '{field_key}': {parsed_value}")
                field_condition = sqlalchemy.and_(*float_conditions)
            else:
                try:
                    # Compare numerically rather than by JSON string so that int-typed storage
                    # (e.g. "2") and non-canonical decimals (e.g. "2.50") match a "2.0" filter.
                    field_condition = sqlalchemy.cast(field_table.value, sqlalchemy.Float) == float(parsed_value)
                except ValueError as exc:
                    raise ValueError(f"Invalid float filter value for '{field_key}': {parsed_value}") from exc
        elif field_type == ExtraFieldType.boolean:
            field_condition = field_table.value == json.dumps(_parse_boolean_filter(parsed_value))
        elif field_type == ExtraFieldType.choice:
            if multi_choice:
                # Multi-choice is stored as a JSON array; match the JSON-encoded token as a
                # substring. json.dumps gives the exact quoted, escaped form the array element
                # is stored as, and LIKE wildcards in the token are escaped.
                token = json.dumps(parsed_value, ensure_ascii=False)
                field_condition = field_table.value.like(f"%{escape_like(token)}%", escape=LIKE_ESCAPE)
            else:
                # Compare against the DB-decoded scalar, independent of JSON encoding.
                field_condition = _JsonScalarText(field_table.value) == parsed_value
        elif field_type == ExtraFieldType.datetime:
            # Compare decoded ISO-8601 strings. Both bounds and stored values are the frontend's
            # canonical toISOString() output, so lexicographic comparison is chronological.
            decoded = _JsonScalarText(field_table.value)
            if "|" in parsed_value:
                start_str, end_str = parsed_value.split("|", 1)
                dt_conditions = []
                if start_str:
                    dt_conditions.append(decoded >= start_str)
                if end_str:
                    dt_conditions.append(decoded <= end_str)
                if not dt_conditions:
                    raise ValueError(
                        f"Invalid datetime range filter for '{field_key}': {parsed_value}. Expected '<start>|<end>'."
                    )
                field_condition = sqlalchemy.and_(*dt_conditions)
            else:
                field_condition = decoded == parsed_value
        elif field_type in (ExtraFieldType.integer_range, ExtraFieldType.float_range):
            if ":" not in parsed_value:
                raise ValueError(
                    f"Invalid range filter value for '{field_key}': {parsed_value}. Expected '<min>:<max>'."
                )
            min_val_str, max_val_str = parsed_value.split(":", 1)
            converter = int if field_type == ExtraFieldType.integer_range else float
            range_conditions = []
            try:
                cast_type = sqlalchemy.Integer if field_type == ExtraFieldType.integer_range else sqlalchemy.Float
                if min_val_str:
                    # stored_min >= filter_min: the range starts at or after the requested minimum.
                    stored_min = sqlalchemy.cast(_JsonArrayFirstElement(field_table.value), cast_type)
                    range_conditions.append(stored_min >= converter(min_val_str))
                if max_val_str:
                    # stored_max <= filter_max: the range ends at or before the requested maximum.
                    stored_max = sqlalchemy.cast(_JsonArraySecondElement(field_table.value), cast_type)
                    range_conditions.append(stored_max <= converter(max_val_str))
            except (ValueError, TypeError) as exc:
                range_kind = "integer" if field_type == ExtraFieldType.integer_range else "float"
                raise ValueError(f"Invalid {range_kind} range filter value for '{field_key}': {parsed_value}") from exc
            if not range_conditions:
                raise ValueError(
                    f"Invalid range filter value for '{field_key}': {parsed_value}. Expected '<min>:<max>'."
                )
            field_condition = sqlalchemy.and_(*range_conditions)
        else:
            raise ValueError(f"Unsupported extra field type for '{field_key}': {field_type}")

        matching_entities = id_select.where(sqlalchemy.and_(field_table.key == field_key, field_condition))
        conditions.append(link_column.in_(matching_entities))

    if not conditions:
        return stmt

    return stmt.where(sqlalchemy.or_(*conditions))


async def find_extra_field_values(
    *,
    db: AsyncSession,
    entity_type: EntityType,
    field_key: str,
) -> list[str]:
    """Find all distinct values currently stored for a scalar extra field.

    Intended for text/single-choice/datetime fields, whose values are stored as JSON string
    scalars. Values are DB-decoded (see _JsonScalarText) so the result is independent of JSON
    encoding, mirroring the built-in distinct-value endpoints (materials, locations, ...). Empty
    and null values are omitted; the result is sorted for a stable option order.
    """
    field_table = _get_field_table_for_entity(entity_type)
    decoded = _JsonScalarText(field_table.value)
    stmt = sqlalchemy.select(decoded).where(field_table.key == field_key).distinct()
    rows = await db.execute(stmt)
    values = {row[0] for row in rows.all() if row[0] is not None and row[0] != ""}
    return sorted(values)


def add_order_by_extra_field(
    stmt: Select,
    base_obj: type[models.Base],
    entity_type: EntityType,
    field_key: str,
    field_type: ExtraFieldType,
    order: SortOrder,
) -> Select:
    """Add an order-by clause to a select statement for an extra field."""
    field_table = _get_field_table_for_entity(entity_type)
    entity_id_column = _get_entity_id_column(field_table)

    value_subq = (
        sqlalchemy.select(field_table.value)
        .where(
            sqlalchemy.and_(
                field_table.key == field_key,
                entity_id_column == base_obj.id,
            )
        )
        .scalar_subquery()
        .correlate(base_obj)
    )

    if field_type == ExtraFieldType.integer:
        sort_expr = sqlalchemy.cast(value_subq, sqlalchemy.Integer)
    elif field_type == ExtraFieldType.float:
        sort_expr = sqlalchemy.cast(value_subq, sqlalchemy.Float)
    elif field_type in (ExtraFieldType.integer_range, ExtraFieldType.float_range):
        cast_type = sqlalchemy.Integer if field_type == ExtraFieldType.integer_range else sqlalchemy.Float
        # Use dialect-specific JSON first-element extraction, then cast to numeric.
        sort_expr = sqlalchemy.cast(_JsonArrayFirstElement(value_subq), cast_type)
    else:
        sort_expr = value_subq

    if order == SortOrder.ASC:
        return stmt.order_by(sort_expr.asc())
    return stmt.order_by(sort_expr.desc())