#!/usr/bin/env python3

import json
import yaml
import os
import sys
from pathvalidate import sanitize_filename
from yaml import CDumper as Dumper


class Loader(yaml.CLoader):
    """Hack to store the stream before passing it off to :class:`yaml.CLoader`"""
    def __init__(self, stream):
        self.stream = stream
        super().__init__(stream)


class IncludeTag:
    yaml_tag = '!include'

    def __init__(self, target):
        self.target = target

    @classmethod
    def to_yaml(cls, dumper, node):
        return dumper.represent_scalar(cls.yaml_tag, node.target, style="'")

    @classmethod
    def from_yaml(cls, loader, node):
        base_dir = os.path.dirname(loader.stream.name)
        target = os.path.join(base_dir, node.value)
        with open(target) as f:
            if target.endswith('.yaml'):
                return yaml.load(f, Loader=Loader)
            else:
                return f.read()


Loader.add_constructor('!include', IncludeTag.from_yaml)
Dumper.add_representer(IncludeTag, IncludeTag.to_yaml)


def uniqueName(obj):
    if 'SaveName' in obj:
        return obj['SaveName']
    return f"{obj['Name']} {obj['Nickname']} {obj['GUID']}"


def shouldRecurse(obj, subobject_prop):
    return subobject_prop in obj \
        and obj.get('Name') not in ('Deck', 'DeckCustom') \
        and len(obj[subobject_prop]) > 1

def normalize_value(value):
    rounded = round(value, 2)
    if rounded == -0.0:
        return 0.0
    else:
        return rounded

def recursively_normalize_values(d):
    if type(d) == dict:
        return {key: recursively_normalize_values(value) for key, value in d.items()}
    elif type(d) == list:
        return [recursively_normalize_values(value) for value in d]
    elif type(d) == float:
        return normalize_value(d)
    else:
        return d

def recursivelyUnpackObject(parent_dir, obj,
                            subobject_prop='ContainedObjects', base_name=None):
    def convertObjectsToIncludes(objects):
        for subObj in objects:
            filename = recursivelyUnpackObject(file_base_path, subObj)
            yield IncludeTag(os.path.join(obj_base_name, filename))

    obj_base_name = base_name or sanitize_filename(uniqueName(obj))
    file_base_path = os.path.join(parent_dir, obj_base_name)

    if shouldRecurse(obj, subobject_prop):
        os.makedirs(file_base_path, exist_ok=True)

        obj[subobject_prop] = list(
            convertObjectsToIncludes(obj[subobject_prop]))

    if 'LuaScript' in obj and len(obj['LuaScript']) > 0:
        with open(file_base_path + '.ttslua', 'w') as f:
            f.write(obj['LuaScript'].replace("\r\n", "\n"))

        obj['LuaScript'] = IncludeTag(obj_base_name + '.ttslua')

    if 'XmlUI' in obj and len(obj['XmlUI']) > 0:
        with open(file_base_path + '.xml', 'w') as f:
            f.write(obj['XmlUI'].replace("\r\n", "\n"))

        obj['XmlUI'] = IncludeTag(obj_base_name + '.xml')

    # round transforms, as TTS seems to slightly change them on each save
    for k in ['Transform', 'AttachedSnapPoints', 'SnapPoints', 'Hands']:
        if k in obj:
            obj[k] = recursively_normalize_values(obj[k])

    with open(file_base_path + '.yaml', 'w') as f:
        yaml.dump(obj, f, Dumper=Dumper)

    return obj_base_name + '.yaml'


def unpackJson(json_file, output_name):
    with open(json_file) as f:
        data = json.load(f)

    if 'ObjectStates' in data:
        prop = 'ObjectStates'
    else:
        prop = 'ContainedObjects'
    recursivelyUnpackObject(
        '', data, subobject_prop=prop, base_name=output_name)


def packYaml(yaml_file, output_json_file):
    with open(yaml_file) as f:
        data = yaml.load(f, Loader=Loader)

    with open(output_json_file, 'w') as f:
        json.dump(data, f, indent=2)


def usage():
    print(f"Usage: {sys.argv[0]} <unpack|pack> <input> <output>")


def main():
    if len(sys.argv) != 4:
        usage()
    elif sys.argv[1] == 'unpack':
        unpackJson(sys.argv[2], sys.argv[3])
    elif sys.argv[1] == 'pack':
        packYaml(sys.argv[2], sys.argv[3])
    else:
        usage()


if __name__ == '__main__':
    main()