# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable, Iterable, Sequence
from datetime import datetime
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Annotated,
    Any,
    Generic,
    Literal,
    TypeVar,
    cast,
    overload,
)

from fastapi import Depends, HTTPException, Query, status
from pendulum.parsing.exceptions import ParserError
from pydantic import AfterValidator, BaseModel, NonNegativeInt
from sqlalchemy import Column, String, and_, func, not_, or_, select as sql_select, true as sql_true
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.inspection import inspect
from sqlalchemy.sql.functions import FunctionElement

from airflow._shared.timezones import timezone
from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
from airflow.api_fastapi.core_api.base import OrmClause
from airflow.api_fastapi.core_api.security import GetUserDep
from airflow.configuration import conf
from airflow.models import Base
from airflow.models.asset import (
    AssetAliasModel,
    AssetEvent,
    AssetModel,
    AssetPartitionDagRun,
    DagScheduleAssetReference,
    TaskInletAssetReference,
    TaskOutletAssetReference,
    association_table,
)
from airflow.models.connection import Connection
from airflow.models.dag import DagModel, DagTag
from airflow.models.dag_favorite import DagFavorite
from airflow.models.dag_version import DagVersion
from airflow.models.dagrun import DagRun
from airflow.models.errors import ParseImportError
from airflow.models.hitl import HITLDetail
from airflow.models.pool import Pool
from airflow.models.taskinstance import TaskInstance
from airflow.models.variable import Variable
from airflow.models.xcom import XComModel
from airflow.typing_compat import Self
from airflow.utils.state import DagRunState, TaskInstanceState
from airflow.utils.types import DagRunType

if TYPE_CHECKING:
    from sqlalchemy.orm.attributes import InstrumentedAttribute
    from sqlalchemy.sql import ColumnElement, Select
    from sqlalchemy.sql.compiler import SQLCompiler

    from airflow.serialization.definitions.dag import SerializedDAG

T = TypeVar("T")

_FALLBACK_PAGE_LIMIT: int = conf.getint("api", "fallback_page_limit")


class _MySQLCollate(FunctionElement):
    """
    Wraps a SQL expression so that on MySQL it is emitted with an explicit ``COLLATE`` clause.

    On every other dialect the expression is passed through unchanged.

    This is needed when a computed expression (e.g. a ``CASE … END`` that mixes
    a stored ``VARCHAR`` column with a ``CAST(integer AS CHAR)``) ends up with
    MySQL coercibility ``NONE`` because the two branches carry different implicit
    collations.  Comparing such an expression with a bound parameter fails with
    "Illegal mix of collations".  Wrapping the expression in an explicit
    ``COLLATE`` gives it ``EXPLICIT`` coercibility, which MySQL accepts in all
    comparison operators.
    """

    type = String()
    inherit_cache = True

    def __init__(self, expr: ColumnElement[Any], collation: str) -> None:
        super().__init__(expr)
        self.collation = collation


@compiles(_MySQLCollate)
def _compile_mysql_collate_default(element: _MySQLCollate, compiler: SQLCompiler, **kw: Any) -> str:
    """Non-MySQL: render the inner expression without any COLLATE clause."""
    (expr,) = element.clauses
    return compiler.process(expr, **kw)


@compiles(_MySQLCollate, "mysql")
def _compile_mysql_collate_mysql(element: _MySQLCollate, compiler: SQLCompiler, **kw: Any) -> str:
    """MySQL: wrap the inner expression with the requested COLLATE clause."""
    (expr,) = element.clauses
    inner = compiler.process(expr, **kw)
    return f"({inner}) COLLATE {element.collation}"


class BaseParam(OrmClause[T], ABC):
    """Base class for path or query parameters with ORM transformation."""

    def __init__(self, value: T | None = None, skip_none: bool = True) -> None:
        super().__init__(value)
        self.attribute: ColumnElement | InstrumentedAttribute | None = None
        self.skip_none = skip_none

    def set_value(self, value: T | None) -> Self:
        self.value = value
        return self

    @classmethod
    @abstractmethod
    def depends(cls, *args: Any, **kwargs: Any) -> Self:
        pass


class LimitFilter(BaseParam[NonNegativeInt]):
    """Filter on the limit."""

    def to_orm(self, select: Select) -> Select:
        if self.value is None and self.skip_none:
            return select

        return select.limit(self.value)

    @classmethod
    def depends(cls, limit: NonNegativeInt = _FALLBACK_PAGE_LIMIT) -> LimitFilter:
        return cls().set_value(min(limit, conf.getint("api", "maximum_page_limit")))


class OffsetFilter(BaseParam[NonNegativeInt]):
    """Filter on offset."""

    def to_orm(self, select: Select) -> Select:
        if self.value is None and self.skip_none:
            return select
        return select.offset(self.value)

    @classmethod
    def depends(
        cls,
        offset: NonNegativeInt = 0,
    ) -> OffsetFilter:
        return cls().set_value(offset)


class _FavoriteFilter(BaseParam[bool]):
    """Filter Dags by favorite status."""

    def __init__(self, user_id: str, value: T | None = None, skip_none: bool = True) -> None:
        super().__init__(skip_none=skip_none)
        self.user_id = user_id

    def to_orm(self, select_stmt: Select) -> Select:
        if self.value is None and self.skip_none:
            return select_stmt

        if self.value:
            select_stmt = select_stmt.join(DagFavorite, DagFavorite.dag_id == DagModel.dag_id).where(
                DagFavorite.user_id == self.user_id
            )
        else:
            select_stmt = select_stmt.where(
                not_(
                    sql_select(DagFavorite)
                    .where(and_(DagFavorite.dag_id == DagModel.dag_id, DagFavorite.user_id == self.user_id))
                    .exists()
                )
            )

        return select_stmt

    @classmethod
    def depends(cls, user: GetUserDep, is_favorite: bool | None = Query(None)) -> _FavoriteFilter:
        instance = cls(user_id=str(user.get_id())).set_value(is_favorite)
        return instance


class _ExcludeStaleFilter(BaseParam[bool]):
    """Filter on is_stale."""

    def to_orm(self, select: Select) -> Select:
        if self.value and self.skip_none:
            return select.where(DagModel.is_stale != self.value)
        return select

    @classmethod
    def depends(cls, exclude_stale: bool = True) -> _ExcludeStaleFilter:
        return cls().set_value(exclude_stale)


