# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved.                     #
# This file is part of the AiiDA code.                                    #
#                                                                         #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file        #
# For further information please visit http://www.aiida.net               #
###########################################################################
"""Definition of caching mechanism and configuration for calculations."""
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import io
import os
import copy
from enum import Enum
from contextlib import contextmanager

import yaml
import six
from wrapt import decorator

from aiida.common import exceptions
from aiida.common.utils import get_object_from_string

__all__ = ('get_use_cache', 'enable_caching', 'disable_caching')


class ConfigKeys(Enum):
    """Valid keys for caching configuration."""

    DEFAULT = 'default'
    ENABLED = 'enabled'
    DISABLED = 'disabled'


DEFAULT_CONFIG = {
    ConfigKeys.DEFAULT.value: False,
    ConfigKeys.ENABLED.value: [],
    ConfigKeys.DISABLED.value: [],
}


def _get_config(config_file):
    """Return the caching configuration.

    :param config_file: the absolute path to the caching configuration file
    :return: the configuration dictionary
    """
    from aiida.manage.configuration import get_profile

    profile = get_profile()

    if profile is None:
        exceptions.ConfigurationError('no profile has been loaded')

    try:
        with io.open(config_file, 'r', encoding='utf8') as handle:
            config = yaml.load(handle)[profile.name]
    # no config file, or no config for this profile
    except (OSError, IOError, KeyError):
        return DEFAULT_CONFIG

    # validate configuration
    for key in config:
        if key not in DEFAULT_CONFIG:
            raise ValueError("Configuration error: Invalid key '{}' in cache_config.yml".format(key))

    # add defaults where config is missing
    for key, default_config in DEFAULT_CONFIG.items():
        config[key] = config.get(key, default_config)

    # load classes
    try:
        for key in [ConfigKeys.ENABLED.value, ConfigKeys.DISABLED.value]:
            config[key] = [get_object_from_string(c) for c in config[key]]
    except (ValueError) as err:
        six.raise_from(
            exceptions.ConfigurationError("Unknown class given in 'cache_config.yml': '{}'".format(err)), err
        )
    return config


_CONFIG = {}


def configure(config_file=None):
    """
    Reads the caching configuration file and sets the _CONFIG variable.
    """
    # pylint: disable=global-statement
    if config_file is None:
        from aiida.manage.configuration import get_config

        config = get_config()
        config_file = os.path.join(config.dirpath, 'cache_config.yml')

    global _CONFIG
    _CONFIG.clear()
    _CONFIG.update(_get_config(config_file=config_file))


@decorator
def _with_config(wrapped, _, args, kwargs):
    """Function decorator to load the caching configuration for the scope of the wrapped function."""
    if not _CONFIG:
        configure()
    return wrapped(*args, **kwargs)


@_with_config
def get_use_cache(node_class=None):
    """Return whether the caching mechanism should be used for the given node class according to the configuration.

    :param node_class: the Node class or sub class to check if enabled for caching
    :return: boolean, True if caching is enabled, False otherwise
    :raises ValueError: if the configuration is invalid by defining the class both enabled and disabled
    """
    if node_class is not None:
        enabled = node_class in _CONFIG[ConfigKeys.ENABLED.value]
        disabled = node_class in _CONFIG[ConfigKeys.DISABLED.value]
        if enabled and disabled:
            raise ValueError('Invalid configuration: Caching for {} is both enabled and disabled.'.format(node_class))
        elif enabled:
            return True
        elif disabled:
            return False
    return _CONFIG[ConfigKeys.DEFAULT.value]


@contextmanager
@_with_config
def _reset_config():
    """Reset the configuration by clearing the contents of the global config variable."""
    # pylint: disable=global-statement
    global _CONFIG
    config_copy = copy.deepcopy(_CONFIG)
    yield
    _CONFIG.clear()
    _CONFIG.update(config_copy)


@contextmanager
def enable_caching(node_class=None):
    """
    Context manager to enable caching, either for a specific node class, or globally.
    Note that this does not affect the behavior of the daemon, only the local Python instance.

    :param node_class: Node class for which caching should be enabled.
    """
    with _reset_config():
        if node_class is None:
            _CONFIG[ConfigKeys.DEFAULT.value] = True
        else:
            _CONFIG[ConfigKeys.ENABLED.value].append(node_class)
            try:
                _CONFIG[ConfigKeys.DISABLED.value].remove(node_class)
            except ValueError:
                pass
        yield


@contextmanager
def disable_caching(node_class=None):
    """
    Context manager to disable caching, either for a specific node class, or globally.
    Note that this does not affect the behavior of the daemon, only the local Python instance.

    :param node_class: Node class for which caching should be disabled.
    """
    with _reset_config():
        if node_class is None:
            _CONFIG[ConfigKeys.DEFAULT.value] = False
        else:
            _CONFIG[ConfigKeys.DISABLED.value].append(node_class)
            try:
                _CONFIG[ConfigKeys.ENABLED.value].remove(node_class)
            except ValueError:
                pass
        yield
