from sqlalchemy import String, ForeignKey, create_engine, select, event
from sqlalchemy.orm import (
    DeclarativeBase, mapped_column, Mapped, relationship, Session,
    joinedload, selectinload, subqueryload,
)
from typing import Optional, List


class Base(DeclarativeBase):
    pass


class Animal(Base):
    __tablename__ = "animal"
    id: Mapped[int] = mapped_column(primary_key=True)
    type: Mapped[str] = mapped_column(String(50))
    name: Mapped[str] = mapped_column(String(50))
    owner_id: Mapped[Optional[int]] = mapped_column(ForeignKey("owner.id"))
    owner: Mapped[Optional["Owner"]] = relationship(back_populates="animals")
    __mapper_args__ = {"polymorphic_on": "type", "polymorphic_identity": "animal"}


class Dog(Animal):
    __tablename__ = "dog"
    id: Mapped[int] = mapped_column(ForeignKey("animal.id"), primary_key=True)
    breed: Mapped[str] = mapped_column(String(50))
    __mapper_args__ = {"polymorphic_identity": "dog"}


class Cat(Animal):
    __tablename__ = "cat"
    id: Mapped[int] = mapped_column(ForeignKey("animal.id"), primary_key=True)
    color: Mapped[str] = mapped_column(String(50))
    __mapper_args__ = {"polymorphic_identity": "cat"}


class Owner(Base):
    __tablename__ = "owner"
    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(50))
    animals: Mapped[List["Animal"]] = relationship(back_populates="owner")


engine = create_engine("sqlite://", echo=False)
Base.metadata.create_all(engine)

with Session(engine) as s:
    o = Owner(id=1, name="Alice")
    s.add(o)
    s.flush()
    s.add(Dog(id=1, name="Rex", breed="Lab", owner=o))
    s.add(Cat(id=2, name="Whiskers", color="orange", owner=o))
    s.commit()


for name, fn in [
    # ("joinedload", joinedload),
    ("selectinload", selectinload),
    ("subqueryload", subqueryload),
]:
    try:
        with Session(engine) as s:
            stmt = select(Owner).options(
                fn(Owner.animals.of_type(Dog).and_(Dog.breed == "Lab"))
            )
            objs = s.execute(stmt).unique().scalars().all()
            for o in objs:
                for a in o.animals:
                    assert isinstance(a, Dog), f'Unexpected type: {type(a)}'
                    if a.breed != 'Lab':
                        raise Exception(f"Unexpected breed: {a.breed}")

            print(f"  {name}: PASS")
    except Exception as e:
        print(f"  {name}: ERROR — {type(e).__name__}: {e}")

print()
for name, fn in [
    # ("joinedload", joinedload),
    ("selectinload", selectinload),
    ("subqueryload", subqueryload),
]:
    try:
        with Session(engine) as s:
            stmt = select(Owner).options(
                fn(Owner.animals.and_(Animal.name == "Rex"))
            )
            objs = s.execute(stmt).unique().scalars().all()
            for o in objs:
                for a in o.animals:
                    if a.name != 'Rex':
                        raise Exception(f"Unexpected name: {a.name}")
            print(f"  {name}: PASS")
    except Exception as e:
        print(f"  {name}: ERROR — {type(e).__name__}: {e}")