class _PrefixPatternParam(BaseParam[str], ABC):
    """
    Shared prefix pattern: pipe ``|`` for OR, ``~`` → empty (match all), Unicode prefix range.

    .. note::
        Trailing non-alphanumeric characters in a search term are stripped before the range
        is computed. A range scan with a punctuation-terminated upper bound is unsafe under
        PostgreSQL's default locale-aware collation (``en_US.utf8`` sorts punctuation in a
        way that breaks the range), and additionally stopping the range at an alphanumeric
        character keeps the upper bound alphanumeric too, so the predicate stays usable by
        default btree indexes. A user who asks for prefix ``"test_"`` gets matches starting
        with ``"test"`` — a small over-match trade-off made explicit in the public
        ``*_prefix_pattern`` query-param description.
    """

    pipe_as_or: bool = True

    @staticmethod
    def _prefix_range_upper(term: str) -> str | None:
        """
        Compute the exclusive upper bound for a prefix range scan.

        Returns ``None`` if the term has no alphanumeric characters. Trailing non-alphanumeric
        characters are dropped before bumping the last character so the resulting upper bound
        is itself alphanumeric and behaves predictably under locale-aware collations. If
        incrementing would land outside the alphanumeric range (e.g. ``'z' → '{'``), we drop
        that character and retry.
        """
        while term and not term[-1].isalnum():
            term = term[:-1]
        if not term:
            return None
        last = ord(term[-1])
        if last >= 0x10FFFF:
            return _PrefixPatternParam._prefix_range_upper(term[:-1])
        bumped = chr(last + 1)
        if not bumped.isalnum():
            return _PrefixPatternParam._prefix_range_upper(term[:-1])
        return term[:-1] + bumped

    @staticmethod
    def _prefix_lower_bound(term: str) -> str:
        """Return the matching lower bound: strip trailing non-alphanumeric chars to pair with the upper."""
        while term and not term[-1].isalnum():
            term = term[:-1]
        return term

    @abstractmethod
    def _prefix_clause(self, term: str):
        """Return the SQL boolean for one prefix term (including empty string after ``~`` alias)."""

    def to_orm(self, select: Select) -> Select:
        # ``skip_none`` only gates the "no value" behavior for the callers that must keep
        # the filter slot present (e.g. ``QueryDagIdPrefixPatternSearchWithNone``); applying
        # a ``None`` value as a filter produces nonsense predicates, so always skip it here.
        if self.value is None:
            return select

        val_str = str(self.value)
        if self.pipe_as_or and "|" in val_str:
            search_terms = [term.strip() for term in val_str.split("|") if term.strip()]
            if search_terms:
                return select.where(or_(*(self._prefix_clause(term) for term in search_terms)))

        return select.where(self._prefix_clause(val_str))

    def transform_aliases(self, value: str | None) -> str | None:
        if value == "~":
            value = ""
        return value


_LIKE_ESCAPE_CHAR = "\\"


def _escape_like_pattern(value: str) -> str:
    r"""
    Escape SQL ``LIKE`` / ``ILIKE`` metacharacters in a user-supplied value.

    Use together with ``column.ilike(f"%{_escape_like_pattern(value)}%", escape="\\")`` on filter
    parameters that intend literal substring matching (so a user-supplied ``%`` or ``_`` does not
    widen the match beyond what the filter semantics promise). Search parameters that explicitly
    expose wildcard semantics (see :class:`_SearchParam`) must not call this — they want the
    metacharacters to pass through.
    """
    return (
        value.replace(_LIKE_ESCAPE_CHAR, _LIKE_ESCAPE_CHAR * 2)
        .replace("%", _LIKE_ESCAPE_CHAR + "%")
        .replace("_", _LIKE_ESCAPE_CHAR + "_")
    )


class _SearchParam(BaseParam[str]):
    """
    Substring search on a column using ``ILIKE '%term%'`` (case-insensitive).

    .. note::
        This full-match substring search most of the time prevents the database
        from using B-tree indexes on ``attribute``, which can be very slow on
        large tables. Prefer :class:`_PrefixSearchParam` (the ``*_prefix_pattern``
        query-param counterpart) when matching from the beginning of the value
        is acceptable.
    """

    def __init__(self, attribute: ColumnElement, skip_none: bool = True, pipe_as_or: bool = True) -> None:
        super().__init__(skip_none=skip_none)
        self.attribute: ColumnElement = attribute
        self.pipe_as_or = pipe_as_or

    def to_orm(self, select: Select) -> Select:
        if self.value is None and self.skip_none:
            return select

        val_str = str(self.value)
        if self.pipe_as_or and "|" in val_str:
            search_terms = [term.strip() for term in val_str.split("|") if term.strip()]
            if search_terms:
                return select.where(or_(*(self.attribute.ilike(f"%{term}%") for term in search_terms)))

        return select.where(self.attribute.ilike(f"%{val_str}%"))

    def transform_aliases(self, value: str | None) -> str | None:
        if value == "~":
            value = "%"
        return value

    @classmethod
    def depends(cls, *args: Any, **kwargs: Any) -> Self:
        raise NotImplementedError("Use search_param_factory instead , depends is not implemented.")


class _PrefixSearchParam(_PrefixPatternParam):
    """
    Prefix search on a column using range comparison (case-sensitive, index-friendly).

    Unlike :class:`_SearchParam`, wildcard characters are treated as literals and the query
    plan can use the column's default B-tree index for the range scan. Trailing
    non-alphanumeric characters in ``term`` are stripped first (see
    :class:`_PrefixPatternParam` for why).
    """

    def __init__(self, attribute: ColumnElement, skip_none: bool = True, pipe_as_or: bool = True) -> None:
        super().__init__(skip_none=skip_none)
        self.attribute: ColumnElement = attribute
        self.pipe_as_or = pipe_as_or

    def _prefix_clause(self, term: str):
        lower = self._prefix_lower_bound(term)
        if not lower:
            return self.attribute.is_not(None)
        upper = self._prefix_range_upper(term)
        if upper is None:
            return self.attribute >= lower
        return and_(self.attribute >= lower, self.attribute < upper)

    @classmethod
    def depends(cls, *args: Any, **kwargs: Any) -> Self:
        raise NotImplementedError("Use prefix_search_param_factory instead, depends is not implemented.")


class _TaskDisplayNamePrefixPatternParam(_PrefixPatternParam):
    """
    Prefix filter equivalent to :attr:`TaskInstance.task_display_name`, rewritten for composite-index use.

    The hybrid expression ``coalesce(_task_display_property_value, task_id)`` cannot use those indexes;
    this implementation applies an equivalent ``OR`` of simpler range predicates instead. Trailing
    non-alphanumeric characters in ``term`` are stripped first (see :class:`_PrefixPatternParam`).
    """

    def _prefix_clause(self, term: str):
        lower = self._prefix_lower_bound(term)
        if not lower:
            return sql_true()
        upper = self._prefix_range_upper(term)
        if upper is None:
            return or_(
                and_(
                    TaskInstance._task_display_property_value.is_(None),
                    TaskInstance.task_id >= lower,
                ),
                and_(
                    TaskInstance._task_display_property_value.is_not(None),
                    TaskInstance._task_display_property_value >= lower,
                ),
            )
        return or_(
            and_(
                TaskInstance._task_display_property_value.is_(None),
                TaskInstance.task_id >= lower,
                TaskInstance.task_id < upper,
            ),
            and_(
                TaskInstance._task_display_property_value.is_not(None),
                TaskInstance._task_display_property_value >= lower,
                TaskInstance._task_display_property_value < upper,
            ),
        )

    @classmethod
    def depends(
        cls,
        task_display_name_prefix_pattern: str | None = Query(
            default=None,
            description=(
                "Prefix match on task display name: optional ``_task_display_property_value`` else "
                "``task_id`` (same as ``coalesce``). Case-sensitive. Index-friendly alternative to "
                "``task_display_name_pattern``. On large databases, combine with ``dag_id_prefix_pattern`` "
                "(or a specific Dag in the path) so ``(dag_id, task_id, ...)`` indexes apply. "
                "Use ``|`` for OR. Use ``~`` to match all. Trailing non-alphanumeric characters in the "
                "term are stripped before matching so the range scan stays index-compatible under "
                "locale-aware collations."
            ),
        ),
    ) -> Self:
        param = cls()
        return param.set_value(param.transform_aliases(task_display_name_prefix_pattern))


