"""DB Connector.

Manages interactions with database schema.
"""

import functools
from contextlib import contextmanager

from oto import response
from sentry_sdk import capture_exception
from sqlalchemy import create_engine, exc, pool
from sqlalchemy.orm import declarative_base, sessionmaker

from timed_release import config


def get_engine():
    """Create engine based on config settings."""
    if config.POOL_CLASS == pool.QueuePool:
        _db_engine = create_engine(
            config.RDS_DB_URL, pool_size=config.POOL_SIZE,
            max_overflow=config.POOL_MAX_OVERFLOW,
            pool_recycle=config.POOL_RECYCLE_MS,
            pool_pre_ping=config.POOL_PRE_PING, poolclass=config.POOL_CLASS)
        return _db_engine
    return create_engine(config.RDS_DB_URL, poolclass=config.POOL_CLASS)


db_engine = get_engine()
db_session_maker = sessionmaker(bind=db_engine)


def get_ar_db_engine():
    """Get AR DB engine."""
    if config.POOL_CLASS == pool.QueuePool:
        _ar_db_engine = create_engine(
            config.AR_DB_URL, pool_size=config.POOL_SIZE,
            max_overflow=config.POOL_MAX_OVERFLOW,
            pool_recycle=config.POOL_RECYCLE_MS,
            pool_pre_ping=config.POOL_PRE_PING)
        return _ar_db_engine
    return create_engine(config.AR_DB_URL, poolclass=config.POOL_CLASS)


ar_db_engine = get_ar_db_engine()
ar_db_session_maker = sessionmaker(bind=ar_db_engine)

base_model = declarative_base()


@contextmanager
def db_session(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.

    Usage:
        with db_session() as session:
            session.execute(query)
    """
    session = db_session_maker()
    try:
        yield session
        if not read_only:
            session.commit()
    except: # noqa
        session.rollback()
        raise
    finally:
        session.close()


@contextmanager
def ar_db_session():
    """Provide a transactional scope around art_relations db operations.

    with ar_db_session() as session:
        session.execute(query)
    """
    session = ar_db_session_maker()
    try:
        yield session
        session.commit()
    except Exception as err:
        session.rollback()
        raise err
    finally:
        session.close()


def db_session_wrap(func):
    """Wrap shared DB Session.

    Creates a new session if one isn't passed in. This lets us share/pass-in
    a common session across multiple functions, making them all transactional.
    """
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        session = kwargs.pop('session', None)
        if session:
            return func(*args, session=session, **kwargs)
        else:
            with db_session() as session:
                return func(*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
    """
    @functools.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