"""Keyset (seek) WHERE-clause construction for efficient deep pagination.

Unlike OFFSET-based pagination, keyset pagination appends a WHERE clause
derived from the cursor — the unique/PK column value of the last-seen row.
For a column ``ORDER BY id ASC`` with a cursor at ``id=5`` the generated
clause is::

    id > 5

This module handles:

* ``extract_column`` — normalise a plain attribute or ``.asc()``/``.desc()``
  expression into ``(column, is_desc)`` pairs.
* ``build_keyset_filter`` — produce the WHERE predicate for a cursor position.
* ``reverse_order_by`` — flip the sort direction for backward pagination.

Type alias ``ColWithDir`` represents a ``(QueryableAttribute, is_desc)`` pair.
"""

from collections.abc import Sequence
from typing import Any, TypeAlias, cast

from sqlalchemy import and_, or_
from sqlalchemy.orm.attributes import QueryableAttribute
from sqlalchemy.sql import operators as sa_operators
from sqlalchemy.sql.elements import ColumnElement, UnaryExpression

from sqlalchemy_relay_pagination.exceptions import InvalidArgumentsError

# A column paired with its sort direction: True → descending.
ColWithDir: TypeAlias = tuple[QueryableAttribute[Any], bool]


def extract_column(expr: QueryableAttribute[Any] | UnaryExpression) -> ColWithDir:
    """Normalise an ORDER BY expression into a ``(column, is_desc)`` pair.

    Args:
        expr: A bare ``Model.col`` attribute **or** a ``Model.col.asc()`` /
            ``Model.col.desc()`` unary expression.

    Returns:
        ``(QueryableAttribute, is_desc)`` where ``is_desc`` is ``True`` for
        descending order.

    Raises:
        InvalidArgumentsError: For unsupported expression types.
    """
    if isinstance(expr, QueryableAttribute):
        return expr, False
    if isinstance(expr, UnaryExpression):
        return cast("QueryableAttribute[Any]", expr.element), expr.modifier is sa_operators.desc_op
    raise InvalidArgumentsError(
        f"Unsupported ORDER BY expression type: {type(expr)}. Use column attributes or .asc()/.desc() expressions."
    )


def build_keyset_filter(
    col: QueryableAttribute[Any],
    cursor_value: Any,
    is_desc: bool,
    is_forward: bool,
) -> ColumnElement[bool]:
    """Return a WHERE predicate that pages past the given cursor value.

    Direction rules:

    +----------+---------+------+
    | forward  | is_desc | op   |
    +==========+=========+======+
    | True     | False   | >    |
    +----------+---------+------+
    | True     | True    | <    |
    +----------+---------+------+
    | False    | False   | <    |
    +----------+---------+------+
    | False    | True    | >    |
    +----------+---------+------+

    Args:
        col: The unique/PK column attribute.
        cursor_value: The decoded cursor value.
        is_desc: ``True`` if the column is sorted descending.
        is_forward: ``True`` for forward (``first``/``after``) pagination.

    Returns:
        A SQLAlchemy ``ColumnElement[bool]`` ready to be passed to ``.filter()``.
    """
    use_gt = is_forward != is_desc
    return col > cursor_value if use_gt else col < cursor_value  # type: ignore[return-value]


def reverse_order_by(col: QueryableAttribute[Any], is_desc: bool) -> UnaryExpression:
    """Return the column expression with its sort direction flipped.

    Used to query the rows immediately *before* the cursor when paginating
    backward; the caller reverses the resulting row list to restore order.
    """
    return col.asc() if is_desc else col.desc()


def normalise_order_by(
    order_by: QueryableAttribute[Any] | UnaryExpression | Sequence[Any],  # type: ignore[type-arg]
) -> tuple[list[ColWithDir], list[Any]]:
    """Normalise a single or multi-column ORDER BY spec.

    Returns:
        ``(cols_with_dir, raw_exprs)`` where ``raw_exprs`` are the original
        expressions suitable for ``query.order_by(*raw_exprs)`` (preserving
        ``.asc()``/``.desc()`` modifiers), and ``cols_with_dir`` is a list of
        ``(QueryableAttribute, is_desc)`` pairs.
    """
    if isinstance(order_by, (list, tuple)):
        cols = [extract_column(e) for e in order_by]
        return cols, list(order_by)
    col_with_dir = extract_column(order_by)  # type: ignore[arg-type]
    return [col_with_dir], [order_by]


def reverse_order_by_list(cols: list[ColWithDir]) -> list[UnaryExpression]:
    """Return each column with its sort direction flipped (for backward pagination)."""
    return [reverse_order_by(col, is_desc) for col, is_desc in cols]


def build_compound_keyset_filter(
    cols: Sequence[ColWithDir],
    cursor_values: Sequence[Any],
    is_forward: bool,
) -> ColumnElement[bool]:
    """Return a staircase OR predicate for compound keyset pagination.

    For ``ORDER BY a ASC, b DESC`` forward, generates::

        (a > va) OR (a = va AND b < vb)

    The operator for each column follows the same rule as single-column:
    ``use_gt = is_forward != is_desc``.
    """
    clauses = []
    for i, ((col, is_desc), val) in enumerate(zip(cols, cursor_values, strict=True)):
        use_gt = is_forward != is_desc
        strict = col > val if use_gt else col < val  # type: ignore[operator]
        if i == 0:
            clauses.append(strict)
        else:
            eq_prefix = and_(*(cols[j][0] == cursor_values[j] for j in range(i)))
            clauses.append(and_(eq_prefix, strict))
    return or_(*clauses)  # type: ignore[return-value]