class QueryTaskInstanceTaskGroupFilter(BaseParam[str]):
    """Task group filter - returns all tasks in the specified group."""

    def __init__(self, dag=None, skip_none: bool = True):
        super().__init__(skip_none=skip_none)
        self._dag: None | SerializedDAG = dag

    @property
    def dag(self) -> None | SerializedDAG:
        return self._dag

    @dag.setter
    def dag(self, value: None | SerializedDAG) -> None:
        self._dag = value

    def to_orm(self, select: Select) -> Select:
        if self.value is None and self.skip_none:
            return select

        if not self.dag:
            raise ValueError("Dag must be set before calling to_orm")

        if not hasattr(self.dag, "task_group"):
            return select

        # Exact matching on group_id
        task_groups = self.dag.task_group.get_task_group_dict()
        task_group = task_groups.get(self.value)
        if not task_group:
            raise HTTPException(
                status.HTTP_404_NOT_FOUND,
                detail={
                    "reason": "not_found",
                    "message": f"Task group {self.value} not found",
                },
            )

        return select.where(TaskInstance.task_id.in_(task.task_id for task in task_group.iter_tasks()))

    @classmethod
    def depends(
        cls,
        value: str | None = Query(
            alias="task_group_id",
            default=None,
            description="Filter by exact task group ID. Returns all tasks within the specified task group.",
        ),
    ) -> QueryTaskInstanceTaskGroupFilter:
        return cls(dag=None).set_value(value)


def search_param_factory(
    attribute: ColumnElement,
    pattern_name: str,
    skip_none: bool = True,
    pipe_as_or: bool = True,
) -> Callable[[str | None], _SearchParam]:
    pipe_clause = (
        "Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). "
        if pipe_as_or
        else "The pipe `|` is matched literally, not as an OR separator. "
    )
    DESCRIPTION = (
        "SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). "
        f"{pipe_clause}"
        "Regular expressions are **not** supported. "
        "\n\n"
        "**Performance note:** this full-match pattern is evaluated as ``ILIKE '%term%'`` and "
        "most of the time prevents the database from using B-tree indexes, which can be very "
        "slow on large tables. Prefer the equivalent "
        f"``{pattern_name.replace('_pattern', '_prefix_pattern')}`` parameter when possible."
    )

    def depends_search(
        value: str | None = Query(alias=pattern_name, default=None, description=DESCRIPTION),
    ) -> _SearchParam:
        search_parm = _SearchParam(attribute, skip_none, pipe_as_or=pipe_as_or)
        value = search_parm.transform_aliases(value)
        return search_parm.set_value(value)

    return depends_search


def prefix_search_param_factory(
    attribute: ColumnElement,
    prefix_pattern_name: str,
    skip_none: bool = True,
    pipe_as_or: bool = True,
) -> Callable[[str | None], _PrefixSearchParam]:
    """
    Build a FastAPI ``Depends`` returning a :class:`_PrefixSearchParam` for prefix matching.

    Prefer this over :func:`search_param_factory` for performance: prefix matching uses a
    B-tree index range scan, while substring matching requires a full table scan.
    """
    pipe_clause = (
        "Use the pipe `|` operator for OR logic (e.g. `dag1|dag2`). "
        if pipe_as_or
        else "The pipe `|` is part of the prefix, not an OR separator. "
    )
    DESCRIPTION = (
        "Prefix match — returns items whose value starts with the given string "
        f"(case-sensitive, index-friendly). {pipe_clause}"
        "Use `~` to match all. Wildcard characters (`%`, `_`) "
        "are treated as literal characters. Trailing non-alphanumeric characters "
        "in the prefix are stripped before matching so the range scan stays "
        "index-compatible under locale-aware collations — e.g. `test_` effectively "
        "matches items starting with `test`, and `s3://` matches items starting with "
        "`s3`."
    )

    def depends_prefix_search(
        value: str | None = Query(alias=prefix_pattern_name, default=None, description=DESCRIPTION),
    ) -> _PrefixSearchParam:
        search_parm = _PrefixSearchParam(attribute, skip_none, pipe_as_or=pipe_as_or)
        value = search_parm.transform_aliases(value)
        return search_parm.set_value(value)

    return depends_prefix_search


class SortParam(BaseParam[list[str]]):
    """Order result by the attribute."""

    MAX_SORT_PARAMS = 10

    def __init__(
        self,
        allowed_attrs: list[str],
        model: Base,
        to_replace: dict[str, str | Column | list[Column]] | None = None,
    ) -> None:
        super().__init__()
        self.allowed_attrs = allowed_attrs
        self.model = model
        self.to_replace = to_replace
        self._cached_resolution: list[tuple[str, ColumnElement, bool]] | None = None

    def set_value(self, value: list[str] | None) -> Self:
        self._cached_resolution = None
        return super().set_value(value)

    def _resolve(self) -> list[tuple[str, ColumnElement, bool]]:
        """Resolve sort columns as (attr_name, column, is_descending) tuples. Cached after first call."""
        if self._cached_resolution is not None:
            return self._cached_resolution

        if self.value is None:
            self.value = [self.get_primary_key_string()]

        order_by_values = self.value
        if len(order_by_values) > self.MAX_SORT_PARAMS:
            raise HTTPException(
                400,
                f"Ordering with more than {self.MAX_SORT_PARAMS} parameters is not allowed. Provided: {order_by_values}",
            )

        resolved: list[tuple[str, ColumnElement, bool]] = []
        for order_by_value in order_by_values:
            lstriped_orderby = order_by_value.lstrip("-")
            # Store the user-facing name in the resolved tuple. ``row_value`` resolves
            # it back to the actual row accessor via ``to_replace`` when reading values
            # for cursor encoding.
            attr_name = lstriped_orderby
            column: Column | None = None
            if self.to_replace:
                replacement = self.to_replace.get(lstriped_orderby, lstriped_orderby)
                if isinstance(replacement, str):
                    lstriped_orderby = replacement
                elif isinstance(replacement, list):
                    # Compound sort: expand the list into multiple sort entries.
                    # Each column's ORM key becomes its attr_name so that
                    # row_value() can read the corresponding attribute via
                    # getattr(row, attr_name) without further to_replace lookups.
                    is_desc = order_by_value.startswith("-")
                    for col in replacement:
                        col_attr_name = col.key
                        resolved.append((col_attr_name, col, is_desc))
                    continue
                else:
                    column = replacement

            if (self.allowed_attrs and lstriped_orderby not in self.allowed_attrs) and column is None:
                raise HTTPException(
                    400,
                    f"Ordering with '{lstriped_orderby}' is disallowed or "
                    f"the attribute does not exist on the model",
                )
            if column is None:
                column = getattr(self.model, lstriped_orderby)

            resolved.append((attr_name, column, order_by_value.startswith("-")))

        primary_key_column = self.get_primary_key_column()
        pk_name = self.get_primary_key_string()
        resolved_column_keys = {getattr(col, "key", None) for _, col, _ in resolved}
        if pk_name not in resolved_column_keys:
            pk_desc = bool(order_by_values and order_by_values[0].startswith("-"))
            resolved.append((pk_name, primary_key_column, pk_desc))

        self._cached_resolution = resolved
        return self._cached_resolution

    def to_orm(self, select: Select, *, reversed: bool = False) -> Select:
        if self.skip_none is False:
            raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")

        resolved = self._resolve()
        if reversed:
            columns = [col.asc() if is_desc else col.desc() for _, col, is_desc in resolved]
        else:
            columns = [col.desc() if is_desc else col.asc() for _, col, is_desc in resolved]
        return select.order_by(None).order_by(*columns)

    def get_resolved_columns(self) -> list[tuple[str, ColumnElement, bool]]:
        """Return resolved sort columns as (attr_name, column_element, is_descending) tuples."""
        return self._resolve()

    def row_value(self, row: Any, name: str) -> Any:
        """
        Extract the sort-key value for ``name`` from a result row.

        Resolves the accessor through ``to_replace`` for string aliases
        (e.g. ``{"dag_run_id": "run_id"}``). For column-form mappings
        (e.g. ``{"run_after": DagRun.run_after}``), resolves through the
        primary model's attribute so association proxies can still be used
        for cursor values. Raises ``NotImplementedError`` when the model
        exposes no such attribute rather than emitting a ``None`` cursor token.
        """
        if self.to_replace:
            replacement = self.to_replace.get(name)
            if isinstance(replacement, str):
                return getattr(row, replacement, None)
            if replacement is not None and not isinstance(replacement, list):
                # Column-form mapping resolves through the primary model's attribute,
                # often an association proxy onto the joined entity
                # (``TaskInstance.run_after`` -> ``dag_run.run_after``). Fail loudly if the
                # model exposes no such attribute, rather than emitting a ``None`` cursor token.
                try:
                    return getattr(row, name)
                except AttributeError:
                    raise NotImplementedError(
                        f"Cursor pagination cannot resolve column-form ``to_replace`` for "
                        f"``{name}``: the primary model exposes no such attribute. Add an "
                        f"association proxy, use a string alias, or sort by a primary-model column."
                    )
            # List-form replacements are expanded in _resolve() into individual entries
            # each using the column's own ORM key as attr_name, so ``name`` at this point
            # is already a concrete model attribute (e.g. ``_rendered_map_index`` or
            # ``map_index``) — fall through to the getattr below.
        return getattr(row, name, None)

    def get_primary_key_column(self) -> Column:
        """Get the primary key column of the model of SortParam object."""
        return inspect(self.model).primary_key[0]

    def get_primary_key_string(self) -> str:
        """Get the primary key string of the model of SortParam object."""
        return self.get_primary_key_column().name

    @classmethod
    def depends(cls, *args: Any, **kwargs: Any) -> Self:
        raise NotImplementedError("Use dynamic_depends, depends not implemented.")

    def dynamic_depends(self, default: str | Sequence[str] | None = None) -> Callable:
        # Include to_replace keys that are not already in allowed_attrs to avoid
        # duplicate entries in the spec description.
        allowed_set = set(self.allowed_attrs)
        to_replace_attrs = [k for k in self.to_replace if k not in allowed_set] if self.to_replace else []

        all_attrs = self.allowed_attrs + to_replace_attrs

        if default is None:
            default_list = [self.get_primary_key_string()]
        elif isinstance(default, str):
            default_list = [default]
        else:
            default_list = list(default)

        _order_by_query = Query(
            default=default_list,
            description=f"Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. "
            f"Supported attributes: `{', '.join(all_attrs) if all_attrs else self.get_primary_key_string()}`",
        )

        def inner(order_by: list[str] = _order_by_query) -> SortParam:
            return SortParam(self.allowed_attrs, self.model, self.to_replace).set_value(order_by)

        return inner


