"""mysql connector."""
from contextlib import contextmanager
from functools import wraps
import time

from flask import g
from sqlalchemy import create_engine
from sqlalchemy import exc
from sqlalchemy import pool
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import sessionmaker

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


def _create_engine(db_url):
    """Create engine based on configuration settings."""
    if config.POOL_CLASS == pool.QueuePool:
        db_engine = create_engine(
            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 db_engine
    else:
        return create_engine(db_url, poolclass=config.POOL_CLASS)


# Do not use these variables directly other than running unit tests
db_engine = _create_engine(config.DB_URL)

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

BaseModel = declarative_base()


@contextmanager
def db_session(db_name=mysql_consts.OWS_VECTOR_JOB_RULES):
    """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()


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

    Creates a new session if one isn't passed in.
    """

    @wraps(fn)
    def wrapper(*args, **kwargs):
        session = kwargs.pop('session', None)
        if session:
            return fn(*args, session=session, **kwargs)
        else:
            with db_session() as session:
                return fn(*args, session=session, **kwargs)

    return wrapper


def db_session_retry_wrap(fn):
    """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(fn)
    def wrapper(*args, **kwargs):
        session = kwargs.pop('session', None)
        try:
            if session:
                return fn(*args, session=session, **kwargs)
            else:
                with db_session() as session:
                    return fn(*args, session=session, **kwargs)
        except exc.InternalError as e:
            # Log warning and sleep for a bit
            g.log.warning(e)
            time.sleep(mysql_consts.INTERNAL_ERROR_SLEEP_TIME)

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

    return wrapper