from typing import cast, reveal_type

from sqlalchemy import Select, select, Integer, Text
from sqlalchemy.engine import Row
from sqlalchemy.orm.query import Query, RowReturningQuery
from sqlalchemy.orm import Session
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy import create_engine


class Base(DeclarativeBase):
    pass


class User(Base):
    id: Mapped[int] = mapped_column(Integer, primary_key=True)
    phone: Mapped[str] = mapped_column(Text)


engine = create_engine("sqlite://")
session = Session(engine)

stmt = select(User.id, User.phone)
# mypy: Select[tuple[int, str]]
#   ty: Select[tuple[Unknown, Unknown]]
reveal_type(stmt)

query = session.query(User.id, User.phone)
# mypy: RowReturningQuery[tuple[int, str]]
#   ty: RowReturningQuery[tuple[Unknown, Unknown]]
reveal_type(query)

for row in session.execute(stmt):
    # mypy: Row[Tuple[int, str]]
    #   ty: Row[tuple[Unknown, Unknown]]
    reveal_type(row)
for row in query.all():
    # mypy: Row[Tuple[int, str]]
    #   ty: Row[tuple[Unknown, Unknown]]
    reveal_type(row)  

user_stmt = select(User)
# mypy: Select[tuple[User]]
#   ty: Select[tuple[Unknown]]
reveal_type(user_stmt)

users = session.scalars(select(User)).all()
# mypy: Sequence[User]
#   ty: Sequence[Unknown]
reveal_type(users)

users_legacy = session.query(User).all()
# mypy: list[User]
#   ty: list[Unknown]
reveal_type(users_legacy)


# None of these give type errors for mypy/ty
def get_user_id_and_phone_base_stmt() -> Select[tuple[int, str]]:
    return select(User.id, User.phone)


def get_user_id_with_only_columns() -> Select[tuple[int]]:
    return get_user_id_and_phone_base_stmt().with_only_columns(User.id)


def get_user_id_and_phone() -> Row[tuple[int, str]]:
    return session.execute(get_user_id_and_phone_base_stmt()).one()


def get_user_stmt() -> Select[tuple[User]]:
    return select(User)


def get_user() -> User:
    return session.scalars(get_user_stmt()).one()


def get_user_id_and_phone_query() -> RowReturningQuery[tuple[int, str]]:
    return session.query(User.id, User.phone)


def get_user_id_with_entities() -> RowReturningQuery[tuple[int]]:
    query = get_user_id_and_phone_query().with_entities(User.id)
    # Not handled well in either, legacy pattern
    # mypy: Query[Any]
    #   ty: Query[Unknown]
    reveal_type(query)
    return cast(RowReturningQuery[tuple[int]], query)


def get_user_query() -> Query[User]:
    return session.query(User)


def get_user_from_query() -> User:
    return get_user_query().one()


# Checking that Row works if we manually specify it
def get_user_id_and_phone_query_all() -> list[Row[tuple[int, str]]]:
    return get_user_id_and_phone_query().all()


rows = get_user_id_and_phone_query_all()
reveal_type(rows[0])  # both: Row[tuple[int, str]]
reveal_type(rows[0].phone)  # both: Any
reveal_type(rows[0][0])  # both: Any
reveal_type(rows[0]._t)  # both: tuple[int, str]