import logging
import os
import time
from contextlib import contextmanager

import sqlalchemy
from alembic.migration import MigrationContext
from alembic.script import ScriptDirectory
from sqlalchemy import sql

# We need to import sqlalchemy.pool to convert poolclass string to class object
from sqlalchemy.pool import (
    AssertionPool,
    AsyncAdaptedQueuePool,
    FallbackAsyncAdaptedQueuePool,
    NullPool,
    QueuePool,
    SingletonThreadPool,
    StaticPool,
)

from mlflow.environment_variables import (
    MLFLOW_MYSQL_SSL_CA,
    MLFLOW_MYSQL_SSL_CERT,
    MLFLOW_MYSQL_SSL_KEY,
    MLFLOW_SQLALCHEMYSTORE_ECHO,
    MLFLOW_SQLALCHEMYSTORE_MAX_OVERFLOW,
    MLFLOW_SQLALCHEMYSTORE_POOL_RECYCLE,
    MLFLOW_SQLALCHEMYSTORE_POOL_SIZE,
    MLFLOW_SQLALCHEMYSTORE_POOLCLASS,
)
from mlflow.exceptions import MlflowException
from mlflow.protos.databricks_pb2 import (
    BAD_REQUEST,
    INTERNAL_ERROR,
    TEMPORARILY_UNAVAILABLE,
)
from mlflow.store.db.db_types import SQLITE
from mlflow.store.model_registry.dbmodels.models import (
    SqlModelVersion,
    SqlModelVersionTag,
    SqlRegisteredModel,
    SqlRegisteredModelAlias,
    SqlRegisteredModelTag,
)
from mlflow.store.tracking.dbmodels.initial_models import Base as InitialBase
from mlflow.store.tracking.dbmodels.models import (
    SqlDataset,
    SqlExperiment,
    SqlExperimentTag,
    SqlInput,
    SqlInputTag,
    SqlLatestMetric,
    SqlMetric,
    SqlParam,
    SqlRun,
    SqlTag,
    SqlTraceInfo,
    SqlTraceRequestMetadata,
    SqlTraceTag,
)

_logger = logging.getLogger(__name__)

MAX_RETRY_COUNT = 10


def _get_package_dir():
    """Returns directory containing MLflow python package."""
    current_dir = os.path.dirname(os.path.abspath(__file__))
    return os.path.normpath(os.path.join(current_dir, os.pardir, os.pardir))


def _all_tables_exist(engine):
    return {
        t
        for t in sqlalchemy.inspect(engine).get_table_names()
        # Filter out alembic tables
        if not t.startswith("alembic_")
    } == {
        SqlExperiment.__tablename__,
        SqlRun.__tablename__,
        SqlMetric.__tablename__,
        SqlParam.__tablename__,
        SqlTag.__tablename__,
        SqlExperimentTag.__tablename__,
        SqlLatestMetric.__tablename__,
        SqlRegisteredModel.__tablename__,
        SqlModelVersion.__tablename__,
        SqlRegisteredModelTag.__tablename__,
        SqlModelVersionTag.__tablename__,
        SqlRegisteredModelAlias.__tablename__,
        SqlDataset.__tablename__,
        SqlInput.__tablename__,
        SqlInputTag.__tablename__,
        SqlTraceInfo.__tablename__,
        SqlTraceTag.__tablename__,
        SqlTraceRequestMetadata.__tablename__,
    }


def _initialize_tables(engine):
    _logger.info("Creating initial MLflow database tables...")
    InitialBase.metadata.create_all(engine)
    _upgrade_db(engine)


def _get_latest_schema_revision():
    """Get latest schema revision as a string."""
    # We aren't executing any commands against a DB, so we leave the DB URL unspecified
    config = _get_alembic_config(db_url="")
    script = ScriptDirectory.from_config(config)
    heads = script.get_heads()
    if len(heads) != 1:
        raise MlflowException(
            f"Migration script directory was in unexpected state. Got {len(heads)} head "
            f"database versions but expected only 1. Found versions: {heads}"
        )
    return heads[0]


def _verify_schema(engine):
    head_revision = _get_latest_schema_revision()
    current_rev = _get_schema_version(engine)
    if current_rev != head_revision:
        raise MlflowException(
            f"Detected out-of-date database schema (found version {current_rev}, "
            f"but expected {head_revision}). Take a backup of your database, then run "
            "'mlflow db upgrade <database_uri>' "
            "to migrate your database to the latest schema. NOTE: schema migration may "
            "result in database downtime - please consult your database's documentation for "
            "more detail."
        )


def _get_managed_session_maker(SessionMaker, db_type):
    """
    Creates a factory for producing exception-safe SQLAlchemy sessions that are made available
    using a context manager. Any session produced by this factory is automatically committed
    if no exceptions are encountered within its associated context. If an exception is
    encountered, the session is rolled back. Finally, any session produced by this factory is
    automatically closed when the session's associated context is exited.
    """

    @contextmanager
    def make_managed_session():
        """Provide a transactional scope around a series of operations."""
        with SessionMaker() as session:
            try:
                if db_type == SQLITE:
                    session.execute(sql.text("PRAGMA foreign_keys = ON;"))
                    session.execute(sql.text("PRAGMA busy_timeout = 20000;"))
                    session.execute(sql.text("PRAGMA case_sensitive_like = true;"))
                yield session
                session.commit()
            except MlflowException:
                session.rollback()
                raise
            except sqlalchemy.exc.OperationalError as e:
                session.rollback()
                _logger.exception(
                    "SQLAlchemy database error. The following exception is caught.\n%s",
                    e,
                )
                raise MlflowException(message=e, error_code=TEMPORARILY_UNAVAILABLE)
            except sqlalchemy.exc.SQLAlchemyError as e:
                session.rollback()
                raise MlflowException(message=e, error_code=BAD_REQUEST)
            except Exception as e:
                session.rollback()
                raise MlflowException(message=e, error_code=INTERNAL_ERROR)

    return make_managed_session


