from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy import event
from sqlalchemy import ForeignKey
from sqlalchemy import Integer
from sqlalchemy import String
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.declarative import declared_attr
from sqlalchemy.orm import object_session
from sqlalchemy.orm import relationship
from sqlalchemy.orm import Session
from sqlalchemy.orm import validates

Base = declarative_base()


class Type(Base):
    """our lookup object."""

    __tablename__ = "type"
    id = Column(Integer, primary_key=True)
    name = Column(String, unique=True)


class HasType(object):
    """Define a class that links to a Type object."""

    @declared_attr
    def type_id(cls):
        return Column(ForeignKey("type.id"), nullable=False)

    @declared_attr
    def _type(cls):
        return relationship("Type")

    type = association_proxy(
        "_type", "name", creator=lambda name: Type(name=name)
    )
    """Define <someobject>.type as the string name of its Type object.

    When <someobject>.type is set to a string, a new, anonymous Type() object
    is created with that name and assigned to <someobject>._type.   However it
    does not have a database id.  This will have to be fixed later when the
    object is associated with a Session where we will replace this
    Type() object with the correct one.

    """

    @validates("_type")
    def _validate_type(self, key, value):
        """Receive the event that occurs when <someobject>._type is set.

        If the object is present in a Session, then make sure it's the Type
        object that we looked up from the database.

        Otherwise, do nothing and we'll fix it later when the object is
        put into a Session.

        """
        sess = object_session(self)
        if sess is not None:
            return _setup_type(sess, value)
        else:
            return value


@event.listens_for(Session, "transient_to_pending")
def _validate_type(session, object_):
    """Receive the HasType object when it gets attached to a Session to correct
    its Type object.

    Note that this discards the existing Type object.

    """

    if (
        isinstance(object_, HasType)  # it's a HasType
        and object_._type is not None  # something set object_._type = Type()
        and object_._type.id is None  # and it has no database id
    ):
        # the id-less Type object that got created
        old_type = object_._type
        # make sure it's not going to be persisted.
        if old_type in session:
            session.expunge(old_type)

        object_._type = _setup_type(session, object_._type)


def _setup_type(session, type_object):
    """Given a Session and a Type object, return
    the correct Type object from the database.

    """

    with session.no_autoflush:
        return session.query(Type).filter_by(name=type_object.name).one()


# demonstrate the pattern.


class A(HasType, Base):
    __tablename__ = "a"

    id = Column(Integer, primary_key=True)


e = create_engine("sqlite://", echo=True)
Base.metadata.create_all(e)

s = Session(e)
t1, t2, t3 = Type(name="typea"), Type(name="typeb"), Type(name="typec")
s.add_all([t1, t2, t3])
s.commit()


a1 = A(type="typeb")
a2 = A(type="typec")
s.add_all([a1, a2])
s.commit()

assert a1._type is t2
assert a1.type == "typeb"
assert a2._type is t3
assert a2.type == "typec"