import math
import typing as t

from sqlalchemy import Select, and_, func, inspect, or_, select, tuple_
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase, Mapper, Session
from sqlalchemy.orm.interfaces import ORMOption

from ash_dal.constants import PAGINATOR_FIRST_PAGE_INDEX
from ash_dal.utils.paginator.paginator_page import CursorPaginatorPage

ModelT = t.TypeVar("ModelT", bound=DeclarativeBase)


class _BaseCursorPaginator(t.Generic[ModelT]):
    """
    Cursor paginator optimized for streaming large result sets with optional eager loading.

    This implementation uses a two-step loop per batch:
    1) scan only primary keys via keyset pagination (`WHERE pk > last_seen_pk ORDER BY pk LIMIT n`)
    2) hydrate full ORM rows for that PK batch with configured load options

    Why this strategy:
    - It keeps memory bounded by `page_size` instead of total row count
    - It avoids OFFSET pagination drift and cost on large tables
    - It works around SQLAlchemy limitations where `yield_per` is incompatible with
      eager loaders that require uniquing/row buffering
    """
    _size: int | None = None

    def __init__(
        self,
        query: Select[t.Any],
        model: type[ModelT],
        page_size: int,
        load_options: t.Sequence[ORMOption] = (),
    ) -> None:
        self._query = query
        self._model = model
        self._page_size = page_size
        self._load_options = load_options

    def _get_pk_scan_stmt(self) -> tuple[Select[t.Any], tuple[t.Any, ...]]:
        mapper = t.cast("Mapper[t.Any]", inspect(self._model))
        pk_columns = tuple(mapper.primary_key)
        # Build a PK-only scan from the original query:
        # - keep existing filters/joins from the incoming query
        # - drop existing ORDER BY to avoid conflicting keyset order
        # - dedupe PKs because joins can produce repeated PK rows
        # - enforce deterministic PK ordering for keyset traversal
        pk_stmt = self._query.with_only_columns(*pk_columns).order_by(None).distinct().order_by(*pk_columns)
        return pk_stmt, pk_columns

    def _get_batch_load_stmt(self, pk_columns: tuple[t.Any, ...], pk_batch: list[t.Any]) -> Select[t.Any]:
        # Single PK uses `col IN (...)`; composite PK uses tuple comparison.
        where_clause = pk_columns[0].in_(pk_batch) if len(pk_columns) == 1 else tuple_(*pk_columns).in_(pk_batch)
        return select(self._model).where(where_clause).options(*self._load_options)

    @staticmethod
    def _apply_keyset_after_pk(
        pk_stmt: Select[t.Any],
        pk_columns: tuple[t.Any, ...],
        last_seen_pk: tuple[t.Any, ...] | None,
    ) -> Select[t.Any]:
        # First batch has no cursor boundary.
        if last_seen_pk is None:
            return pk_stmt
        # Fast path for single-column PK.
        if len(pk_columns) == 1:
            return pk_stmt.where(pk_columns[0] > last_seen_pk[0])

        # Composite keyset condition:
        # (k1 > v1) OR (k1 == v1 AND k2 > v2) OR ...
        disjunction_clauses: list[t.Any] = []
        for idx, column in enumerate(pk_columns):
            equality_prefix = [pk_columns[prefix_idx] == last_seen_pk[prefix_idx] for prefix_idx in range(idx)]
            disjunction_clauses.append(and_(*equality_prefix, column > last_seen_pk[idx]))
        return pk_stmt.where(or_(*disjunction_clauses))

    @staticmethod
    def _rows_to_pk_values(rows: t.Sequence[t.Any], pk_columns: tuple[t.Any, ...]) -> list[t.Any]:
        # Convert SQLAlchemy row tuples into PK values used by `IN (...)` hydration.
        raw_pk_values = [row[0] for row in rows] if len(pk_columns) == 1 else [tuple(row) for row in rows]
        # Preserve first-seen order while deduplicating repeated PK rows.
        return list(dict.fromkeys(raw_pk_values))

    @staticmethod
    def _normalize_pk_value(pk_value: t.Any, pk_columns: tuple[t.Any, ...]) -> tuple[t.Any, ...]:
        # Keyset comparison logic expects a tuple shape for both single and composite PKs.
        return (pk_value,) if len(pk_columns) == 1 else t.cast("tuple[t.Any, ...]", pk_value)

    @staticmethod
    def _index_db_items_by_pk(
        db_items: t.Sequence[ModelT],
        pk_columns: tuple[t.Any, ...],
    ) -> dict[t.Any, ModelT]:
        # Index hydrated rows by PK so we can emit them in PK batch order.
        if len(pk_columns) == 1:
            return {getattr(db_item, pk_columns[0].key): db_item for db_item in db_items}
        return {tuple(getattr(db_item, column.key) for column in pk_columns): db_item for db_item in db_items}

    @staticmethod
    def _order_items_by_pk_batch(
        pk_batch: list[t.Any],
        items_by_pk: dict[t.Any, ModelT],
    ) -> tuple[ModelT, ...]:
        # `IN (...)` does not guarantee return order; restore deterministic PK order.
        # Missing keys are ignored defensively in case rows disappear mid-stream.
        return tuple(items_by_pk[pk_value] for pk_value in pk_batch if pk_value in items_by_pk)