class FilterOptionEnum(Enum):
    """Filter options for FilterParam."""

    EQUAL = "eq"
    NOT_EQUAL = "ne"
    LESS_THAN = "lt"
    LESS_THAN_EQUAL = "le"
    GREATER_THAN = "gt"
    GREATER_THAN_EQUAL = "ge"
    IN = "in"
    NOT_IN = "not_in"
    ANY_EQUAL = "any_eq"
    ALL_EQUAL = "all_eq"
    IS_NONE = "is_none"
    CONTAINS = "contains"


class FilterParam(BaseParam[T]):
    """Filter on attribute."""

    def __init__(
        self,
        attribute: InstrumentedAttribute,
        value: T | None = None,
        filter_option: FilterOptionEnum = FilterOptionEnum.EQUAL,
        skip_none: bool = True,
    ) -> None:
        super().__init__(value, skip_none)
        self.attribute: InstrumentedAttribute = attribute
        self.value: T | None = value
        self.filter_option: FilterOptionEnum = filter_option

    def to_orm(self, select: Select) -> Select:
        if isinstance(self.value, (list, str)) and not self.value and self.skip_none:
            return select
        if self.value is None and self.skip_none:
            return select

        if isinstance(self.value, list):
            if self.filter_option == FilterOptionEnum.IN:
                return select.where(self.attribute.in_(self.value))
            if self.filter_option == FilterOptionEnum.NOT_IN:
                return select.where(self.attribute.notin_(self.value))
            if self.filter_option == FilterOptionEnum.ANY_EQUAL:
                conditions = [self.attribute == val for val in self.value]
                return select.where(or_(*conditions))
            if self.filter_option == FilterOptionEnum.ALL_EQUAL:
                conditions = [self.attribute == val for val in self.value]
                return select.where(and_(*conditions))
            raise HTTPException(
                400, f"Invalid filter option {self.filter_option} for list value {self.value}"
            )

        if self.filter_option == FilterOptionEnum.EQUAL:
            return select.where(self.attribute == self.value)
        if self.filter_option == FilterOptionEnum.NOT_EQUAL:
            return select.where(self.attribute != self.value)
        if self.filter_option == FilterOptionEnum.LESS_THAN:
            return select.where(self.attribute < self.value)
        if self.filter_option == FilterOptionEnum.LESS_THAN_EQUAL:
            return select.where(self.attribute <= self.value)
        if self.filter_option == FilterOptionEnum.GREATER_THAN:
            return select.where(self.attribute > self.value)
        if self.filter_option == FilterOptionEnum.GREATER_THAN_EQUAL:
            return select.where(self.attribute >= self.value)
        if self.filter_option == FilterOptionEnum.IS_NONE:
            if self.value is None:
                return select
            if self.value is False:
                return select.where(self.attribute.is_not(None))
            if self.value is True:
                return select.where(self.attribute.is_(None))
        if self.filter_option == FilterOptionEnum.CONTAINS:
            # For JSON/JSONB columns, convert to text before applying LIKE
            from sqlalchemy import Text, cast

            if str(self.attribute.type).upper() in ("JSON", "JSONB"):
                return select.where(cast(self.attribute, Text).contains(self.value))
            return select.where(self.attribute.contains(self.value))
        raise ValueError(f"Invalid filter option {self.filter_option} for value {self.value}")

    @classmethod
    def depends(cls, *args: Any, **kwargs: Any) -> Self:
        raise NotImplementedError("Use filter_param_factory instead , depends is not implemented.")


def filter_param_factory(
    attribute: ColumnElement | InstrumentedAttribute,
    _type: type,
    filter_option: FilterOptionEnum = FilterOptionEnum.EQUAL,
    filter_name: str | None = None,
    default_value: T | None = None,
    default_factory: Callable[[], T | None] | None = None,
    skip_none: bool = True,
    transform_callable: Callable[[T | None], Any] | None = None,
    *,
    description: str | None = None,
) -> Callable[[T | None], FilterParam[T | None]]:
    # if filter_name is not provided, use the attribute name as the default
    filter_name = filter_name or getattr(attribute, "name", str(attribute))
    # can only set either default_value or default_factory
    query = (
        Query(alias=filter_name, default_factory=default_factory, description=description)
        if default_factory is not None
        else Query(alias=filter_name, default=default_value, description=description)
    )

    def depends_filter(value: T | None = query) -> FilterParam[T | None]:
        if transform_callable:
            value = transform_callable(value)
        # Cast to InstrumentedAttribute for type compatibility
        attr = cast("InstrumentedAttribute", attribute)
        return FilterParam(attr, value, filter_option, skip_none)

    # add type hint to value at runtime
    depends_filter.__annotations__["value"] = _type

    return depends_filter


