"""MySQL database connector.

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

from contextlib import contextmanager
from functools import wraps
from typing import Optional

from sqlalchemy.orm import declarative_base, sessionmaker
from sqlalchemy.orm.session import Session

from collaborator import config
from collaborator.utils import db as db_utils

_db_engine = db_utils.create_engine(config.COLLABORATOR_DB_URL)
_session_maker = sessionmaker(bind=_db_engine, expire_on_commit=False)

BaseModel = declarative_base()


@contextmanager
def create_session(session: Optional[Session] = None):
    """Mysql session context manager."""
    if session:
        yield session
    else:
        with db_utils.create_session(_session_maker) as session:
            yield session


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

    return wrapper