class SyncCursorPaginator(_BaseCursorPaginator[ModelT]):
    def __init__(
        self,
        session: Session,
        query: Select[t.Any],
        model: type[ModelT],
        page_size: int,
        load_options: t.Sequence[ORMOption] = (),
    ) -> None:
        super().__init__(query=query, model=model, page_size=page_size, load_options=load_options)
        self._session = session

    def paginate(self) -> t.Iterator[CursorPaginatorPage[ModelT]]:
        pk_stmt, pk_columns = self._get_pk_scan_stmt()
        last_seen_pk: tuple[t.Any, ...] | None = None
        page_index = PAGINATOR_FIRST_PAGE_INDEX
        while True:
            # Fetch the next PK window strictly after the last emitted PK tuple.
            current_pk_stmt = self._apply_keyset_after_pk(
                pk_stmt=pk_stmt, pk_columns=pk_columns, last_seen_pk=last_seen_pk
            )
            pk_rows = self._session.execute(current_pk_stmt.limit(self._page_size)).all()
            if not pk_rows:
                break

            # Hydrate full ORM rows for this batch with eager-load options enabled.
            pk_batch = self._rows_to_pk_values(rows=pk_rows, pk_columns=pk_columns)
            db_items = (
                self._session
                .scalars(self._get_batch_load_stmt(pk_columns=pk_columns, pk_batch=pk_batch))
                .unique()
                .all()
            )
            # Reorder hydrated rows to match PK batch order.
            items_by_pk = self._index_db_items_by_pk(db_items=db_items, pk_columns=pk_columns)
            page_items = self._order_items_by_pk_batch(pk_batch=pk_batch, items_by_pk=items_by_pk)
            yield CursorPaginatorPage(index=page_index, items=page_items)
            # Advance cursor to the last emitted PK tuple.
            last_seen_pk = self._normalize_pk_value(pk_value=pk_batch[-1], pk_columns=pk_columns)
            page_index += 1

    @property
    def size(self) -> int:
        if self._size is None:
            stmt = select(func.count()).select_from(self._query.subquery())
            items_count: int = self._session.scalar(stmt) or 0
            self._size = math.ceil(items_count / self._page_size)
        return self._size


class AsyncCursorPaginator(_BaseCursorPaginator[ModelT]):
    def __init__(
        self,
        session: AsyncSession,
        query: Select[t.Any],
        model: type[ModelT],
        page_size: int,
        load_options: t.Sequence[ORMOption] = (),
    ) -> None:
        super().__init__(query=query, model=model, page_size=page_size, load_options=load_options)
        self._session = session

    async def paginate(self) -> t.AsyncIterator[CursorPaginatorPage[ModelT]]:
        pk_stmt, pk_columns = self._get_pk_scan_stmt()
        last_seen_pk: tuple[t.Any, ...] | None = None
        page_index = PAGINATOR_FIRST_PAGE_INDEX
        while True:
            # Fetch the next PK window strictly after the last emitted PK tuple.
            current_pk_stmt = self._apply_keyset_after_pk(
                pk_stmt=pk_stmt, pk_columns=pk_columns, last_seen_pk=last_seen_pk
            )
            pk_rows = (await self._session.execute(current_pk_stmt.limit(self._page_size))).all()
            if not pk_rows:
                break

            # Hydrate full ORM rows for this batch with eager-load options enabled.
            pk_batch = self._rows_to_pk_values(rows=pk_rows, pk_columns=pk_columns)
            db_items = (
                (await self._session.scalars(self._get_batch_load_stmt(pk_columns=pk_columns, pk_batch=pk_batch)))
                .unique()
                .all()
            )
            # Reorder hydrated rows to match PK batch order.
            items_by_pk = self._index_db_items_by_pk(db_items=db_items, pk_columns=pk_columns)
            page_items = self._order_items_by_pk_batch(pk_batch=pk_batch, items_by_pk=items_by_pk)
            yield CursorPaginatorPage(index=page_index, items=page_items)
            # Advance cursor to the last emitted PK tuple.
            last_seen_pk = self._normalize_pk_value(pk_value=pk_batch[-1], pk_columns=pk_columns)
            page_index += 1

    @property
    async def size(self) -> int:
        if self._size is None:
            stmt = select(func.count()).select_from(self._query.subquery())
            items_count: int = await self._session.scalar(stmt) or 0
            self._size = math.ceil(items_count / self._page_size)
        return self._size
