# -*- coding:utf-8 -*-
""" pretty serialization
    original from <http://djangosnippets.org/snippets/2397/>
"""
from io import StringIO
from datetime import datetime
import yaml

try:
    from yaml import CSafeLoader as SafeLoader
except ImportError:
    from yaml import SafeLoader
from django.utils import timezone
from django.core.serializers.base import DeserializationError
from django.core.serializers.pyyaml import (
    Serializer as YamlSerializer, DjangoSafeDumper)
from django.core.serializers.python import (
    Deserializer as PythonDeserializer,
)


# noinspection PyClassHasNoInit
class Serializer(YamlSerializer):
    """ utf8-friendly dumpdata management command """

    def end_serialization(self):
        yaml.dump(self.objects, self.stream, allow_unicode=True,
                  default_flow_style=False,
                  Dumper=DjangoSafeDumper, **self.options)


def load_yaml(stream_or_string):
    """ laod yaml data with adding timezone """
    output = yaml.load(stream_or_string, SafeLoader)
    for a_model in output:
        for key, value in a_model.items():
            if key == 'fields':
                for vkey, vvalue in value.items():
                    if isinstance(vvalue, datetime):
                        value[vkey] = vvalue.replace(tzinfo=timezone.utc)
    return output


# noinspection PyPep8Naming
def Deserializer(stream_or_string, **options):  # pylint:disable=invalid-name
    """Deserialize a stream or string of YAML data."""
    if isinstance(stream_or_string, bytes):
        stream_or_string = stream_or_string.decode()
    if isinstance(stream_or_string, str):
        stream = StringIO(stream_or_string)
    else:
        stream = stream_or_string
    try:
        yield from PythonDeserializer(load_yaml(stream), **options)
    except (GeneratorExit, DeserializationError):
        raise
    except Exception as exc:
        raise DeserializationError() from exc