def _get_alembic_config(db_url, alembic_dir=None):
    """
    Constructs an alembic Config object referencing the specified database and migration script
    directory.

    Args:
        db_url: Database URL, like sqlite:///<absolute-path-to-local-db-file>. See
            https://docs.sqlalchemy.org/en/13/core/engines.html#database-urls for a full list of
            valid database URLs.
        alembic_dir: Path to migration script directory. Uses canonical migration script
            directory under mlflow/alembic if unspecified. TODO: remove this argument in MLflow 1.1,
            as it's only used to run special migrations for pre-1.0 users to remove duplicate
            constraint names.
    """
    from alembic.config import Config

    final_alembic_dir = (
        os.path.join(_get_package_dir(), "store", "db_migrations")
        if alembic_dir is None
        else alembic_dir
    )
    # Escape any '%' that appears in a db_url. This could be in a password,
    # url, or anything that is part of a potentially complex database url
    db_url = db_url.replace("%", "%%")
    config = Config(os.path.join(final_alembic_dir, "alembic.ini"))
    config.set_main_option("script_location", final_alembic_dir)
    config.set_main_option("sqlalchemy.url", db_url)
    return config


def _upgrade_db(engine):  # noqa: D417
    """
    Upgrade the schema of an MLflow tracking database to the latest supported version.
    Note that schema migrations can be slow and are not guaranteed to be transactional -
    we recommend taking a backup of your database before running migrations.

    Args:
        url: Database URL, like sqlite:///<absolute-path-to-local-db-file>. See
            https://docs.sqlalchemy.org/en/13/core/engines.html#database-urls for a full list of
            valid database URLs.
    """
    # alembic adds significant import time, so we import it lazily
    from alembic import command

    db_url = str(engine.url)
    _logger.info("Updating database tables")
    config = _get_alembic_config(db_url)
    # Initialize a shared connection to be used for the database upgrade, ensuring that
    # any connection-dependent state (e.g., the state of an in-memory database) is preserved
    # for reference by the upgrade routine. For more information, see
    # https://alembic.sqlalchemy.org/en/latest/cookbook.html#sharing-a-
    # connection-with-a-series-of-migration-commands-and-environments
    with engine.begin() as connection:
        config.attributes["connection"] = connection
        command.upgrade(config, "heads")


def _get_schema_version(engine):
    with engine.connect() as connection:
        mc = MigrationContext.configure(connection)
        return mc.get_current_revision()


def create_sqlalchemy_engine_with_retry(db_uri):
    attempts = 0
    while True:
        attempts += 1
        engine = create_sqlalchemy_engine(db_uri)
        try:
            sqlalchemy.inspect(engine)
            return engine
        except Exception as e:
            if attempts < MAX_RETRY_COUNT:
                sleep_duration = 0.1 * ((2**attempts) - 1)
                _logger.warning(
                    "SQLAlchemy engine could not be created. The following exception is caught.\n"
                    "%s\nOperation will be retried in %.1f seconds",
                    e,
                    sleep_duration,
                )
                time.sleep(sleep_duration)
                continue
            raise


def create_sqlalchemy_engine(db_uri):
    pool_size = MLFLOW_SQLALCHEMYSTORE_POOL_SIZE.get()
    pool_max_overflow = MLFLOW_SQLALCHEMYSTORE_MAX_OVERFLOW.get()
    pool_recycle = MLFLOW_SQLALCHEMYSTORE_POOL_RECYCLE.get()
    echo = MLFLOW_SQLALCHEMYSTORE_ECHO.get()
    poolclass = MLFLOW_SQLALCHEMYSTORE_POOLCLASS.get()
    kwargs = {}
    # Send argument only if they have been injected.
    # Some engine does not support them (for example sqllite)
    if pool_size:
        kwargs["pool_size"] = pool_size
    if pool_max_overflow:
        kwargs["max_overflow"] = pool_max_overflow
    if pool_recycle:
        kwargs["pool_recycle"] = pool_recycle
    if echo:
        kwargs["echo"] = echo
    if poolclass:
        pool_class_map = {
            "AssertionPool": AssertionPool,
            "AsyncAdaptedQueuePool": AsyncAdaptedQueuePool,
            "FallbackAsyncAdaptedQueuePool": FallbackAsyncAdaptedQueuePool,
            "NullPool": NullPool,
            "QueuePool": QueuePool,
            "SingletonThreadPool": SingletonThreadPool,
            "StaticPool": StaticPool,
        }
        if poolclass not in pool_class_map:
            list_str = " ".join(pool_class_map.keys())
            err_str = (
                f"Invalid poolclass parameter: {poolclass}. Set environment variable "
                f"poolclass to empty or one of the following values: {list_str}"
            )
            _logger.warning(err_str)
            raise ValueError(err_str)
        kwargs["poolclass"] = pool_class_map[poolclass]
    if kwargs:
        _logger.info("Create SQLAlchemy engine with pool options %s", kwargs)

    # Handle MySQL SSL certificates via connect_args
    if db_uri.startswith("mysql"):
        connect_args = {
            k: v
            for k, v in {
                "ssl_ca": MLFLOW_MYSQL_SSL_CA.get(),
                "ssl_cert": MLFLOW_MYSQL_SSL_CERT.get(),
                "ssl_key": MLFLOW_MYSQL_SSL_KEY.get(),
            }.items()
            if v
        }
        if connect_args:
            kwargs["connect_args"] = connect_args

    return sqlalchemy.create_engine(db_uri, pool_pre_ping=True, **kwargs)