class _TagFilterModel(BaseModel):
    """Tag Filter Model with a match mode parameter."""

    tags: list[str]
    tags_match_mode: Literal["any", "all"] | None


class _TagsFilter(BaseParam[_TagFilterModel]):
    """Filter on tags."""

    def to_orm(self, select: Select) -> Select:
        if self.skip_none is False:
            raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")

        if not self.value or not self.value.tags:
            return select

        conditions = [DagModel.tags.any(DagTag.name == tag) for tag in self.value.tags]
        operator = or_ if not self.value.tags_match_mode or self.value.tags_match_mode == "any" else and_
        return select.where(operator(*conditions))

    @classmethod
    def depends(
        cls,
        tags: list[str] = Query(default_factory=list),
        tags_match_mode: Literal["any", "all"] | None = None,
    ) -> _TagsFilter:
        return cls().set_value(_TagFilterModel(tags=tags, tags_match_mode=tags_match_mode))


class _OwnersFilter(BaseParam[list[str]]):
    """Filter on owners."""

    def to_orm(self, select: Select) -> Select:
        if self.skip_none is False:
            raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")

        if not self.value:
            return select

        conditions = [
            DagModel.owners.ilike(f"%{_escape_like_pattern(owner)}%", escape=_LIKE_ESCAPE_CHAR)
            for owner in self.value
        ]
        return select.where(or_(*conditions))

    @classmethod
    def depends(cls, owners: list[str] = Query(default_factory=list)) -> _OwnersFilter:
        return cls().set_value(owners)


def _safe_parse_datetime(date_to_check: str) -> datetime:
    """
    Parse datetime and raise error for invalid dates.

    :param date_to_check: the string value to be parsed
    """
    if not date_to_check:
        raise ValueError(f"{date_to_check} cannot be None.")
    return _safe_parse_datetime_optional(date_to_check)


@overload
def _safe_parse_datetime_optional(date_to_check: str) -> datetime: ...


@overload
def _safe_parse_datetime_optional(date_to_check: None) -> None: ...


def _safe_parse_datetime_optional(date_to_check: str | None) -> datetime | None:
    """
    Parse datetime and raise error for invalid dates.

    Allow None values.

    :param date_to_check: the string value to be parsed
    """
    if date_to_check is None:
        return None
    try:
        return timezone.parse(date_to_check, strict=True)
    except (TypeError, ParserError):
        raise HTTPException(
            400, f"Invalid datetime: {date_to_check!r}. Please check the date parameter have this value."
        )


class _DagIdAssetReferenceFilter(BaseParam[list[str]]):
    """Search on dag_id."""

    def __init__(self, skip_none: bool = True) -> None:
        super().__init__(skip_none=skip_none)

    @classmethod
    def depends(cls, dag_ids: list[str] = Query(None)) -> _DagIdAssetReferenceFilter:
        # needed to handle cases where dag_ids=a1,b1
        if dag_ids and len(dag_ids) == 1 and "," in dag_ids[0]:
            dag_ids = dag_ids[0].split(",")
        return cls().set_value(dag_ids)

    def to_orm(self, select: Select) -> Select:
        if self.value is None and self.skip_none:
            return select

        # At this point, self.value is either a list[str] or None -> coerce falsy None to an empty list
        dag_ids = self.value or []
        return select.where(
            (AssetModel.scheduled_dags.any(DagScheduleAssetReference.dag_id.in_(dag_ids)))
            | (AssetModel.producing_tasks.any(TaskOutletAssetReference.dag_id.in_(dag_ids)))
            | (AssetModel.consuming_tasks.any(TaskInletAssetReference.dag_id.in_(dag_ids)))
        )


class Range(BaseModel, Generic[T]):
    """Range with a lower and upper bound."""

    lower_bound_gte: T | None
    lower_bound_gt: T | None
    upper_bound_lte: T | None
    upper_bound_lt: T | None


class RangeFilter(BaseParam[Range]):
    """Filter on range in between the lower and upper bound."""

    def __init__(self, value: Range | None, attribute: InstrumentedAttribute) -> None:
        super().__init__(value)
        self.attribute: InstrumentedAttribute = attribute

    def to_orm(self, select: Select) -> Select:
        if self.skip_none is False:
            raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")

        if self.value is None:
            return select

        if self.value.lower_bound_gte:
            select = select.where(self.attribute >= self.value.lower_bound_gte)
        if self.value.lower_bound_gt:
            select = select.where(self.attribute > self.value.lower_bound_gt)
        if self.value.upper_bound_lte:
            select = select.where(self.attribute <= self.value.upper_bound_lte)
        if self.value.upper_bound_lt:
            select = select.where(self.attribute < self.value.upper_bound_lt)

        return select

    @classmethod
    def depends(cls, *args: Any, **kwargs: Any) -> Self:
        raise NotImplementedError("Use the `range_filter_factory` function to create the dependency")

    def is_active(self) -> bool:
        """Check if the range filter has any active bounds."""
        return self.value is not None and (
            self.value.lower_bound_gte is not None
            or self.value.lower_bound_gt is not None
            or self.value.upper_bound_lte is not None
            or self.value.upper_bound_lt is not None
        )


class NullableDatetimeRangeFilter(RangeFilter):
    """
    RangeFilter for nullable datetime columns (``start_date``, ``end_date``), rewritten for index use.

    ``COALESCE(column, now())`` wraps the column in a function call that prevents PostgreSQL from
    using btree indexes, forcing sequential scans on large tables. This class emits equivalent
    ``OR`` predicates so each branch can be satisfied by an independent index scan.

    NULL semantics: ``start_date=NULL`` means the task has not started yet; ``end_date=NULL`` means
    the task is still running. For lower bounds the NULL branch passes unconditionally — a not-yet-
    started/ended task will eventually satisfy any past lower bound. For upper bounds the NULL branch
    is ``col IS NULL AND now() <= x``, preserving the COALESCE(col, now()) semantics without the
    function-wrap index penalty.
    """

    def to_orm(self, select: Select) -> Select:
        if self.skip_none is False:
            raise ValueError(f"Cannot set 'skip_none' to False on a {type(self)}")

        if self.value is None:
            return select

        if self.value.lower_bound_gte:
            x = self.value.lower_bound_gte
            select = select.where(or_(self.attribute >= x, self.attribute.is_(None)))
        if self.value.lower_bound_gt:
            x = self.value.lower_bound_gt
            select = select.where(or_(self.attribute > x, self.attribute.is_(None)))
        if self.value.upper_bound_lte:
            x = self.value.upper_bound_lte
            select = select.where(or_(self.attribute <= x, and_(self.attribute.is_(None), func.now() <= x)))
        if self.value.upper_bound_lt:
            x = self.value.upper_bound_lt
            select = select.where(or_(self.attribute < x, and_(self.attribute.is_(None), func.now() < x)))

        return select


def datetime_range_filter_factory(
    filter_name: str, model: Base, attribute_name: str | None = None
) -> Callable[[datetime | None, datetime | None, datetime | None, datetime | None], RangeFilter]:
    def depends_datetime(
        lower_bound_gte: datetime | None = Query(alias=f"{filter_name}_gte", default=None),
        lower_bound_gt: datetime | None = Query(alias=f"{filter_name}_gt", default=None),
        upper_bound_lte: datetime | None = Query(alias=f"{filter_name}_lte", default=None),
        upper_bound_lt: datetime | None = Query(alias=f"{filter_name}_lt", default=None),
    ) -> RangeFilter:
        attr = getattr(model, attribute_name or filter_name)
        range_val = Range(
            lower_bound_gte=lower_bound_gte,
            lower_bound_gt=lower_bound_gt,
            upper_bound_lte=upper_bound_lte,
            upper_bound_lt=upper_bound_lt,
        )
        if filter_name in ("start_date", "end_date"):
            return NullableDatetimeRangeFilter(range_val, attr)
        return RangeFilter(range_val, attr)

    return depends_datetime


