# -*- coding: utf-8 -*-
# :-
# :Copyright: (c) 2015-2020, Franz Glasner. All rights reserved.
# :License:   3-clause BSD. See LICENSE.txt for details.
# :-
"""Simple wrapper for :mod:`yaml` to support all-unicode strings when
loading configuration files.

"""

from __future__ import division, print_function, absolute_import

try:
    from collections import OrderedDict
except ImportError:
    try:
        from ordereddict import OrderedDict
    except ImportError:
        OrderedDict = None
import yaml
import yaml.constructor

from .compat import u


__all__ = ["safe_load", "safe_load_all", "load", "load_all"]


DictImpl = OrderedDict or dict


class ConfigLoader(yaml.Loader):

    """A YAML loader, which makes all ``!!str`` strings to Unicode.
    Standard PyYAML does this only in the non-ASCII case.

    If an `OrderedDict` implementation is available then all "map" and
    "omap" nodes are constructed as `OrderedDict`.
    This is against YAML specs but within configuration files it seems
    more natural.

    """

    def construct_yaml_str(self, node):
        return self.construct_scalar(node)

    if OrderedDict:

        #
        # From https://pypi.python.org/pypi/yamlordereddictloader/0.1.1
        # (MIT License)
        #

        def construct_yaml_map(self, node):
            data = OrderedDict()
            yield data
            value = self.construct_mapping(node)
            data.update(value)

        def construct_mapping(self, node, deep=False):
            if isinstance(node, yaml.MappingNode):
                self.flatten_mapping(node)
            else:
                raise yaml.constructor.ConstructorError(
                    None,
                    None,
                    'expected a mapping node, but found %s' % node.id,
                    node.start_mark)

            mapping = OrderedDict()
            for key_node, value_node in node.value:
                key = self.construct_object(key_node, deep=deep)
                try:
                    hash(key)
                except TypeError as err:
                    raise yaml.constructor.ConstructorError(
                        'while constructing a mapping', node.start_mark,
                        'found unacceptable key (%s)' % (err,
                                                         key_node.start_mark)
                    )
                value = self.construct_object(value_node, deep=deep)
                mapping[key] = value
            return mapping


ConfigLoader.add_constructor(
    u("tag:yaml.org,2002:str"),
    ConfigLoader.construct_yaml_str)
if OrderedDict:
    ConfigLoader.add_constructor(
        u("tag:yaml.org,2002:map"),
        ConfigLoader.construct_yaml_map)
    ConfigLoader.add_constructor(
        u("tag:yaml.org,2002:omap"),
        ConfigLoader.construct_yaml_map)


class ConfigSafeLoader(yaml.SafeLoader):

    """A safe YAML loader, which makes all ``!!str`` strings to Unicode.
    Standard PyYAML does this only in the non-ASCII case.

    If an `OrderedDict` implementation is available then all "map" and
    "omap" nodes are constructed as `OrderedDict`.
    This is against YAML specs but within configuration files it seems
    more natural.

    """

    def construct_yaml_str(self, node):
        return self.construct_scalar(node)

    if OrderedDict:

        #
        # From https://pypi.python.org/pypi/yamlordereddictloader/0.1.1
        # (MIT License)
        #

        def construct_yaml_map(self, node):
            data = OrderedDict()
            yield data
            value = self.construct_mapping(node)
            data.update(value)

        def construct_mapping(self, node, deep=False):
            if isinstance(node, yaml.MappingNode):
                self.flatten_mapping(node)
            else:
                raise yaml.constructor.ConstructorError(
                    None,
                    None,
                    'expected a mapping node, but found %s' % node.id,
                    node.start_mark)

            mapping = OrderedDict()
            for key_node, value_node in node.value:
                key = self.construct_object(key_node, deep=deep)
                try:
                    hash(key)
                except TypeError as err:
                    raise yaml.constructor.ConstructorError(
                        'while constructing a mapping', node.start_mark,
                        'found unacceptable key (%s)' % (err,
                                                         key_node.start_mark)
                    )
                value = self.construct_object(value_node, deep=deep)
                mapping[key] = value
            return mapping


ConfigSafeLoader.add_constructor(
    u("tag:yaml.org,2002:str"),
    ConfigSafeLoader.construct_yaml_str)
if OrderedDict:
    ConfigSafeLoader.add_constructor(
        u("tag:yaml.org,2002:map"),
        ConfigSafeLoader.construct_yaml_map)
    ConfigSafeLoader.add_constructor(
        u("tag:yaml.org,2002:omap"),
        ConfigSafeLoader.construct_yaml_map)


def load(stream, Loader=ConfigLoader):
    """Parse the given `stream` and return a Python object constructed
    from for the first document in the stream.

    """
    data = yaml.load(stream, Loader)
    # Map an empty document to an empty dict
    if data is None:
        return DictImpl()
    if not isinstance(data, DictImpl):
        raise TypeError("YAML root object must be a mapping")
    return data


def load_all(stream, Loader=ConfigLoader):
    """Parse the given `stream` and return a sequence of Python objects
    corresponding to the documents in the `stream`.

    """
    data_all = yaml.load_all(stream, Loader)
    rdata = []
    for data in data_all:
        if data is None:
            rdata.append(DictImpl())
        else:
            if not isinstance(data, DictImpl):
                raise TypeError("YAML root object must be a mapping")
            rdata.append(data)
    return rdata


def safe_load(stream):
    """Parse the given `stream` and return a Python object constructed
    from for the first document in the stream.

    Recognizes only standard YAML tags and cannot construct an
    arbitrary Python object.

    """
    data = yaml.load(stream, Loader=ConfigSafeLoader)
    # Map an empty document to an empty dict
    if data is None:
        return DictImpl()
    if not isinstance(data, DictImpl):
        raise TypeError("YAML root object must be a mapping")
    return data


def safe_load_all(stream):
    """Return the list of all decoded YAML documents in the file `stream`.

    Recognizes only standard YAML tags and cannot construct an
    arbitrary Python object.

    """
    data_all = yaml.load_all(stream, Loader=ConfigSafeLoader)
    rdata = []
    for data in data_all:
        if data is None:
            rdata.append(DictImpl())
        else:
            if not isinstance(data, DictImpl):
                raise TypeError("YAML root object must be a mapping")
            rdata.append(data)
    return data_all