from sqlalchemy import String, select, ForeignKey
from sqlalchemy.orm import (
    declarative_base, sessionmaker, mapped_column, Mapped, relationship, with_polymorphic,
    joinedload
)
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
# from src.db.base_class import Base

Base = declarative_base()


class Meta(Base):
    __tablename__ = 'meta'
    id: Mapped[str] = mapped_column(String(), primary_key=True)
    name: Mapped[str] = mapped_column(String())


class Top(Base):
    __tablename__ = 'top'

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

    meta: Mapped[Meta] = relationship(Meta)

    __mapper_args__ = dict(
        polymorphic_on=type
    )


class Foo(Top):
    __tablename__ = 'foo'

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

    __mapper_args__ = dict(
        polymorphic_identity='FOO'
    )


class Bar(Top):
    __tablename__ = 'bar'

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

    foo_id: Mapped[str] = mapped_column(String(), ForeignKey(Foo.id))

    foo: Mapped[Meta] = relationship(Foo, foreign_keys=[foo_id])

    __mapper_args__ = dict(
        polymorphic_identity='BAR'
    )


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

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)


async def test():

    async with AsyncSessionFactory() as sess:
        sess.bind.echo = True
        meta = Meta()
        meta.id = '1'
        meta.name = 'some'

        top = Foo(id='1', meta_id='1')
        sess.add(top)

        bar = Bar(id='2', meta_id='1', foo_id='1')
        sess.add(bar)

        try:
            await sess.commit()
        except Exception as exc:
            print(exc)

        tp: Top = with_polymorphic(Top, '*', flat=True)

        q = (
            select(tp)
        ).options(
            joinedload(tp.Bar.foo)
           .joinedload(Foo.meta)  #  has no effect - see sql-echo below
        )
        orms = (await sess.execute(q)).scalars().all()
        print(orms)


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


import asyncio

asyncio.run(main())