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 SortOrder
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}")


# Escape character for LIKE patterns. Deliberately not backslash: a backslash ESCAPE clause is
# ambiguous under MySQL/MariaDB string parsing. '/' renders safely on all four dialects.
_LIKE_ESCAPE = "/"


def _escape_like(value: str) -> str:
    """Escape LIKE wildcards so user input is matched literally, not as a wildcard pattern."""
    return (
        value.replace(_LIKE_ESCAPE, _LIKE_ESCAPE * 2).replace("%", f"{_LIKE_ESCAPE}%").replace("_", f"{_LIKE_ESCAPE}_")
    )


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