from itertools import chain

from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.schema import CreateColumn

from puddl import get_config


def get_alchemy_session(conf=None, **kwargs):
    """
    Convenience. Don't know if this is wise. :>
    """
    if conf is None:
        conf = get_config()
    import sqlalchemy
    from sqlalchemy.orm import sessionmaker
    engine = sqlalchemy.create_engine(conf.db_url, **kwargs)

    Session = sessionmaker(bind=engine)
    return Session()


# https://docs.sqlalchemy.org/en/13/dialects/postgresql.html#postgresql-10-identity-columns
@compiles(CreateColumn, 'postgresql')
def use_identity(element, compiler, **kw):
    text = compiler.visit_create_column(element, **kw)
    text = text.replace("SERIAL", "INT GENERATED BY DEFAULT AS IDENTITY")
    return text


def make_base(schema):
    class Base:
        __table_args__ = {"schema": schema}

    return declarative_base(cls=Base)


class Upserter:
    def __init__(self, session, model):
        self.session = session
        self.model = model
        inspector = Inspector.from_engine(session.bind)
        _xs = inspector.get_unique_constraints(self.model.__tablename__,
                                               schema=self.model.__table__.schema)
        list_of_lists = [x['column_names'] for x in _xs]
        self.unique_columns = set(chain.from_iterable(list_of_lists))

    def upsert(self, **data):
        # read this carefully:
        # https://docs.sqlalchemy.org/en/13/dialects/postgresql.html#sqlalchemy.dialects.postgresql.dml.Insert.on_conflict_do_update
        columns = set(data.keys())
        update_columns = columns - self.unique_columns
        insert_stmt = insert(self.model).values(**data)
        excluded_mapping = {  # LHS: column name, RHS: Column instance
            col: getattr(insert_stmt.excluded, col)
            for col in update_columns
        }
        do_update_stmt = insert_stmt.on_conflict_do_update(
            index_elements=self.unique_columns,
            set_=excluded_mapping,
        )
        self.session.execute(do_update_stmt)
