"""
MySQL Connector.

Manages interactions with MySQL
for art_relations and direct_delivery database.
"""

from contextlib import contextmanager
from typing import Generator

from sqlalchemy import Engine, create_engine, text
from sqlalchemy.orm import (
    DeclarativeMeta,
    Session,
    declarative_base,
    scoped_session,
    sessionmaker,
)

from assets import config
from assets.api import app


def _create_engine(db_url: str, engine_name: str) -> Engine:
    """Create engine based on config settings."""
    db_engine = create_engine(
        url=db_url,
        pool_logging_name=engine_name,
        **config.CONNECT_ARGS,
    )

    return db_engine


ar_db_engine = _create_engine(config.AR_DB_URL, "ar-db")
ar_session_factory = sessionmaker(bind=ar_db_engine)
ar_database_session = scoped_session(ar_session_factory)

au_db_engine = _create_engine(config.RDS_DB_URL, "au-db")
au_session_factory = sessionmaker(bind=au_db_engine)
au_database_session = scoped_session(au_session_factory)
sessions = (ar_database_session, au_database_session)

ArModel: DeclarativeMeta = declarative_base()

AuModel: DeclarativeMeta = declarative_base()


@app.teardown_request
def remove_session(error: object) -> None:
    """
    For scoped_session, explicitly close session at end of each request.

    Args:
        error (object): When a teardown function was called because of an
            exception it will be passed an error object.
    """
    for session in sessions:
        session.remove()


@contextmanager
def ar_db_session(
    read_only: bool = False, turn_off_foreign_key_checks: bool = False
) -> Generator[Session, None, None]:
    """
    Provide transactional scope around series of operations for art_relations.

    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 ar_database_session() as session:
            session.execute(query)
    """
    session = ar_database_session()
    try:
        if turn_off_foreign_key_checks:
            session.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
        yield session
        if not read_only:
            session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        if turn_off_foreign_key_checks:
            session.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))

    # cleanup by closing session
    session.close()


@contextmanager
def au_db_session(
    read_only: bool = False,
    turn_off_foreign_key_checks: bool = False,
) -> Generator[Session, None, None]:
    """
    Provide a transactional scope around a series of operations.

    for asset_upload database
    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 au_database_session() as session:
            session.execute(query)
    """
    session = au_database_session()
    try:
        if turn_off_foreign_key_checks:
            session.execute(text("SET FOREIGN_KEY_CHECKS = 0;"))
        yield session
        if not read_only:
            session.commit()
    except Exception:
        session.rollback()
        raise
    finally:
        if turn_off_foreign_key_checks:
            session.execute(text("SET FOREIGN_KEY_CHECKS = 1;"))

    # cleanup by closing session
    session.close()


@contextmanager
def db_lock(
    session: Session, lock_name: str, lock_timeout: int
) -> Generator[int, None, None]:
    """
    Obtain a lock and release it on context exit.

        session (sqlalchemy.orm.Session): current db session
        lock_name (str): unique id of the lock
        lock_timeout (int): seconds to wait for lock to release
    """
    thread_id_on_get_lock = session.connection().connection.thread_id()
    lock = session.execute(
        text(
            "SELECT GET_LOCK(:name, :timeout)",
        ),
        {
            "name": lock_name,
            "timeout": lock_timeout,
        },
    ).scalar()
    if not lock:
        raise Exception(f"Unable to acquire lock {lock_name}")
    try:
        yield lock
    except Exception:
        raise
    finally:
        thread_id_on_release_lock = session.connection().connection.thread_id()
        if thread_id_on_get_lock != thread_id_on_release_lock:
            raise Exception("thread id on release != on lock")
        session.execute(
            text(
                "SELECT RELEASE_LOCK(:name)",
            ),
            {
                "name": lock_name,
            },
        )