def float_range_filter_factory(
    filter_name: str, model: Base
) -> Callable[[float | None, float | None, float | None, float | None], RangeFilter]:
    def depends_float(
        lower_bound_gte: float | None = Query(alias=f"{filter_name}_gte", default=None),
        lower_bound_gt: float | None = Query(alias=f"{filter_name}_gt", default=None),
        upper_bound_lte: float | None = Query(alias=f"{filter_name}_lte", default=None),
        upper_bound_lt: float | None = Query(alias=f"{filter_name}_lt", default=None),
    ) -> RangeFilter:
        return RangeFilter(
            Range(
                lower_bound_gte=lower_bound_gte,
                lower_bound_gt=lower_bound_gt,
                upper_bound_lte=upper_bound_lte,
                upper_bound_lt=upper_bound_lt,
            ),
            getattr(model, filter_name),
        )

    return depends_float


# Common Safe DateTime
DateTimeQuery = Annotated[str, AfterValidator(_safe_parse_datetime)]
OptionalDateTimeQuery = Annotated[str | None, AfterValidator(_safe_parse_datetime_optional)]

# Dag
QueryLimit = Annotated[LimitFilter, Depends(LimitFilter.depends)]
QueryOffset = Annotated[OffsetFilter, Depends(OffsetFilter.depends)]
QueryPausedFilter = Annotated[
    FilterParam[bool | None],
    Depends(filter_param_factory(DagModel.is_paused, bool | None, filter_name="paused")),
]
QueryHasImportErrorsFilter = Annotated[
    FilterParam[bool | None],
    Depends(
        filter_param_factory(
            DagModel.has_import_errors,
            bool | None,
            filter_name="has_import_errors",
            description="Filter Dags by having import errors. Only Dags that have been successfully loaded before will be returned.",
        )
    ),
]
QueryFavoriteFilter = Annotated[_FavoriteFilter, Depends(_FavoriteFilter.depends)]
QueryExcludeStaleFilter = Annotated[_ExcludeStaleFilter, Depends(_ExcludeStaleFilter.depends)]
QueryDagIdPatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(DagModel.dag_id, "dag_id_pattern"))
]
QueryDagIdPrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(DagModel.dag_id, "dag_id_prefix_pattern"))
]
QueryDagDisplayNamePatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(DagModel.dag_display_name, "dag_display_name_pattern"))
]
QueryDagDisplayNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(DagModel.dag_display_name, "dag_display_name_prefix_pattern")),
]
QueryBundleNameFilter = Annotated[
    FilterParam[str | None],
    Depends(filter_param_factory(DagModel.bundle_name, str | None, filter_name="bundle_name")),
]
QueryBundleVersionFilter = Annotated[
    FilterParam[str | None],
    Depends(filter_param_factory(DagModel.bundle_version, str | None, filter_name="bundle_version")),
]
QueryDagIdPatternSearchWithNone = Annotated[
    _SearchParam, Depends(search_param_factory(DagModel.dag_id, "dag_id_pattern", False))
]
QueryDagIdPrefixPatternSearchWithNone = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(DagModel.dag_id, "dag_id_prefix_pattern", False)),
]
QueryTagsFilter = Annotated[_TagsFilter, Depends(_TagsFilter.depends)]
QueryOwnersFilter = Annotated[_OwnersFilter, Depends(_OwnersFilter.depends)]


class _HasAssetScheduleFilter(BaseParam[bool]):
    """Filter Dags that have asset-based scheduling."""

    def to_orm(self, select: Select) -> Select:
        if self.value is None and self.skip_none:
            return select

        asset_ref_subquery = sql_select(DagScheduleAssetReference.dag_id).distinct()

        if self.value:
            # Filter Dags that have asset-based scheduling
            return select.where(DagModel.dag_id.in_(asset_ref_subquery))

        # Filter Dags that do NOT have asset-based scheduling
        return select.where(DagModel.dag_id.notin_(asset_ref_subquery))

    @classmethod
    def depends(
        cls,
        has_asset_schedule: bool | None = Query(None, description="Filter Dags with asset-based scheduling"),
    ) -> _HasAssetScheduleFilter:
        return cls().set_value(has_asset_schedule)


class _AssetDependencyFilter(BaseParam[str]):
    """Filter Dags by specific asset dependencies."""

    def to_orm(self, select: Select) -> Select:
        if self.value is None:
            return select

        escaped = _escape_like_pattern(self.value)
        asset_dag_subquery = (
            sql_select(DagScheduleAssetReference.dag_id)
            .join(AssetModel, DagScheduleAssetReference.asset_id == AssetModel.id)
            .where(
                or_(
                    AssetModel.name.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
                    AssetModel.uri.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
                )
            )
            .distinct()
        )

        return select.where(DagModel.dag_id.in_(asset_dag_subquery))

    @classmethod
    def depends(
        cls,
        asset_dependency: str | None = Query(
            None, description="Filter Dags by asset dependency (name or URI)"
        ),
    ) -> _AssetDependencyFilter:
        return cls().set_value(asset_dependency)


QueryHasAssetScheduleFilter = Annotated[_HasAssetScheduleFilter, Depends(_HasAssetScheduleFilter.depends)]
QueryAssetDependencyFilter = Annotated[_AssetDependencyFilter, Depends(_AssetDependencyFilter.depends)]


class _ConsumingAssetFilter(BaseParam[str | None]):
    """Filter Dag runs by consuming asset (name or URI)."""

    def to_orm(self, select: Select) -> Select:
        if not self.value:
            return select

        escaped = _escape_like_pattern(self.value)
        event_subquery = (
            sql_select(AssetEvent.id)
            .join(AssetModel, AssetEvent.asset_id == AssetModel.id)
            .where(
                or_(
                    AssetModel.name.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
                    AssetModel.uri.ilike(f"%{escaped}%", escape=_LIKE_ESCAPE_CHAR),
                )
            )
            .distinct()
        )

        dagrun_subquery = (
            sql_select(association_table.c.dag_run_id)
            .where(association_table.c.event_id.in_(event_subquery))
            .distinct()
        )

        return select.where(DagRun.id.in_(dagrun_subquery))

    @classmethod
    def depends(
        cls,
        consuming_asset_pattern: str | None = Query(
            None, description="Filter by consuming asset name or URI using pattern matching"
        ),
    ) -> _ConsumingAssetFilter:
        return cls().set_value(consuming_asset_pattern)


QueryConsumingAssetPatternSearch = Annotated[_ConsumingAssetFilter, Depends(_ConsumingAssetFilter.depends)]


class _PendingActionsFilter(BaseParam[bool]):
    """Filter Dags by having pending HITL actions (more than 1)."""

    def to_orm(self, select: Select) -> Select:
        if self.value is None and self.skip_none:
            return select

        from airflow.models.hitl import HITLDetail
        from airflow.models.taskinstance import TaskInstance

        # Join with HITLDetail and TaskInstance to find Dags
        pending_actions_count_subquery = (
            sql_select(func.count(HITLDetail.ti_id))
            .join(TaskInstance, HITLDetail.ti_id == TaskInstance.id)
            .where(
                HITLDetail.responded_at.is_(None),
                TaskInstance.state.in_((TaskInstanceState.DEFERRED, TaskInstanceState.AWAITING_INPUT)),
            )
            .where(TaskInstance.dag_id == DagModel.dag_id)
            .scalar_subquery()
        )

        if self.value is True:
            # Filter to show only Dags with pending actions
            where_clause = pending_actions_count_subquery >= 1
        else:
            # Filter to show only Dags without pending actions
            where_clause = pending_actions_count_subquery == 0

        return select.where(where_clause)

    @classmethod
    def depends(cls, has_pending_actions: bool | None = Query(None)) -> _PendingActionsFilter:
        return cls().set_value(has_pending_actions)


