from typing import Dict, List, Union

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, scoped_session, sessionmaker

from src.config import Config
from src.config import config as _config
from src.constants import DB_CREDENTIAL_KEYS, DB_CREDENTIALS, DB_ENGINE_PREFIX, DBType
from src.utils.common import get_db_type
from src.utils.db_type_dispatcher import db_type_dispatch


def get_db_link_prepared_for_credentials(database_type: DBType = DBType.MYSQL) -> str:
    """Get dialect[+driver] + prepared string for DB credentials to be passed based on DB Type

    The string form of the URL is
        dialect:  is a database name such as ``mysql``, ``postgresql``, etc.,
        driver: the name of a DBAPI, such as ``psycopg2``, ``pymysql``, etc.
        credentials:

    Args:
        database_type: one of [DBType.MYSQL, DBType.PSQL]

    Returns:
        dialect[+driver]://{user}:{password}@{host}:{port}/{database_name}?binary_prefix=true

    """
    return f"{DB_ENGINE_PREFIX[database_type]}{DB_CREDENTIALS[database_type]}"


def get_db_config_key(name: str, slave: bool = False) -> str:
    """Generate a key to get proper value from Config based on a name

    If not slave:
        port -> DB_PORT
        host -> DB_HOST
        password -> DB_PASS
        database_name -> DB_NAME

    If slave:
        port -> SLAVE_DB_PORT
        host -> SLAVE_DB_HOST
        password -> SLAVE_DB_PASS
        database_name -> SLAVE_DB_NAME

    Args:
        name: key name to check and convert to a Config key
        slave: is used to point at read replica if one exists

    Returns:

    """
    index = name.find("_")
    return ("DB_" if not slave else "DB_SLAVE_") + name[index + 1 : index + 5].upper()


def get_db_config_attr(config: Config, name: str, slave: bool = False) -> Union[int, str]:
    """GET attr from config by name

    Args:
        config: instance of a Config class (Configuration File)
        name: key name to be converted to a Config key to get proper Config value
        slave: is used to point at read replica if one exists

    Returns:

    """
    return config.__getattribute__(get_db_config_key(name, slave))


def get_db_credentials(
    config: Config, slave: bool = False, db_credential_keys: List[str] = DB_CREDENTIAL_KEYS
) -> Dict[str, str]:
    """Get credentials for a prepared string to be filled with"""
    return {key: get_db_config_attr(config, key, slave) for key in db_credential_keys}


def get_db_link_with_credentials(config: Config, db_type: DBType = DBType.MYSQL, slave: bool = False):
    """Get a fully configured string for a db engine to be created

        Args:
        config: instance of a Config class
        db_type: one of [DBType.MYSQL, DBType.PSQL]
        slave: is used to point at read replica if one exists

    Returns:

    """
    credentials = get_db_credentials(config, slave)

    return get_db_link_prepared_for_credentials(db_type).format(**credentials)


def get_engine(config: Config, db_type: DBType = DBType.MYSQL, slave: bool = False) -> Engine:
    """Get a specified database engine

    Args:
        config: instance of a Config class
        db_type: one of [DBType.MYSQL, DBType.PSQL]
        slave: is used to point at read replica if one exists

    Returns:
        sqlalchemy.engine.Engine
    """
    return create_engine(
        get_db_link_with_credentials(config, db_type, slave),
        pool_recycle=config.DB_POOL_RECYCLE,
        pool_size=config.DB_POOL_SIZE,
    )


class RoutingSession(Session):
    """Redirects SELECT statements to a read replica of a database if the one exists"""

    def get_bind(self, mapper=None, clause=None, *args, **kwargs):
        if clause is not None and clause.is_selectable and _config.ALLOW_DB_REPLICATION:
            return get_engine(_config, get_db_type(_config.DB_TYPE), True)
        return get_engine(_config, get_db_type(_config.DB_TYPE))


def get_base_sessionmaker(*args, **kwargs):
    """Get base sessionmaker
    Is used for MYSQL
    """
    return sessionmaker(class_=RoutingSession)


def get_psql_sessionmaker(*args, **kwargs):
    """Get scoped session for PSQL"""
    return scoped_session(get_base_sessionmaker())


@db_type_dispatch((DBType.MYSQL, get_base_sessionmaker), (DBType.PSQL, get_psql_sessionmaker))
def get_sessionmaker(db_type: DBType):
    raise NotImplementedError("Unknown DB_TYPE")


def get_session(db_type: DBType = DBType.MYSQL):
    """Create a session to a specific database

    Args:
        db_type: one of [DBType.MYSQL, DBType.PSQL]
    """
    session_maker = get_sessionmaker(db_type=db_type)
    return session_maker()