from sqlalchemy import String, select, ForeignKey, event
from sqlalchemy.orm import (
    declarative_base, sessionmaker, mapped_column, Mapped, relationship, with_polymorphic,
    joinedload
)
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession

Base = declarative_base()


class TopABC(Base):
    __tablename__ = 'top_abc'

    id: Mapped[str] = mapped_column(String(), primary_key=True)
    name: Mapped[str] = mapped_column(String())

    # top: Mapped['Top'] = relationship(lambda: with_polymorphic(Top, '*', flat=True))
    top_mid: Mapped['TopMid'] = relationship(lambda: TopMid, uselist=False, innerjoin=True, back_populates='top_abc')


class TopMid(Base):
    __tablename__ = 'top_mid'

    id: Mapped[str] = mapped_column(String(), ForeignKey(TopABC.id), primary_key=True)

    top_abc: Mapped[TopABC] = relationship(
        TopABC, lazy='joined', uselist=False, innerjoin=True, back_populates='top_mid'
    )
    top: Mapped[Top] = relationship(lambda: Top, lazy='joined', uselist=False, innerjoin=True, back_populates='top_mid')


class Top(Base):
    __tablename__ = 'top'

    id: Mapped[str] = mapped_column(String(), ForeignKey(TopMid.id), primary_key=True)
    type: Mapped[str] = mapped_column(String())

    top_mid: Mapped[TopMid] = relationship(TopMid, lazy='joined', uselist=False, innerjoin=True, back_populates='top')
    __mapper_args__ = dict(
        polymorphic_on=type
    )


class Foo(Top):
    __tablename__ = 'foo'

    id: Mapped[str] = mapped_column(String(), ForeignKey(Top.id), primary_key=True)

    foo_name: Mapped[str] = mapped_column(String())

    __mapper_args__ = dict(
        polymorphic_identity='FOO'
    )


class Bar(Top):
    __tablename__ = 'bar'

    id: Mapped[str] = mapped_column(String(), ForeignKey(Top.id), primary_key=True)
    bar_name: Mapped[str] = mapped_column(String())

    __mapper_args__ = dict(
        polymorphic_identity='BAR'
    )


async_engine = create_async_engine("sqlite+aiosqlite://", pool_pre_ping=True)

@event.listens_for(async_engine.sync_engine, "connect")
def enable_sqlite_fks(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()
    print("Foreign keys enabled for connection.")

AsyncSessionFactory = sessionmaker(
    bind=async_engine, class_=AsyncSession, expire_on_commit=True, autoflush=False
)


async def create_tables() -> None:
    async with async_engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)


def test():

    top_poly = with_polymorphic(Top, '*', flat=True)
    q_bug = (
        select(TopABC)
        .options(
            joinedload(TopABC.top_mid)
            .options(
                joinedload(TopMid.top.of_type(top_poly))
            )
        )
    )
    q_ok = (
        select(TopABC)
        .options(
            joinedload(TopABC.top_mid)
            .joinedload(TopMid.top.of_type(top_poly))
        )
    )
    print(f'{q_bug=!s}')
    print()
    print(f'{q_ok=!s}')


async def main():
    await create_tables()
    test()


import asyncio

asyncio.run(main())