# -*- coding: utf-8 -*-
r"""Simple wrapper for 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"]


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):
    return yaml.load(stream, Loader)


def load_all(stream, Loader=ConfigLoader):
    return yaml.load_all(stream, Loader)


def safe_load(stream):
    return yaml.load(stream, Loader=ConfigSafeLoader)


def safe_load_all(stream):
    return yaml.load_all(stream, Loader=ConfigSafeLoader)