from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy import ForeignKey
from sqlalchemy import Integer
from sqlalchemy import literal
from sqlalchemy import null
from sqlalchemy import select
from sqlalchemy import String
from sqlalchemy import union
from sqlalchemy.ext.declarative import AbstractConcreteBase
from sqlalchemy.orm import aliased
from sqlalchemy.orm import composite
from sqlalchemy.orm import contains_eager
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import declared_attr
from sqlalchemy.orm import relationship
from sqlalchemy.orm import Session
from sqlalchemy.orm import with_polymorphic


Base = declarative_base()


class Metadata(Base):
    __tablename__ = "metadata"
    id = Column(
        Integer,
        primary_key=True,
        unique=True,
    )
    some_data = Column(String)


def make_statement(*filter_cond, include_metadata=False):

    a_stmt = (
        select(
            A.id,
            A.thing1,
            A.x1,
            A.y1,
            null().label("thing2"),
            null().label("x2"),
            null().label("y2"),
            literal("a").label("type"),
        )
        .join(Metadata)
        .filter(*filter_cond)
    )
    if include_metadata:
        a_stmt = a_stmt.add_columns(
            Metadata.__table__
        )

    b_stmt = (
        select(
            B.id,
            null().label("thing1"),
            null().label("x1"),
            null().label("y1"),
            B.thing2,
            B.x2,
            B.y2,
            literal("b").label("type"),
        )
        .join(Metadata)
        .filter(*filter_cond)
    )
    if include_metadata:
        b_stmt = b_stmt.add_columns(Metadata.__table__)

    return union(a_stmt, b_stmt)


class BaseObj(AbstractConcreteBase, Base):
    @declared_attr
    def id(cls):
        return Column(ForeignKey(Metadata.id), primary_key=True)

    @classmethod
    def _create_polymorphic_union(cls, mappers, discriminator_name):
        return make_statement().subquery()

    @declared_attr
    def related_metadata(cls):
        return relationship(Metadata)


class XYThing:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __composite_values__(self):
        return (self.x, self.y)

    def __repr__(self):
        return f"XYThing({self.x}, {self.y})"


class A(BaseObj):
    __tablename__ = "a"
    thing1 = Column(String)
    comp1 = composite(XYThing, Column("x1", Integer), Column("y1", Integer))

    @declared_attr
    def __mapper_args__(cls):
        return {"polymorphic_identity": cls.__tablename__, "concrete": True}

    def __repr__(self):
        return f"A(id={self.id}, thing1={self.thing1}, comp1={self.comp1})"


class B(BaseObj):
    __tablename__ = "b"
    thing2 = Column(String)
    comp2 = composite(XYThing, Column("x2", Integer), Column("y2", Integer))

    @declared_attr
    def __mapper_args__(cls):
        return {"polymorphic_identity": cls.__tablename__, "concrete": True}

    def __repr__(self):
        return f"B(id={self.id}, thing2={self.thing2}, comp2={self.comp2})"


e = create_engine("postgresql://scott:tiger@localhost/test", echo="debug")
Base.metadata.drop_all(e)
Base.metadata.create_all(e)

with Session(e) as sess:
    sess.add_all([Metadata(id=1, some_data="m1"), Metadata(id=2, some_data="m2")])
    sess.flush()

    sess.add_all(
        [
            A(id=1, thing1="thing1", comp1=XYThing(1, 2)),
            B(id=2, thing2="thing2", comp2=XYThing(3, 4)),
        ]
    )
    sess.commit()


# version one - use selectinload

with Session(e) as sess:

    alias = make_statement(Metadata.id < 3, include_metadata=True).subquery()
    ac = with_polymorphic(
        BaseObj,
        [A, B],
        selectable=alias,
        adapt_on_names=True,
    )

    mt = aliased(Metadata, alias=alias)

    for obj in sess.scalars(
        select(ac).options(
            contains_eager(ac.A.related_metadata.of_type(mt)),
            contains_eager(ac.B.related_metadata.of_type(mt)),
        )
    ):

        print(obj)
        print(obj.related_metadata)