import inspect
import subprocess
from itertools import chain

import sqlalchemy
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.engine import Engine
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy.exc import ProgrammingError
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.schema import CreateColumn
from sqlalchemy.sql.functions import GenericFunction
from sqlalchemy.types import Text

from puddl import get_config
from puddl.conf import DBConfig


def _get_caller_app():
    # We need to go down the stack two times:
    #
    # - we are called by something (our caller)
    # - someone else called our caller - that's what we want
    frm = inspect.stack()[2]
    mod_name = inspect.getmodule(frm[0]).__name__
    conf = get_config()
    return conf.get_app_from_module_name(mod_name)


def get_engine(conf=None, **kwargs):
    """
    Convenience function to get an engine based on puddl config.
    """

    if conf is None:
        conf = get_config()

    return sqlalchemy.create_engine(conf.db_url, **kwargs)


def get_alchemy_session(conf=None, **kwargs):
    """
    Convenience. Don't know if this is wise. :>
    """
    engine = get_engine(conf, **kwargs)
    session_cls = sessionmaker(bind=engine)
    return session_cls()


# 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


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)


class Schema:
    """
    This has two usages: in models for declarative_base and interactively (CLI,
    notebook, ...).

    In models, you care about the `declarative_base` and some meta info like `name`.
    You do not have an engine, so you cannot supply a bind. In that case, just don't::

        >>> schema = Schema('foo')

    When you use the schema interactively, you want to operate against the database.
    You must have a bind, so provide it. The schema gets created automatically and
    switches the connection to use it:

        >>> engine = ...
        >>> schema = Schema('foo', bind=engine)

    You can then use the engine like this:

        >>> schema.engine.execute('SELECT 1').first()
        (1,)
    """
    engine = None
    Session = None

    def __init__(self, schema_name, bind=None):
        self.name = schema_name
        if bind is not None:
            self.bind(bind)

    def declarative_base(self):
        class Base:
            __table_args__ = {"schema": self.name}

        return declarative_base(cls=Base)

    def bind(self, engine):
        self.engine = engine
        self.engine.execute(f'CREATE SCHEMA IF NOT EXISTS {self.name}')
        self.engine.execute(f'SET search_path TO {self.name}')
        self.Session = sessionmaker(bind=self.engine)


# https://docs.sqlalchemy.org/en/13/core/functions.html
# noinspection PyPep8Naming
class puddl_upsert_role(GenericFunction):
    type = Text
    package = 'puddl'
    identifier = 'upsert_role'


class App:
    """
    >>> app = App('sunrise')

    This creates a Postgres ROLE and SCHEMA with the name "sunrise".

    Please instantiate this class once per process, because it uses SQLAlchemy's engine,
    thus similar rules apply. "[...] the Engine is most efficient when created just once
    at the module level of an application, not per-object or per-function call"
    -- https://docs.sqlalchemy.org/en/13/core/connections.html#basic-usage

    An app holds DB configuration and an engine:

    >>> app.db_config
    DBConfig('sunrise')
    >>> app.engine
    Engine(postgresql://sunrise:***@localhost:13370/puddl?application_name=sunrise)
    """
    name: str
    db_config: DBConfig
    engine: Engine

    def __init__(self, name):
        self.name = name

        # get global credentials (to upsert the app itself)
        puddl_root_db_config = DBConfig()
        my_password = f'{self.name}-{puddl_root_db_config.PGPASSWORD}'
        root_engine = sqlalchemy.create_engine(puddl_root_db_config.url)
        session_cls = sessionmaker(root_engine)
        session = session_cls()
        try:
            session.execute(f"CREATE SCHEMA IF NOT EXISTS {name}")
            session.execute(sqlalchemy.func.puddl.upsert_role(name, my_password))
            session.execute(f"GRANT ALL ON SCHEMA {name} TO {name}")
            # puddl may write to any scope
            session.execute(f"GRANT ALL ON SCHEMA {name} TO puddl")
            session.commit()
        except ProgrammingError:
            session.rollback()
            raise
        self.db_config = DBConfig(self.name)
        self.engine = sqlalchemy.create_engine(self.db_config.url)

    def has_table(self, table_name):
        return self.engine.dialect.has_table(self.engine, table_name)

    def df_dump(self, df, table_name, index=True):
        """
        Dumps a DataFrame to table_name.
        WARNING! This truncates the table if it exists.
        """
        if self.has_table(table_name):
            self.engine.execute(f'TRUNCATE TABLE {table_name}')
        if index:
            # follow SQL conventions for the index
            df.index += 1
        return df.to_sql(
            table_name, self.engine, self.name, if_exists='replace',
            index=index
        )

    def psql(self):
        db_env = self.db_config.to_dict()
        subprocess.Popen('psql', env=db_env).wait()
