from typing import NamedTuple, Optional, Type, TypeVar, cast

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import DeclarativeBase, aliased, selectinload
from sqlalchemy.sql.expression import literal
from sqlalchemy.sql.selectable import Select

# import sqlparse


# def format_sql(sql):
#     return sqlparse.format(str(sql), reindent=True, keyword_case='upper')


class SortField(NamedTuple):
    field: str
    desc: bool = False


Sort = list[SortField]

_Id = TypeVar('_Id')


async def load_models(
    session: AsyncSession,
    Model: Type[DeclarativeBase],
    id_field: str,
    *,
    sort: Optional[Sort] = None,
    before: Optional[_Id] = None,
    after: Optional[_Id] = None,
    first: Optional[int] = None,
    last: Optional[int] = None,
):
    stmt = compose_statement(
        Model,
        id_field,
        sort=sort,
        before=before,
        after=after,
        first=first,
        last=last,
    )

    # TODO: Make this optional
    stmt = stmt.options(selectinload('*'))

    models = await session.scalars(stmt)
    return models


def compose_statement(
    Model: Type[DeclarativeBase],
    id_field: str,
    *,
    sort: Optional[Sort] = None,
    before: Optional[_Id] = None,
    after: Optional[_Id] = None,
    first: Optional[int] = None,
    last: Optional[int] = None,
) -> Select:
    '''Return a SELECT statement object to be given to session.scalars'''

    forward = (after is not None) or (first is not None)
    backward = (before is not None) or (last is not None)

    if forward and backward:
        raise ValueError('Only either after/first or before/last is allowed')

    sort = sort or []

    if id_field not in [s.field for s in sort]:
        sort.append(SortField(id_field))

    def sorting_fields(Model, reverse=False):
        return [
            f.desc() if reverse ^ d else f
            for f, d in [(getattr(Model, s.field), s.desc) for s in sort]
        ]

    if not (forward or backward):
        return select(Model).order_by(*sorting_fields(Model))

    cursor = after if forward else before
    limit = first if forward else last

    if cursor is None:
        stmt = select(Model).order_by(*sorting_fields(Model, reverse=backward))
    else:
        cte = select(
            Model,
            func.row_number()
            .over(order_by=sorting_fields(Model, reverse=backward))
            .label('row_number'),
        ).cte()

        subq = select(cte.c.row_number.label('cursor'))
        subq = subq.where(getattr(cte.c, id_field) == cursor)
        subq = cast(Select[tuple], subq.subquery())

        Alias = aliased(Model, cte)  # type: ignore
        stmt = select(Alias).select_from(cte)
        stmt = stmt.join(subq, literal(True))  # type: ignore # cartesian product
        stmt = stmt.order_by(*sorting_fields(Alias, reverse=backward))
        stmt = stmt.where(cte.c.row_number > subq.c.cursor)

    if limit is not None:
        stmt = stmt.limit(limit)

    if backward:
        Alias = aliased(Model, stmt.subquery())
        stmt = select(Alias).order_by(*sorting_fields(Alias))

    return stmt
