"""DB Connector.

Manages interactions with Relational Database (MySQL and SQLite).
"""

from contextlib import contextmanager
from functools import wraps
import os
from time import sleep

from flask import g
from oto import response
from sentry_sdk import capture_exception

from sqlalchemy import create_engine
from sqlalchemy import event
from sqlalchemy import exc
from sqlalchemy import pool
from sqlalchemy import select
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

from backend import config
from backend.constants import mysql as mysql_consts


def _create_engine(db_url):
    """Create engine based on configuration settings."""
    if config.POOL_CLASS != pool.QueuePool:
        return create_engine(db_url, poolclass=config.POOL_CLASS)
    result = create_engine(
        db_url, pool_size=config.POOL_SIZE,
        max_overflow=config.POOL_MAX_OVERFLOW,
        pool_recycle=config.POOL_RECYCLE_MS)

    # SQLAlchemy 1.2 supports pessimistic disconnect handling out of box.
    # Once version 1.2 is out of beta, it is recommended to use
    # the `pool_pre_ping` param for `create_engine` and remove any code
    # in this file related to `_ping_connection`
    # See http://docs.sqlalchemy.org/en/latest/core/pooling.html
    if config.POOL_PRE_PING:
        event.listen(result, 'engine_connect', _ping_connection)
    return result


def _ping_connection(connection, branch):
    """Ping database connection after engine_connect event.

    This function is copied verbatim from
    http://docs.sqlalchemy.org/en/latest/core/pooling.html
    """
    if branch:
        # "branch" refers to a sub-connection of a connection,
        # we don't want to bother pinging on these.
        return

    # turn off "close with result".  This flag is only used with
    # "connectionless" execution, otherwise will be False in any case
    save_should_close_with_result = connection.should_close_with_result
    connection.should_close_with_result = False

    try:
        # run a SELECT 1.   use a core select() so that
        # the SELECT of a scalar value without a table is
        # appropriately formatted for the backend
        connection.scalar(select([1]))
    except exc.DBAPIError as err:
        # catch SQLAlchemy's DBAPIError, which is a wrapper
        # for the DBAPI's exception.  It includes a .connection_invalidated
        # attribute which specifies if this connection is a "disconnect"
        # condition, which is based on inspection of the original exception
        # by the dialect in use.
        if err.connection_invalidated:
            # run the same SELECT again - the connection will re-validate
            # itself and establish a new connection.  The disconnect detection
            # here also causes the whole connection pool to be invalidated
            # so that all stale connections are discarded.
            connection.scalar(select([1]))
        else:
            raise
    finally:
        # restore "close with result"
        connection.should_close_with_result = save_should_close_with_result


# Do not use these variables directly other than running unit tests
# This engine connects to art_relations DB
db_engine = _create_engine(config.DB_URL)

# This engine connects to local service DB
ows_track_db_engine = _create_engine(config.OWS_TRACK_DB_URL)

db_engines = (db_engine, ows_track_db_engine)

# Fix for multiprocessing on pools
# Ensure connections are disposed
try:
    from uwsgidecorators import postfork
except ImportError:
    # not in uUWSGI context, no need for postfork event
    pass
else:
    @postfork
    def dispose():
        """Dispose existing connection pool on engine."""
        for db_engine_item in db_engines:
            db_engine_item.dispose()


# please don't use sessions directly;
# instead use db_session
sessions = {
    mysql_consts.ART_RELATIONS: sessionmaker(bind=db_engine),
    mysql_consts.OWS_TRACK: sessionmaker(bind=ows_track_db_engine)
}


BaseModel = declarative_base()
OwsTrackBaseModel = declarative_base()


@contextmanager
def db_session(db_name=mysql_consts.ART_RELATIONS):
    """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
    Usage:
        with db_session() as session:
            session.execute(query)
    """
    session = sessions[db_name]()
    try:
        yield session
        session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def ows_track_db_session():
    """Transactional scope for operations with ows-track own DB."""
    with db_session(mysql_consts.OWS_TRACK) as session:
        yield session


def db_session_wrap(function):
    """DB Session Wrappper.

    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() as session:
                return function(*args, session=session, **kwargs)

    return wrapper


def ows_track_db_session_wrap(function):
    """ows-track DB Session Wrappper.

    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 ows_track_db_session() as session:
                return function(*args, session=session, **kwargs)

    return wrapper


def db_session_retry_wrap(function):
    """DB Session Wrappper 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() as session:
                    return function(*args, session=session, **kwargs)
        except exc.InternalError as e:
            # Log warning and sleep for a bit
            g.log.warning(e)
            sleep(mysql_consts.INTERNAL_ERROR_SLEEP_TIME)

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

    return wrapper


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:
            capture_exception(exception)
            return response.create_fatal_response()

        return function_return
    return call_function_with_error_handling


def connect(dbapi_connection, connection_record):
    """Set connection record PID."""
    connection_record.info['pid'] = os.getpid()


def checkout(dbapi_connection, connection_record, connection_proxy):
    """Ensure current PID matches connection record PID."""
    pid = os.getpid()
    if connection_record.info['pid'] != pid:
        connection_record.connection = connection_proxy.connection = None
        raise exc.DisconnectionError(
            'Connection record belongs to pid %s, '
            'attempting to check out in pid %s' %
            (connection_record.info['pid'], pid)
        )


for db_engine_item in db_engines:
    event.listen(db_engine_item, 'connect', connect)
    event.listen(db_engine_item, 'checkout', checkout)