QueryPendingActionsFilter = Annotated[_PendingActionsFilter, Depends(_PendingActionsFilter.depends)]

# DagRun
QueryLastDagRunStateFilter = Annotated[
    FilterParam[DagRunState | None],
    Depends(filter_param_factory(DagRun.state, DagRunState | None, filter_name="last_dag_run_state")),
]


def _transform_dag_run_states(states: Iterable[str] | None) -> list[DagRunState | None] | None:
    try:
        if not states:
            return None
        return [None if s in ("none", None) else DagRunState(s) for s in states]
    except ValueError:
        raise HTTPException(
            status_code=HTTP_422_UNPROCESSABLE_CONTENT,
            detail=f"Invalid value for state. Valid values are {', '.join(DagRunState)}",
        )


QueryDagRunStateFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(
            DagRun.state,
            list[str],
            FilterOptionEnum.ANY_EQUAL,
            default_factory=list,
            transform_callable=_transform_dag_run_states,
        )
    ),
]


def _transform_dag_run_types(types: list[str] | None) -> list[DagRunType | None] | None:
    try:
        if not types:
            return None
        return [None if run_type in ("none", None) else DagRunType(run_type) for run_type in types]
    except ValueError:
        raise HTTPException(
            status_code=HTTP_422_UNPROCESSABLE_CONTENT,
            detail=f"Invalid value for run type. Valid values are {', '.join(DagRunType)}",
        )


QueryDagRunRunTypesFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(
            attribute=DagRun.run_type,
            _type=list[str],
            filter_option=FilterOptionEnum.ANY_EQUAL,
            default_factory=list,
            transform_callable=_transform_dag_run_types,
        )
    ),
]

QueryDagRunTriggeringUserSearch = Annotated[
    _SearchParam, Depends(search_param_factory(DagRun.triggering_user_name, "triggering_user"))
]
QueryDagRunTriggeringUserPrefixSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(DagRun.triggering_user_name, "triggering_user_prefix")),
]
QueryDagRunPartitionKeySearch = Annotated[
    _SearchParam,
    Depends(search_param_factory(DagRun.partition_key, "partition_key_pattern", pipe_as_or=False)),
]
QueryDagRunPartitionKeyPrefixSearch = Annotated[
    _PrefixSearchParam,
    Depends(
        prefix_search_param_factory(DagRun.partition_key, "partition_key_prefix_pattern", pipe_as_or=False)
    ),
]

# DagTags
QueryDagTagPatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(DagTag.name, "tag_name_pattern"))
]
QueryDagTagPrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(DagTag.name, "tag_name_prefix_pattern"))
]


# TI
def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | None] | None:
    """Transform a list of state strings into a list of TaskInstanceState enums handling special 'None' cases."""
    if not states:
        return None

    try:
        return [None if s in ("no_status", "none", None) else TaskInstanceState(s) for s in states]
    except ValueError:
        raise HTTPException(
            status_code=HTTP_422_UNPROCESSABLE_CONTENT,
            detail=f"Invalid value for state. Valid values are {', '.join(TaskInstanceState)}",
        )


QueryTIStateFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(
            TaskInstance.state,
            list[str],
            FilterOptionEnum.ANY_EQUAL,
            default_factory=list,
            transform_callable=_transform_ti_states,
        )
    ),
]
QueryTIPoolFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(TaskInstance.pool, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list)
    ),
]
QueryTIQueueFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(TaskInstance.queue, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list)
    ),
]
QueryTIPoolNamePatternSearch = Annotated[
    _SearchParam,
    Depends(search_param_factory(TaskInstance.pool, "pool_name_pattern")),
]
QueryTIPoolNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(TaskInstance.pool, "pool_name_prefix_pattern")),
]

QueryTIQueueNamePatternSearch = Annotated[
    _SearchParam,
    Depends(search_param_factory(TaskInstance.queue, "queue_name_pattern")),
]
QueryTIQueueNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(TaskInstance.queue, "queue_name_prefix_pattern")),
]
QueryTIExecutorFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(
            TaskInstance.executor, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list
        )
    ),
]
QueryTITaskDisplayNamePatternSearch = Annotated[
    _SearchParam,
    Depends(search_param_factory(TaskInstance.task_display_name, "task_display_name_pattern")),
]
QueryTITaskDisplayNamePrefixPatternSearch = Annotated[
    _TaskDisplayNamePrefixPatternParam, Depends(_TaskDisplayNamePrefixPatternParam.depends)
]
QueryTITaskGroupFilter = Annotated[
    QueryTaskInstanceTaskGroupFilter, Depends(QueryTaskInstanceTaskGroupFilter.depends)
]
QueryTIDagVersionFilter = Annotated[
    FilterParam[list[int]],
    Depends(
        filter_param_factory(
            DagVersion.version_number,
            list[int],
            FilterOptionEnum.ANY_EQUAL,
            default_factory=list,
        )
    ),
]
QueryDagRunVersionFilter = Annotated[
    FilterParam[list[int]],
    Depends(
        filter_param_factory(
            DagVersion.version_number,
            list[int],
            FilterOptionEnum.ANY_EQUAL,
            default_factory=list,
            filter_name="dag_version",
        )
    ),
]
QueryTITryNumberFilter = Annotated[
    FilterParam[list[int]],
    Depends(
        filter_param_factory(
            TaskInstance.try_number, list[int], FilterOptionEnum.ANY_EQUAL, default_factory=list
        )
    ),
]

QueryTIOperatorFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(
            TaskInstance.operator, list[str], FilterOptionEnum.ANY_EQUAL, default_factory=list
        )
    ),
]
QueryTIOperatorNamePatternSearch = Annotated[
    _SearchParam,
    Depends(
        search_param_factory(
            TaskInstance.custom_operator_name,
            "operator_name_pattern",
        )
    ),
]
QueryTIOperatorNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(
        prefix_search_param_factory(
            TaskInstance.custom_operator_name,
            "operator_name_prefix_pattern",
        )
    ),
]

QueryTIMapIndexFilter = Annotated[
    FilterParam[list[int]],
    Depends(
        filter_param_factory(
            TaskInstance.map_index, list[int], FilterOptionEnum.ANY_EQUAL, default_factory=list
        )
    ),
]
# On MySQL the CASE expression that backs rendered_map_index mixes a stored
# VARCHAR column (utf8mb4_bin, IMPLICIT) with CAST(map_index AS CHAR)
# (utf8mb4_0900_ai_ci, IMPLICIT), which gives the whole expression NONE
# coercibility.  Comparing it against a bound parameter then fails with
# "Illegal mix of collations".  _MySQLCollate wraps the expression so that
# on MySQL an explicit COLLATE clause is emitted (giving EXPLICIT coercibility);
# on PostgreSQL and SQLite the wrapper is transparent.
_rendered_map_index_collated = _MySQLCollate(
    cast("ColumnElement[Any]", TaskInstance.rendered_map_index), "utf8mb4_0900_ai_ci"
)

