# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import os

import alembic
import alembic.config
import alembic.migration
import alembic.script
import sqlalchemy as sa
import sqlalchemy.schema

from rally import exceptions
from rally.common.db import api
from rally.common.db import models


INITIAL_REVISION_UUID = "ca3626f62937"


def _alembic_config():
    path = os.path.join(os.path.dirname(__file__), "alembic.ini")
    config = alembic.config.Config(path)
    return config


def schema_cleanup():
    """Drop all database objects.

    Drops all database objects remaining on the default schema of the given
    engine. Per-db implementations will also need to drop items specific to
    those systems, such as sequences, custom types (e.g. pg ENUM), etc.
    """
    engine = api.get_engine()
    with engine.begin() as conn:
        inspector = sa.inspect(engine)
        metadata = sa.schema.MetaData()
        tbs = []
        all_fks = []

        for table_name in inspector.get_table_names():
            fks = []
            for fk in inspector.get_foreign_keys(table_name):
                if not fk["name"]:
                    continue
                fks.append(
                    sa.schema.ForeignKeyConstraint((), (), name=fk["name"])
                )
            table = sa.schema.Table(table_name, metadata, *fks)
            tbs.append(table)
            all_fks.extend(fks)

        if engine.name != "sqlite":
            for fkc in all_fks:
                conn.execute(sa.schema.DropConstraint(fkc))
        for table in tbs:
            conn.execute(sa.schema.DropTable(table))

        if engine.name == "postgresql":
            sqla_100 = int(sa.__version__.split(".")[0]) >= 1

            if sqla_100:
                enums = [e["name"] for e in sa.inspect(conn).get_enums()]
            else:
                enums = conn.dialect._load_enums(conn).keys()

            for e in enums:
                conn.execute("DROP TYPE %s" % e)


def schema_revision(config=None, engine=None, detailed=False):
    """Current database revision.

    :param config: Instance of alembic config
    :param engine: Instance of DB engine
    :param detailed: whether to return a dict with detailed data
    :rtype detailed: bool
    :returns: Database revision
    :rtype: string
    :rtype: dict
    """
    engine = engine or api.get_engine()
    with engine.connect() as conn:
        context = alembic.migration.MigrationContext.configure(conn)
        revision = context.get_current_revision()
    if detailed:
        config = config or _alembic_config()
        sc_dir = alembic.script.ScriptDirectory.from_config(config)
        return {
            "revision": revision,
            "current_head": sc_dir.get_current_head(),
        }
    return revision


def schema_upgrade(revision=None, config=None, engine=None):
    """Used for upgrading database.

    :param revision: Desired database version
    :type revision: string
    :param config: Instance of alembic config
    :param engine: Instance of DB engine
    """
    revision = revision or "head"
    config = config or _alembic_config()
    engine = engine or api.get_engine()

    if schema_revision(engine=engine) is None:
        schema_stamp(INITIAL_REVISION_UUID, config=config)

    alembic.command.upgrade(config, revision or "head")


def schema_create(config=None, engine=None):
    """Create database schema from models description.

    Can be used for initial installation instead of upgrade('head').
    :param config: Instance of alembic config
    :param engine: Instance of DB engine
    """
    engine = engine or api.get_engine()

    # NOTE(viktors): If we will use metadata.create_all() for non empty db
    #                schema, it will only add the new tables, but leave
    #                existing as is. So we should avoid of this situation.
    if schema_revision(engine=engine) is not None:
        raise exceptions.DBMigrationError(
            "DB schema is already under version control. Use upgrade() instead"
        )

    models.BASE.metadata.create_all(engine)
    schema_stamp("head", config=config)


def schema_stamp(revision, config=None):
    """Stamps database with provided revision.

    Don't run any migrations.
    :param revision: Should match one from repository or head - to stamp
                     database with most recent revision
    :type revision: string
    :param config: Instance of alembic config
    """
    config = config or _alembic_config()
    return alembic.command.stamp(config, revision)
