"""
immediateload() silently falls back to lazyload in these cases:
  A) As a level-2 (nested) sub-option, even with non-polymorphic entities:
       selectinload(Owner.animals).options(immediateload(Animal.meta))
       -> meta is lazyloaded on access instead of during execute()
  B) At level 1, when the queried entity is a with_polymorphic():
       select(with_polymorphic(Animal, '*')).options(immediateload(ap.meta))
       -> meta is lazyloaded on access
  In all cases joinedload, selectinload, and subqueryload work correctly.
Is this using documented patterns?
  YES. The immediateload() docstring says: "This function is part of the
  :class:`_orm.Load` interface and supports both method-chained and
  standalone operation." Chaining is expected to work.
  For the polymorphic case: there IS a known limitation documented as a
  TODO in test/orm/inheritance/test_assorted_poly.py:2721:
    "immediateload (and lazyload) do not support the target item
     being a with_polymorphic"
  But this is not documented in the public API docs.
  The chaining failure (case A) does NOT appear to be documented anywhere.
Workaround:
  Use selectinload instead of immediateload (the immediateload docstring
  itself recommends this: "superseded in general by selectinload").
Tested on SQLAlchemy 2.1.0b2 (sqlite in-memory)
"""

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


class Base(DeclarativeBase):
    pass


class Meta(Base):
    __tablename__ = "meta"
    id: Mapped[int] = mapped_column(primary_key=True)
    name: Mapped[str] = mapped_column(String(50))


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))
    meta_id: Mapped[Optional[int]] = mapped_column(ForeignKey("meta.id"))
    meta: Mapped[Optional["Meta"]] = relationship()
    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:
    m1 = Meta(id=1, name="alpha")
    m2 = Meta(id=2, name="beta")
    s.add_all([m1, m2])
    s.flush()
    o = Owner(id=1, name="Alice")
    s.add(o)
    s.flush()
    s.add_all([
        Dog(id=1, name="Rex", breed="Lab", meta=m1, owner=o),
        Cat(id=2, name="Whiskers", color="orange", meta=m2, owner=o),
    ])
    s.commit()

queries_during_access = []


def track(conn, clauseelement, multiparams, params, execution_options):
    queries_during_access.append(str(clauseelement))


event.listen(engine, "before_execute", track)


def test(label, stmt, accessor):
    queries_during_access.clear()
    with Session(engine) as s:
        objs = s.execute(stmt).unique().scalars().all()
        queries_during_access.clear()
        for obj in objs:
            accessor(obj)
        n = len(queries_during_access)
    status = "PASS" if n == 0 else f"FAIL ({n} lazy loads during access)"
    print(f"  {label}: {status}")


loaders = [
    ("joinedload", joinedload),
    ("selectinload", selectinload),
    ("subqueryload", subqueryload),
    ("immediateload", immediateload),
]

# -------------------------------------------------------------------------
print("=" * 70)
print("CASE A: immediateload as level-2 sub-option (non-polymorphic too)")
print("  Owner.animals -> Animal.meta")
print("=" * 70)
print()
print("--- via .options() ---")
for name, fn in loaders:
    for name2, fn2 in loaders:
        test(
            f"{name} + {name2}",
            select(Owner).options(fn(Owner.animals).options(fn2(Animal.meta))),
            lambda o: [a.meta for a in o.animals],
        )

print()
print("--- via direct chain ---")
for name, fn in loaders:
    for name2, fn2 in loaders:
        test(
            f"{name}.{name2}",
            select(Owner).options(getattr(fn(Owner.animals), name2)(Animal.meta)),
            lambda o: [a.meta for a in o.animals],
        )

# -------------------------------------------------------------------------
print()
print("=" * 70)
print("CASE B: immediateload at level 1 on with_polymorphic entity")
print("  select(poly_animal).options(loader(poly.meta))")
print("=" * 70)
print()
ap = with_polymorphic(Animal, "*", flat=True)
for name, fn in loaders:
    test(
        f"{name}(poly.meta)",
        select(ap).options(fn(ap.meta)),
        lambda o: o.meta,
    )

# -------------------------------------------------------------------------
print()
print("=" * 70)
print("CASE C: immediateload at level 1 on plain entity (should work)")
print("  select(Owner).options(loader(Owner.animals))")
print("=" * 70)
print()
for name, fn in loaders:
    test(
        f"{name}(Owner.animals)",
        select(Owner).options(fn(Owner.animals)),
        lambda o: list(o.animals),
    )