"""Batched child-resource pagination using window functions.

``paginate_many`` resolves paginated children for multiple parent keys in a
single SQL query, using ``ROW_NUMBER() OVER (PARTITION BY fk_col)`` and
optionally ``COUNT(*) OVER (PARTITION BY fk_col)`` to avoid N+1 queries.

This is the natural batch-load function for a dataloader pattern.

When per-key ``after``/``before`` cursors are supplied, the function falls back
to individual ``paginate()`` calls (N queries instead of one), which is
correct but loses the batching efficiency.
"""

from collections.abc import Mapping, Sequence
from typing import Any, TypeVar

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import aliased
from sqlalchemy.sql.selectable import Select

from sqlalchemy_relay_pagination.config import PaginationConfig, get_config
from sqlalchemy_relay_pagination.exceptions import InvalidArgumentsError
from sqlalchemy_relay_pagination.keyset import normalise_order_by
from sqlalchemy_relay_pagination.pagination import OrderByExpr, _assemble_connection, paginate
from sqlalchemy_relay_pagination.types import Connection, PaginationArgs, RequestedFields

T = TypeVar("T")
K = TypeVar("K")


def _assert_clean_query(query: Select) -> None:
    """Raise ``InvalidArgumentsError`` if the query has ORDER BY, LIMIT, or OFFSET.

    ``paginate_many`` builds its own subquery structure around the supplied
    ``Select``; an existing ORDER BY or LIMIT would be silently carried into
    the inner subquery and corrupt the window-function results.
    """
    if getattr(query, "_order_by_clauses", None):
        raise InvalidArgumentsError(
            "query must not contain an ORDER BY clause — paginate_many applies its own window ordering"
        )
    if getattr(query, "_limit_clause", None) is not None:
        raise InvalidArgumentsError(
            "query must not contain a LIMIT clause — paginate_many applies its own row limit via ROW_NUMBER"
        )
    if getattr(query, "_offset_clause", None) is not None:
        raise InvalidArgumentsError(
            "query must not contain an OFFSET clause — paginate_many applies its own row limit via ROW_NUMBER"
        )


