"""
SQLalchemy helper functions
"""
from pyaddons import misc_outputs
from simple_progress import progress
from sqlalchemy.sql import select
import configparser
import sys
import re
# import os
import sqlalchemy
import sqlalchemy_utils
import pprint as pp


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_sqlalchemy_connection_info(config_file, section, entry_prefix):
    """
    Get SQLAlchemy connection information from a config file section. This
    allows the user to supply all sorts of connection parameters to SQLAlchemy
    via a config file, see:
    https://docs.sqlalchemy.org/en/13/core/engines.html#sqlalchemy.engine_from_config

    Parameters
    ----------
    config : :obj:`configparser`
        Access to the config file
    section : :obj:`str`
        The section containing the database connection info
    entry_prefix : :obj:`str`
        The prefix for the section entries, these will be used in the SQL
        Alchemy connection. If an entry does not start with the prefix is is
        assumed to be an additional DBAPI specific connection option, see:
        https://docs.sqlalchemy.org/en/13/core/engines.html#custom-dbapi-a

    Returns
    -------
    conn_info :`dict`
        The connection info that will be used to connect to an SQL Alchemy
        creat_engine statement

    Raises
    ------
    KeyError
        If the section specifing the connection information does not exist
    """

    # Any DBAPI specific args
    connect_args = {}

    # SQLAlchemy config args
    config_args = {}

    config = config_file
    if not isinstance(config_file, configparser.ConfigParser):
        config = configparser.ConfigParser()
        config.read(config_file)

    # First make sure the section exists, this will raise a KeyError if not
    config[section]

    # Loop through all the entries in the section to build the options. If an
    # entry does not start with the prefix is is assumed to be an additional
    # DBAPI specific connection option, see:
    # https://docs.sqlalchemy.org/en/13/core/engines.html#custom-dbapi-args
    for e, v in config.items(section):
        # If we are prefixed then it is an arg to SQLAlchemy, otherwise it is
        # an arg to the DBAPI
        if e.startswith(entry_prefix):
            config_args[e] = v
        else:
            connect_args[e] = v

    # Make sure we add the connect args if we have them
    if len(connect_args) > 0:
        config_args['connect_args'] = connect_args

    return config_args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def print_orm_obj(obj, exclude=[]):
    """
    prints the contents of the SQL Alchemy ORM object

    This is designed to be called from an SQLAlchemy ORM model object and it
    in the __repr__ method. It simply frovides a formatted string that has the
    data in the object.

    Parameters
    ----------
    obj : :obj:`Base`
        The SQLAlchemy declarative base model object. This should have a
        __table__ attribute that contains the corresponding SQLAlchemy table
        object
    exclude : list, optional
        A list of attributes (these correspond to table columns) that you do
        not want to include in the string

    Returns
    -------
    str
        A formatted string that can be printed. The string contains the
        attributes of the object and the data they contain.
    """

    # Loop through all the columns in the table
    return misc_outputs.obj_repr(obj,
                                 attr=[i.name for i in obj.__table__.columns
                                       if i.name not in exclude])


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def clean_password(connstr):
    """replaces a password in an SQLAlchemy connection string with XXXXX

    This generates an SQL Alchemy connection string that can safely be printed
    on the command line.

    Parameters
    ----------
    connstr : :obj:`str`
        An SQL Alchemy connection string of the format:
        driver://user:pwd@host:port/dbname. If there is no password then no
        substitutions are made and the original string is returned

    Returns
    -------
    str
        A connection string with the password section replaced by XXXXX or the
        original string.
    """

    return re.sub(r'(://.+?):.*?@', r'\1:XXXXX@', connstr)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_database_connection(conn_info, conn_prefix, must_exist=False):
    """
    Connect to the database depending on the database type

    Parameters
    ----------

    conn_info :`str`
        The defaults heading or path without any "|" delimiter or an
        SQLite database path
    must_exist :`bool`
        Only applies to 'sqlite' types. If the database file does not exist and
        this is True then the error launcher is fired. Otherwise it is silently
        created. For MySQL types, it is always an error if the database does
        not exist

    Returns
    -------

    Session : :obj:`Session`
        An SQLAlchemy session object
    """

    # create the database engine, note that this does not connect until it is
    # asked to do something, so we use sqlalchemy_utils to enforce the
    # must_exist
    engine = sqlalchemy.engine_from_config(conn_info, prefix=conn_prefix)

    # This essentially tests the connection before we have done anything
    dbexists = sqlalchemy_utils.database_exists(engine.url)

    # Now if must_exist is True then we use the engine to determine if the
    # database actually exists
    if must_exist is True and dbexists is False:
        raise FileNotFoundError("database must exist and doesn't")

    # If we get here then we are good to continue, so we create a session
    return sqlalchemy.orm.sessionmaker(bind=engine)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def copy_table(from_session, from_table, to_session, to_table=None,
               checkfirst=True, verbose=False, commit_every=100000):
    """
    Copy the contents of a table from one session to another. Note that this
    has only been used to copy databases from engines with the same dialect
    and some copying from SQLite to MySQL. You could have dialect issues when
    creating tables on the target. Especially if you are using reflected
    objects.

    Parameters
    ----------
    from_session : :obj:`Session`
        An SQLAlchemy session object where the table is being copied from
    from_table : :obj:`sqlalchemy.declarative.Base` or `sqlalchemy.Table`
        An ORM class representing the table to be copied or a `Table` class.
        if it is a `Table` class then to_table can't be `NoneType` as we need
        an ORM class to do the bulk updates. This is to cope with the edge
        cases where the `from_table` has no primary keys, so the `to_table`
        orm class should have an autoincrement one set that will be updated
        automatically.
    to_session : :obj:`Session`
        An SQLAlchemy session object where the table is being copied to
    to_table : :obj:`sqlalchemy.Base` or NoneType, optional, default: NoneType
        An ORM class representing the table to be copied to. If `NoneType` then
        `from_table` is the same as `to_table`. See exception when `from_table`
        is a `sqlalchemy.Table` class.
    checkfirst : bool, optional, default: True
        We will attempt to create `to_table`, if check first is `False` and
        `to_table` already exists this will result in an error
    verbose : bool, optional, default: False
        Output copying progress to STDERR
    commit_every : int, optional, default: 100000
        The number of rows to load up before committing to `to_table`
    """
    # The way we gather the table name from the source table will differ
    # depending on if it is an ORM object or a Table object
    from_tablename = None

    # This will hold a function that is used to extract a dict of the data
    # from the results object, this is because the method to do this will
    # vary if from_table is a Table() object as opposed to a ORM object
    # As we will will use SQLAlchemy query expression syntax to query the
    # database in the case of a Table() object and Declarative Base in the
    # case of an ORM object
    data_ext_func = None

    # If we have a Table object
    if isinstance(from_table, sqlalchemy.Table):
        # Make sure to_table id defined
        if to_table is None:
            raise TypeError(
                "if 'from_table' is a Table() to_table can't be None"
            )

        # Set the data extraction function and from table name
        data_ext_func = _get_exp_row_dict
        from_tablename = from_table.name

        # Query all the rows. Note that we are streaming the results and
        # not buffering, this helps save memory if the table is large
        # I am not sure all DB APIs can do this
        s = select([from_table])
        conn = from_session.get_bind().connect()
        sql = conn.execution_options(stream_results=True).execute(s)
    else:
        # We presumably have an ORM object so set the data extraction function
        # table name and issue the query
        data_ext_func = _get_orm_row_dict
        from_tablename = from_table.__tablename__
        sql = from_session.query(from_table)

    # In the case of ORM objects set the to_table if not set already
    if to_table is None:
        to_table = from_table

    # Create the table we are copying to
    to_table.__table__.create(
        to_session.get_bind(),
        checkfirst=checkfirst
    )

    msg = "copying data from {0} to {1}".format(
        from_tablename, to_table.__tablename__
    )
    prog = progress.RateProgress(default_msg=msg, file=sys.stderr,
                                 verbose=verbose)

    row_cache = []
    counter = 0
    # Loop through the query results from the source table
    for row in prog.progress(sql):
        # Extract the data as a dict into the cache that will be used until
        # we want to commit
        row_cache.append(data_ext_func(row))
        counter += 1

        # Have we reached our cache limit, if so, commit and empty
        if counter == commit_every:
            to_session.bulk_insert_mappings(to_table, row_cache)
            to_session.commit()
            row_cache = []
            counter = 0

    # A final commit of data in the cache
    if counter > 1:
        to_session.bulk_insert_mappings(to_table, row_cache)
        to_session.commit()
        row_cache = []
        counter = 0


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _get_orm_row_dict(row):
    """
    Return a dict of the results row when the query has been performed by the
    ORM

    Parameters
    ----------
    row : `sqlalchemy.ResultsObject`
        A results row

    Returns
    -------
    results_dict : dict
        The keys of the dict are column names and the values are data values
        in the row
    """
    return row.__dict__


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _get_exp_row_dict(row):
    """
    Return a dict of the results row when the query has been performed by the
    SQLAlchemy SQL expression API

    Parameters
    ----------
    row : `sqlalchemy.ResultsProxy`
        A results row

    Returns
    -------
    results_dict : dict
        The keys of the dict are column names and the values are data values
        in the row
    """
    return dict(row)