QueryTIRenderedMapIndexPatternSearch = Annotated[
    _SearchParam,
    Depends(
        search_param_factory(
            _rendered_map_index_collated,
            "rendered_map_index_pattern",
        )
    ),
]
QueryTIRenderedMapIndexPrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(
        prefix_search_param_factory(
            _rendered_map_index_collated,
            "rendered_map_index_prefix_pattern",
        )
    ),
]

# XCom
QueryXComKeyPatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(XComModel.key, "xcom_key_pattern"))
]
QueryXComKeyPrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(XComModel.key, "xcom_key_prefix_pattern"))
]

QueryXComDagDisplayNamePatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(DagModel.dag_display_name, "dag_display_name_pattern"))
]
QueryXComDagDisplayNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(DagModel.dag_display_name, "dag_display_name_prefix_pattern")),
]
QueryXComRunIdPatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(XComModel.run_id, "run_id_pattern"))
]
QueryXComRunIdPrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(XComModel.run_id, "run_id_prefix_pattern"))
]
QueryXComTaskIdPatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(XComModel.task_id, "task_id_pattern"))
]
QueryXComTaskIdPrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(XComModel.task_id, "task_id_prefix_pattern"))
]

# Assets
QueryAssetNamePatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(AssetModel.name, "name_pattern"))
]
QueryAssetNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(AssetModel.name, "name_prefix_pattern"))
]
QueryUriPatternSearch = Annotated[_SearchParam, Depends(search_param_factory(AssetModel.uri, "uri_pattern"))]
QueryUriPrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(AssetModel.uri, "uri_prefix_pattern"))
]
QueryAssetAliasNamePatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(AssetAliasModel.name, "name_pattern"))
]
QueryAssetAliasNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(AssetAliasModel.name, "name_prefix_pattern"))
]
QueryAssetDagIdPatternSearch = Annotated[
    _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter.depends)
]
QueryPartitionedDagRunHasCreatedDagRunIdFilter = Annotated[
    FilterParam[bool | None],
    Depends(
        filter_param_factory(
            AssetPartitionDagRun.created_dag_run_id,
            bool | None,
            FilterOptionEnum.IS_NONE,
            filter_name="has_created_dag_run_id",
            transform_callable=lambda v: not v if v is not None else None,
        )
    ),
]
QueryPartitionedDagRunDagIdFilter = Annotated[
    FilterParam[str | None],
    Depends(
        filter_param_factory(
            AssetPartitionDagRun.target_dag_id,
            str | None,
            filter_name="dag_id",
        )
    ),
]

# Variables
QueryVariableKeyPatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(Variable.key, "variable_key_pattern"))
]
QueryVariableKeyPrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(Variable.key, "variable_key_prefix_pattern")),
]

# Pools
QueryPoolNamePatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(Pool.pool, "pool_name_pattern"))
]
QueryPoolNamePrefixPatternSearch = Annotated[
    _PrefixSearchParam, Depends(prefix_search_param_factory(Pool.pool, "pool_name_prefix_pattern"))
]


# UI Shared
def _optional_boolean(value: bool | None) -> bool | None:
    return value if value is not None else False


QueryIncludeUpstream = Annotated[bool, AfterValidator(_optional_boolean)]
QueryIncludeDownstream = Annotated[bool, AfterValidator(_optional_boolean)]

state_priority: list[None | TaskInstanceState] = [
    TaskInstanceState.FAILED,
    TaskInstanceState.UPSTREAM_FAILED,
    TaskInstanceState.UP_FOR_RETRY,
    TaskInstanceState.UP_FOR_RESCHEDULE,
    TaskInstanceState.RUNNING,
    TaskInstanceState.RESTARTING,
    TaskInstanceState.DEFERRED,
    TaskInstanceState.AWAITING_INPUT,
    TaskInstanceState.QUEUED,
    TaskInstanceState.SCHEDULED,
    None,
    TaskInstanceState.SUCCESS,
    TaskInstanceState.SKIPPED,
    TaskInstanceState.REMOVED,
]

# Connections
QueryConnectionIdPatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(Connection.conn_id, "connection_id_pattern"))
]
QueryConnectionIdPrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(Connection.conn_id, "connection_id_prefix_pattern")),
]

# Human in the loop
QueryHITLDetailDagIdPatternSearch = Annotated[
    _SearchParam,
    Depends(
        search_param_factory(
            TaskInstance.dag_id,
            "dag_id_pattern",
        )
    ),
]
QueryHITLDetailDagIdPrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(
        prefix_search_param_factory(
            TaskInstance.dag_id,
            "dag_id_prefix_pattern",
        )
    ),
]
QueryHITLDetailTaskIdPatternSearch = Annotated[
    _SearchParam,
    Depends(
        search_param_factory(
            TaskInstance.task_id,
            "task_id_pattern",
        )
    ),
]
QueryHITLDetailTaskIdPrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(
        prefix_search_param_factory(
            TaskInstance.task_id,
            "task_id_prefix_pattern",
        )
    ),
]
QueryHITLDetailTaskIdFilter = Annotated[
    FilterParam[str | None],
    Depends(
        filter_param_factory(
            TaskInstance.task_id,
            str | None,
            filter_name="task_id",
        )
    ),
]
QueryHITLDetailMapIndexFilter = Annotated[
    FilterParam[int | None],
    Depends(
        filter_param_factory(
            TaskInstance.map_index,
            int | None,
            filter_name="map_index",
        )
    ),
]
QueryHITLDetailSubjectSearch = Annotated[
    _SearchParam,
    Depends(
        search_param_factory(
            HITLDetail.subject,
            "subject_search",
        )
    ),
]
QueryHITLDetailBodySearch = Annotated[
    _SearchParam,
    Depends(
        search_param_factory(
            HITLDetail.body,
            "body_search",
        )
    ),
]
QueryHITLDetailResponseReceivedFilter = Annotated[
    FilterParam[bool | None],
    Depends(
        filter_param_factory(
            HITLDetail.response_received,
            bool | None,
            filter_name="response_received",
        )
    ),
]
QueryHITLDetailRespondedUserIdFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(
            HITLDetail.responded_by_user_id,
            list[str],
            FilterOptionEnum.ANY_EQUAL,
            default_factory=list,
            filter_name="responded_by_user_id",
        )
    ),
]
QueryHITLDetailRespondedUserNameFilter = Annotated[
    FilterParam[list[str]],
    Depends(
        filter_param_factory(
            HITLDetail.responded_by_user_name,
            list[str],
            FilterOptionEnum.ANY_EQUAL,
            default_factory=list,
            filter_name="responded_by_user_name",
        )
    ),
]

# Parse Import Errors
QueryParseImportErrorFilenamePatternSearch = Annotated[
    _SearchParam, Depends(search_param_factory(ParseImportError.filename, "filename_pattern"))
]
QueryParseImportErrorFilenamePrefixPatternSearch = Annotated[
    _PrefixSearchParam,
    Depends(prefix_search_param_factory(ParseImportError.filename, "filename_prefix_pattern")),
]
QueryParseImportErrorFilenameFilter = Annotated[
    FilterParam,
    Depends(
        filter_param_factory(
            ParseImportError.filename,
            str | None,
            filter_name="filename",
            description="Exact filename match. Returns only the import error for this specific file path.",
        )
    ),
]
QueryParseImportErrorBundleNameFilter = Annotated[
    FilterParam,
    Depends(
        filter_param_factory(
            ParseImportError.bundle_name,
            str | None,
            filter_name="bundle_name",
            description="Exact bundle name match. Returns only import errors from this specific bundle.",
        )
    ),
]