async def paginate_many(
    session: AsyncSession,
    query: Select[tuple[T]],
    model: type[T],
    partition_by: Any,
    order_by: OrderByExpr,
    keys: Sequence[K],
    first: int | None = None,
    last: int | None = None,
    after: Mapping[K, str | None] | None = None,
    before: Mapping[K, str | None] | None = None,
    requested_fields: RequestedFields | None = None,
    config: PaginationConfig | None = None,
) -> dict[K, Connection[T]]:
    """Paginate child resources for multiple parent keys in a single SQL query.

    Uses ``ROW_NUMBER() OVER (PARTITION BY fk_col)`` to fetch paginated results
    for all parent keys at once, making it suitable for dataloader batch loading.

    When ``after`` or ``before`` are provided the function falls back to
    individual :func:`paginate` calls (one query per key) because each parent
    may be at a different cursor position.  Results are still correct; only
    the single-query optimisation is lost.

    Args:
        session: An open ``AsyncSession``.
        query: A ``select(Model)`` statement — must **not** have ORDER BY,
            LIMIT, or OFFSET (those are applied internally).
        model: The ORM model class (used for ``aliased()`` remapping).
        partition_by: ORM column attribute for grouping, e.g. ``Post.user_id``.
        order_by: Within-partition sort expression, e.g. ``Post.id`` or
            ``Post.id.desc()``.
        keys: Parent IDs to fetch children for.
        first: Return first *n* edges per group (forward pagination).
        last: Return last *n* edges per group (backward pagination).
        after: Per-key opaque cursor strings (forward).  Triggers N individual
            queries instead of a single batched query.  ``None`` values in
            the mapping are treated as "no cursor" for that key.
        before: Per-key opaque cursor strings (backward).  Triggers N
            individual queries instead of a single batched query.  ``None``
            values in the mapping are treated as "no cursor" for that key.
        requested_fields: Declare which fields are needed for query optimization.
        config: Per-call config override. Falls back to the global config.

    Returns:
        A dict mapping each key to its ``Connection``. Keys are in
        caller-supplied insertion order. Keys with no children map to empty
        Connections.

    Raises:
        InvalidArgumentsError: If ``query`` contains ORDER BY, LIMIT, or
            OFFSET; or if the pagination arguments are otherwise invalid.
    """
    _assert_clean_query(query)

    if not keys:
        return {}

    # --- cursor path: fall back to individual paginate() calls (N queries) ---
    if after is not None or before is not None:
        output: dict[K, Connection[T]] = {}
        for key in keys:
            key_after = (after or {}).get(key)
            key_before = (before or {}).get(key)
            key_query = query.filter(partition_by == key)
            conn = await paginate(
                session,
                key_query,
                order_by=order_by,
                first=first,
                last=last,
                after=key_after,
                before=key_before,
                requested_fields=requested_fields,
                config=config,
            )
            output[key] = conn
        return output

    # --- batched window-function path ---
    cfg = config or get_config()
    args = PaginationArgs(first=first, last=last, config=config)  # type: ignore[arg-type]
    fields = requested_fields or RequestedFields()
    cols, _ = normalise_order_by(order_by)

    if not fields.needs_page_query and not fields.total_count:
        empty = _assemble_connection(
            rows=None,
            total=None,
            cols=cols,
            codec=cfg.cursor_codec,
            limit=args.effective_limit,
            is_forward=args.is_forward,
            fetch_extra=False,
            fields=fields,
        )
        return dict.fromkeys(keys, empty)

    fetch_extra = (args.is_forward and fields.has_next_page) or (not args.is_forward and fields.has_previous_page)
    fetch_limit = args.effective_limit + 1 if fetch_extra else args.effective_limit

    partition_key = partition_by.key

    # Filter down to only the requested keys
    inner_subq = query.filter(partition_by.in_(keys)).subquery("_base")

    # Window ORDER BY direction:
    #   is_forward=True  → same direction as order_by
    #   is_forward=False → reversed (so _assemble_connection can un-reverse)
    if args.is_forward:
        window_order = [
            inner_subq.c[col.key].desc() if is_desc else inner_subq.c[col.key].asc() for col, is_desc in cols
        ]
    else:
        window_order = [
            inner_subq.c[col.key].asc() if is_desc else inner_subq.c[col.key].desc() for col, is_desc in cols
        ]

    rn_expr = (
        func.row_number()
        .over(
            partition_by=inner_subq.c[partition_key],
            order_by=window_order,
        )
        .label("_rn")
    )

    select_exprs: list[Any] = [inner_subq, rn_expr]
    if fields.total_count:
        # COUNT computed over inner_subq before the _rn filter — always the full partition count
        total_expr = func.count().over(partition_by=inner_subq.c[partition_key]).label("_total")
        select_exprs.append(total_expr)

    ranked_subq = select(*select_exprs).subquery("_ranked")

    # Re-attach ORM mapping so execute() returns model instances
    aliased_model = aliased(model, ranked_subq)

    stmt_cols: list[Any] = [aliased_model, ranked_subq.c["_rn"]]
    if fields.total_count:
        stmt_cols.append(ranked_subq.c["_total"])

    stmt = (
        select(*stmt_cols)
        .where(ranked_subq.c["_rn"] <= fetch_limit)
        .order_by(ranked_subq.c[partition_key], ranked_subq.c["_rn"])
    )

    result = await session.execute(stmt)
    rows_raw = result.all()

    # Group rows by partition key value
    groups: dict[Any, list[Any]] = {}
    for row in rows_raw:
        entity = row[0]
        key_val = getattr(entity, partition_key)
        if key_val not in groups:
            groups[key_val] = []
        groups[key_val].append(row)

    result_out: dict[K, Connection[T]] = {}
    for key in keys:
        group = groups.get(key, [])
        rows = [r[0] for r in group]
        if fields.total_count:
            total: int | None = group[0][2] if group else 0
        else:
            total = None

        conn = _assemble_connection(
            rows=rows,
            total=total,
            cols=cols,
            codec=cfg.cursor_codec,
            limit=args.effective_limit,
            is_forward=args.is_forward,
            fetch_extra=fetch_extra,
            fields=fields,
        )
        result_out[key] = conn  # type: ignore[assignment]

    return result_out  # type: ignore[return-value]
