import json
import os

import LookFileBakeAPI
from Katana import FnAttribute

class JsonLookFileBakeOutputFormat(
        LookFileBakeAPI.BaseLookFileBakeOutputFormat):
    """
    Class implementing a simple JSON LookFileBake output format.
    """

    # Class Variables ---------------------------------------------------------

    DisplayName = "as JSON"
    FileExtension = ""
    PassFileExtension = "json"

    # Instance Methods --------------------------------------------------------

    def writeSinglePass(self, passData):
        """
        @type passData: :py:obj:`LookFileBakeAPI.LookFilePassData`
        @rtype: :py:obj:`list` of :py:obj:`str`
        @param passData: The data representing a single Look File pass to be
            baked.
        @return: A list of paths to files which have been written.
        """
        # Create a JSON document from the given data.
        jsonPassData = {}
        jsonPassData["passName"] = passData.passName
        materialDict = {}
        jsonPassData["materialDict"] = materialDict
        for materialLocation, (locationType, attributes) in \
                passData.materialDict.items():
            materialDict[materialLocation] = {
                "type": locationType,
                "attributes": JsonLookFileBakeOutputFormat.attrToDict(
                    attributes)
            }
        # TODO: Also add passData.outputDictList, passData.rootOverrideDict
        # and passData.sharedOverridesDict.

        # Generate text that will ultimately be written into the text file.
        text = json.dumps(jsonPassData, indent=4)

        # Get the file path for this pass from the given pass data.
        filePath = passData.filePath

        # If the enclosing directory doesn't exist, then try to create it.
        dirPath = os.path.dirname(filePath)
        if not os.path.exists(dirPath):
            os.makedirs(dirPath)

        # Finally write the string containing the pass data into the file.
        with open(filePath, "w") as outFile:
            outFile.write(text)

        return [filePath]

    # Static Methods ----------------------------------------------------------

    @staticmethod
    def attrToDict(attribute):
        """
        Helper static method that converts an :py:obj:`FnAttribute` attribute into a
        Python dictionary.

        @type attribute: :py:obj:`FnAttribute` or :py:obj:`None`
        @param attribute: The attribute to be converted.
        @rtype: :py:obj:`dict`
        @return: The dictionary representing the given attribute.
        """
        result = {}

        # Early out if not a valid attribute.
        if attribute is None:
            return result

        if isinstance(attribute, FnAttribute.DataAttribute):
            typeName = None
            if isinstance(attribute, FnAttribute.IntAttribute):
                typeName = "IntAttr"
            elif isinstance(attribute, FnAttribute.FloatAttribute):
                typeName = "FloatAttr"
            elif isinstance(attribute, FnAttribute.DoubleAttribute):
                typeName = "DoubleAttr"
            elif isinstance(attribute, FnAttribute.StringAttribute):
                typeName = "StringAttr"

            if typeName is not None:
                result["type"] = typeName

                # Only add if not the default.
                if attribute.getTupleSize() != 1:
                    result["tupleSize"] = attribute.getTupleSize()

                # If only one value, just add it as a scalar. Also, if only
                # samples at 0.0 are available, don't add times: it helps
                # readability.
                samples = attribute.getSamples()
                if len(samples) == 1 and list(samples.keys())[0] == 0.0:
                    values = samples[0.0]
                    if len(values) == 1:
                        result["value"] = values[0]
                    else:
                        result["values"] = list(values)
                else:
                    result["samples"] = {}
                    for timeSample, values in samples.items():
                        if len(values) == 1:
                            result["samples"][str(timeSample)] = values[0]
                        else:
                            result["samples"][str(timeSample)] = list(values)
        elif isinstance(attribute, FnAttribute.GroupAttribute):
            result["type"] = "GroupAttr"

            # Only add if not the default.
            if not attribute.getGroupInherit():
                result["groupInherit"] = False

            # Recursively call this function to write the attributes for its
            # children
            children = {}
            for childName, childAttribute in attribute.childList():
                children[childName] = JsonLookFileBakeOutputFormat.attrToDict(
                    childAttribute)
            result["children"] = children

        return result

# Register the output format.
LookFileBakeAPI.RegisterOutputFormat(JsonLookFileBakeOutputFormat)