"""DB Connector.

Manages interactions with Relational Database (MySQL and SQLite).
"""
from contextlib import contextmanager
from functools import wraps, partial
from time import sleep
from sys import exit as sysexit

from integration_scripts import logger
from integration_scripts.connectors import sentry
from integration_scripts.constants import mysql as mysql_consts

from owsresponse import response

from sqlalchemy import create_engine
from sqlalchemy import exc
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import scoped_session
from sqlalchemy.orm import sessionmaker

import integration_scripts.config as orch_config

from .pidguard import add_engine_pidguard

AR_MYSQL_USER = orch_config.AR_MYSQL_USER
AR_MYSQL_PASSWORD = orch_config.AR_MYSQL_PASSWORD
AR_MYSQL_HOST = orch_config.AR_MYSQL_HOST
AR_MYSQL_DB = orch_config.AR_MYSQL_DB

ENVIRONMENT = orch_config.ENVIRONMENT
TEST_ENVIRONMENT = orch_config.TEST_ENVIRONMENT

env_present = all(x not in ['', None] for x in [
    AR_MYSQL_USER,
    AR_MYSQL_PASSWORD,
    AR_MYSQL_HOST,
    AR_MYSQL_DB
])

if not env_present and ENVIRONMENT != TEST_ENVIRONMENT:
    sysexit(
        'AR_MYSQL_USER, AR_MYSQL_PASSWORD, AR_MYSQL_HOST, AR_MYSQL_DB must not'
        ' be blank.\nTarget environment: {}'.format(ENVIRONMENT))


def _create_engine(db_url):
    """Create engine based on configuration settings."""
    engine = create_engine(
        db_url, poolclass=orch_config.POOL_CLASS, pool_pre_ping=True)

    add_engine_pidguard(engine)
    return engine


# DEBUG - SPIT OUT THE HOSTNAME
if orch_config.DEBUG_LOG:
    logger.debug('Using Host: {}', AR_MYSQL_HOST)


# Do not use these variables directly other than running unit tests
ar_db_engine = _create_engine(mysql_consts.AR_DB_URL)
rds_db_engine = _create_engine(mysql_consts.RDS_DB_URL)

# please don't use sessions directly;
# instead use db_session
sessions = {
    mysql_consts.ART_RELATIONS:
        scoped_session(sessionmaker(bind=ar_db_engine)),
    mysql_consts.RDS:
        scoped_session(sessionmaker(bind=rds_db_engine)),
}


BaseModel = declarative_base()
ARModel = declarative_base()
RDSModel = declarative_base()


@contextmanager
def db_session(db_name, read_only=False):
    """Provide a transactional scope around a series of operations.

    Taken from http://docs.sqlalchemy.org/en/latest/orm/session_basics.html.
    This handles rollback and closing of session, so there is no need
    to do that throughout the code.

    Args:
        db_name (session): Name of database to connect to
        read_only (bool): Determines if commit is performed.

    Usage:
        with db_session() as session:
            session.execute(query)
    """
    session = sessions[db_name]()
    try:
        yield session
        if not read_only:
            session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


def _db_session_wrap(function, hostname):
    """DB Session Wrapper.

    Creates a new session if one isn't passed in.
    """
    @wraps(function)
    def wrapper(*args, **kwargs):
        session = kwargs.pop('session', None)
        if session:
            return function(*args, session=session, **kwargs)
        else:
            with db_session(hostname) as session:
                return function(*args, session=session, **kwargs)

    return wrapper


def _db_session_retry_wrap(function, hostname):
    """DB Session Wrapper with retry.

    Creates a new session if one isn't passed in, and does 1 retry if an
    internal error occurs while executing query. db_session_wrap is still
    preferred since this can rollback the session.
    """
    @wraps(function)
    def wrapper(*args, **kwargs):
        session = kwargs.pop('session', None)
        try:
            if session:
                return function(*args, session=session, **kwargs)
            else:
                with db_session(hostname) as session:
                    return function(*args, session=session, **kwargs)
        except exc.InternalError as e:
            logger.warning(e)
            sleep(mysql_consts.INTERNAL_ERROR_SLEEP_TIME)

            if session:
                session.rollback()
                return function(*args, session=session, **kwargs)
            else:
                with db_session(hostname) as session:
                    return function(*args, session=session, **kwargs)

    return wrapper


ar_db_session_wrap = partial(
    _db_session_wrap, hostname=mysql_consts.ART_RELATIONS)

rds_db_session_wrap = partial(
    _db_session_wrap, hostname=mysql_consts.RDS)


def wrap_db_errors(function):
    """Decorate the given function with logic to handle SQLAlchemy errors.

    If a SQLAlchemy exception is thrown, it will be caught and logged and the
    function will return a fatal response.

    Args:
        function (func): the function to decorate

    Returns:
        func: function decorated with error-handling logic
    """
    @wraps(function)
    def call_function_with_error_handling(*args, **kwargs):
        try:
            function_return = function(*args, **kwargs)
        except exc.SQLAlchemyError as exception:
            sentry.sentry_client.captureMessage(exception, stack=True)
            return response.create_fatal_response()

        return function_return
    return call_function_with_error_handling


def invalidate_pool(db_engine):
    """Dispose of the current pool. Useful for multiprocessing."""
    db_engine.dispose()