import yaml
from copy import deepcopy
from yaml.constructor import ConstructorError, SafeConstructor

from load.tools.yaml_injection.resource_loader import URLLoader, FileLoader, BaseResourceLoader

try:
    from collections.abc import Hashable
except ImportError:
    from collections import Hashable

REFERENCE_SOURCE_KEY = 'ref'


class RefLoader(BaseResourceLoader):
    def __init__(self):
        self.mutable_data_root = None

    def download_resource(self, path_raw):
        assert self.mutable_data_root is not None
        path = path_raw.split()
        data = self.mutable_data_root
        path_got = []
        for k in path:
            path_got.append(k)
            try:
                data = data[k]
            except KeyError:
                raise KeyError('There is no path {%s}.' % path_got)
        return yaml.dump(data)

    def set_data_root(self, data_root):
        assert self.mutable_data_root is None
        self.mutable_data_root = data_root

    def copy(self):
        return self.__class__()


class InjectionLoader(yaml.SafeLoader):
    def __init__(self, stream):
        super(self.__class__, self).__init__(stream)
        self.data = None
        self.main_node = None
        self.resources_loaders_map = {
            'url': URLLoader(),
            'file': FileLoader(),
            'ref': RefLoader(),
        }

    @classmethod
    def with_resources(cls, *resources_loaders):
        def constructor(stream):
            naked_loader = cls(stream)
            for resource_loader in resources_loaders:
                key = resource_loader.key()
                if key == REFERENCE_SOURCE_KEY:
                    raise ValueError('"{%s}" is reserved key for resource loader' % REFERENCE_SOURCE_KEY)
                naked_loader.resources_loaders_map[key] = resource_loader
            return naked_loader

        return constructor

    def _copy(self, stream):
        new_loader = self.__class__(stream)
        new_loader.resources_loaders_map = {k: l.copy()
                                            for k, l in self.resources_loaders_map.items()}
        return new_loader

    def get_single_data(self):
        # Ensure that the stream contains a single document and construct it.
        node = self.get_single_node()
        self.main_node = node
        if node is not None:
            return self.construct_document(node)
        return None

    def _construct_preloaded(self):
        while self.state_generators:
            state_generators = self.state_generators
            self.state_generators = []
            for generator in state_generators:
                for dummy in generator:
                    pass

    def _inject(self, mapping, value_node, source=None):
        if source:
            items = self.construct_object(value_node, deep=True)
            if not isinstance(items, list):
                items = [items]
            sources = {source: items}
        else:
            sources = self.construct_mapping(value_node, deep=True)

        for source_key, sources_specifications in sources.items():
            if not isinstance(sources_specifications, list):
                sources_specifications = [sources_specifications]

            for specification in sources_specifications:
                self._construct_preloaded()

                source_data_text = self.resources_loaders_map[source_key].download_resource(specification)

                inner_loader = self._copy(source_data_text)
                inner_node = inner_loader.get_single_node()
                if inner_node is None:
                    continue
                inner_loader.main_node = inner_node
                inner_data = inner_loader.construct_mapping(inner_node,
                                                            prepared=mapping)
                inner_loader.resources_loaders_map[REFERENCE_SOURCE_KEY].set_data_root(inner_data)
                inner_loader._construct_preloaded()
                if inner_data:
                    mapping.update(inner_data)

    def construct_mapping(self, node, deep=False, prepared=None):
        if isinstance(node, yaml.MappingNode):
            self.flatten_mapping(node)
        if not isinstance(node, yaml.MappingNode):
            raise ConstructorError(None, None,
                                   'expected a mapping node, but found %s' % node.id,
                                   node.start_mark)
        mapping = {}
        if prepared is not None:
            mapping = deepcopy(prepared)

        for key_node, value_node in node.value:
            key = self.construct_object(key_node, deep=deep)
            if not isinstance(key, Hashable):
                raise ConstructorError('while constructing a mapping', node.start_mark,
                                       'found unhashable key', key_node.start_mark)
            if key_node.tag == '!inject':
                self._inject(mapping, value_node, source=key)
            elif key in mapping and isinstance(value_node, yaml.MappingNode):
                generator = self.construct_yaml_map_prepared(value_node, mapping[key])
                mapping[key] = next(generator)
                self.state_generators.append(generator)
            elif isinstance(value_node, yaml.MappingNode):
                generator = self.construct_yaml_map_implicit_preparation(value_node)
                mapping[key] = next(generator)
                self.state_generators.append(generator)
            else:
                value = self.construct_object(value_node, deep=deep)
                mapping[key] = value
        return mapping

    @staticmethod
    def update_from_prepared(data, value, prepared):
        """
        :param data:
        :param value: новые данные. Считаем их приоритетными.
        :param prepared: данные, которые уже были загружены
        :return:
        """
        for k, v in prepared.items():
            if k not in value:
                value[k] = prepared[k]
                continue
            if not (isinstance(value[k], dict)):
                continue
            value[k].update(prepared[k])
        data.update(value)

    def construct_yaml_map_prepared(self, node, prepared):
        data = {}
        yield data
        value = self.construct_mapping(node)
        self.update_from_prepared(data, value, prepared)

    def construct_yaml_map_implicit_preparation(self, node):
        data = {}
        yield data
        value = self.construct_mapping(node)
        self.update_from_prepared(data, value, data)

    def construct_object(self, node, deep=False):
        data = super(self.__class__, self).construct_object(node, deep=deep)
        if node is self.main_node:
            self.resources_loaders_map[REFERENCE_SOURCE_KEY].set_data_root(data)
        return data


InjectionLoader.add_constructor('!inject', SafeConstructor.construct_yaml_str)
