# -*- coding: utf-8 -*-
# pylint: disable=import-error,no-name-in-module
"""Utilities and configuration of the SqlAlchemy database schema."""
from __future__ import absolute_import

import os
import contextlib

from alembic import command
from alembic.config import Config
from alembic.runtime.environment import EnvironmentContext
from alembic.script import ScriptDirectory

from sqlalchemy.orm.exc import NoResultFound

from aiida.backends.sqlalchemy import get_scoped_session
from aiida.common import NotExistent
from ..manager import BackendManager, SettingsManager, Setting

ALEMBIC_FILENAME = 'alembic.ini'
ALEMBIC_REL_PATH = 'migrations'

# The database schema version required to perform schema reset for a given code schema generation
SCHEMA_VERSION_RESET = {'1': None}


class SqlaBackendManager(BackendManager):
    """Class to manage the database schema."""

    @staticmethod
    @contextlib.contextmanager
    def alembic_config():
        """Context manager to return an instance of an Alembic configuration with the current connection inserted.

        :return: instance of :py:class:`alembic.config.Config`
        """
        from . import ENGINE

        with ENGINE.begin() as connection:
            dir_path = os.path.dirname(os.path.realpath(__file__))
            config = Config(os.path.join(dir_path, ALEMBIC_FILENAME))
            config.set_main_option('script_location', os.path.join(dir_path, ALEMBIC_REL_PATH))
            config.attributes['connection'] = connection  # pylint: disable=unsupported-assignment-operation
            yield config

    def get_settings_manager(self):
        """Return an instance of the `SettingsManager`.

        :return: `SettingsManager`
        """
        if self._settings_manager is None:
            self._settings_manager = SqlaSettingsManager()

        return self._settings_manager

    def _load_backend_environment(self):
        """Load the backend environment."""
        from . import reset_session
        reset_session()

    def reset_backend_environment(self):
        """Reset the backend environment."""
        from aiida.backends import sqlalchemy
        if sqlalchemy.ENGINE is not None:
            sqlalchemy.ENGINE.dispose()
        sqlalchemy.SCOPED_SESSION_CLASS = None

    def is_database_schema_ahead(self):
        """Determine whether the database schema version is ahead of the code schema version.

        .. warning:: this will not check whether the schema generations are equal

        :return: boolean, True if the database schema version is ahead of the code schema version.
        """
        from alembic.util import CommandError

        # In the case of SqlAlchemy, if the database revision is ahead of the code, that means the revision stored in
        # the database is not even present in the code base. Therefore we cannot locate it in the revision graph and
        # determine whether it is ahead of the current code head. We simply try to get the revision and if it does not
        # exist it means it is ahead.
        with self.alembic_config() as config:
            try:
                script = ScriptDirectory.from_config(config)
                script.get_revision(self.get_schema_version_database())
            except CommandError:
                # Raised when the revision of the database is not present in the revision graph.
                return True
            else:
                return False

    def get_schema_version_code(self):
        """Return the code schema version."""
        with self.alembic_config() as config:
            script = ScriptDirectory.from_config(config)
            schema_version_code = script.get_current_head()

        return schema_version_code

    def get_schema_version_reset(self, schema_generation_code):
        """Return schema version the database should have to be able to automatically reset to code schema generation.

        :param schema_generation_code: the schema generation of the code.
        :return: schema version
        """
        return SCHEMA_VERSION_RESET[schema_generation_code]

    def get_schema_version_database(self):
        """Return the database schema version.

        :return: `distutils.version.StrictVersion` with schema version of the database
        """

        def get_database_version(revision, _):
            """Get the current revision."""
            if isinstance(revision, tuple) and revision:
                config.attributes['rev'] = revision[0]  # pylint: disable=unsupported-assignment-operation
            else:
                config.attributes['rev'] = None  # pylint: disable=unsupported-assignment-operation
            return []

        with self.alembic_config() as config:

            script = ScriptDirectory.from_config(config)

            with EnvironmentContext(config, script, fn=get_database_version):
                script.run_env()
                return config.attributes['rev']  # pylint: disable=unsubscriptable-object

    def set_schema_version_database(self, version):
        """Set the database schema version.

        :param version: string with schema version to set
        """
        # pylint: disable=cyclic-import
        from aiida.manage.manager import get_manager
        backend = get_manager()._load_backend(schema_check=False)  # pylint: disable=protected-access
        backend.execute_raw(r"""UPDATE alembic_version SET version_num='{}';""".format(version))

    def _migrate_database_version(self):
        """Migrate the database to the current schema version."""
        super(SqlaBackendManager, self)._migrate_database_version()
        with self.alembic_config() as config:
            command.upgrade(config, 'head')


class SqlaSettingsManager(SettingsManager):
    """Class to get, set and delete settings from the `DbSettings` table."""

    table_name = 'db_dbsetting'

    def validate_table_existence(self):
        """Verify that the `DbSetting` table actually exists.

        :raises: `~aiida.common.exceptions.NotExistent` if the settings table does not exist
        """
        from sqlalchemy.engine import reflection
        inspector = reflection.Inspector.from_engine(get_scoped_session().bind)
        if self.table_name not in inspector.get_table_names():
            raise NotExistent('the settings table does not exist')

    def get(self, key):
        """Return the setting with the given key.

        :param key: the key identifying the setting
        :return: Setting
        :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist
        """
        from aiida.backends.sqlalchemy.models.settings import DbSetting
        self.validate_table_existence()

        try:
            setting = get_scoped_session().query(DbSetting).filter_by(key=key).one()
        except NoResultFound:
            raise NotExistent('setting `{}` does not exist'.format(key))

        return Setting(key, setting.getvalue(), setting.description, setting.time)

    def set(self, key, value, description=None):
        """Return the settings with the given key.

        :param key: the key identifying the setting
        :param value: the value for the setting
        :param description: optional setting description
        """
        from aiida.backends.sqlalchemy.models.settings import DbSetting
        from aiida.orm.utils.node import validate_attribute_extra_key

        self.validate_table_existence()
        validate_attribute_extra_key(key)

        other_attribs = dict()
        if description is not None:
            other_attribs['description'] = description

        DbSetting.set_value(key, value, other_attribs=other_attribs)

    def delete(self, key):
        """Delete the setting with the given key.

        :param key: the key identifying the setting
        :raises: `~aiida.common.exceptions.NotExistent` if the settings does not exist
        """
        from aiida.backends.sqlalchemy.models.settings import DbSetting
        self.validate_table_existence()

        try:
            setting = get_scoped_session().query(DbSetting).filter_by(key=key).one()
            setting.delete()
        except NoResultFound:
            raise NotExistent('setting `{}` does not exist'.format(key))
