#!/usr/bin/env python
"""
give a json element file, create pbrt scenes and objects. This should create the main element/element.pbrt,
element/geometry.pbrt and element/materials.pbrt. In addition, if the element has instanced primitives, then
for each primitive, a primitive subdirectory is created, along with a primitive.pbrt file and an archive_geometry
pbrt file for each archive referenced by that primitive description:

element/element.pbrt
element/element_geometry.pbrt
element/material.pbrt
element/objects.pbrt

element/primitive/primitive.pbrt
element/primitive/archive_geometry.pbrt

"""
import sys
import os
import argparse
import json
import pbrtutils as pu

ISLANDDIR = os.getcwd()


def parseArgs():

    parser = argparse.ArgumentParser(
        description="Takes an element and creates pbrt element version of the geometry and scene.")

    parser.add_argument('element', metavar='ELEM', help="Element to load.")
    parser.add_argument('--frameRange', type=str, default="1", help="Frame range to export, i.e. 1 or 1-3 or 1-5:2")
    parser.add_argument('-o', '--outdir', default=ISLANDDIR, help="Where to write out the pbrt files(s).")
    parser.add_argument('--skipObj', action='store_true', default=False, help="Skip Obj Conversion")
    parser.add_argument('--skipCopies', action='store_true', default=False, help="Skip Instanced Copies Obj Generation")
    parser.add_argument('--skipArchives', action='store_true', default=False, help="Skip Archive Conversion")
    parser.add_argument('--skipMat', action='store_true', default=False, help="Skip Material Generation")
    parser.add_argument('--skipPrims', action='store_true', default=False, help="Skip Primitive(s) PBRT Generation")
    parser.add_argument('--skipMain', action='store_true', default=False, help="Skip Main PBRT Generation")
    parser.add_argument('--prim', default="", help="Specific json primitive file to load, i.e. xgGrass")

    return parser.parse_args()


def main():
    args = parseArgs()

    if args.element == "island":
        # create the main island file.
        pu.writeMainPbrtFile()
        return

    jsonFile = os.path.join("json", args.element, args.element + ".json")
    with open(jsonFile, "r") as jf:
        jsonDict = json.load(jf)

    # pull data from dict
    element = jsonDict["name"]
    objFile = jsonDict["geomObjFile"]
    matFile = jsonDict["matFile"] if "matFile" in jsonDict else ""
    instancedPrimitives = jsonDict["instancedPrimitiveJsonFiles"] if "instancedPrimitiveJsonFiles" in jsonDict else {}
    instancedCopies = jsonDict["instancedCopies"] if "instancedCopies" in jsonDict else {}
    primDesc = args.prim

    archiveFiles = {}
    primitiveFiles = []

    print "Preparing pbrt export for element: ", element
    elementDir = os.path.join(args.outdir, "pbrt", element)

    # make all the directories we're going to need
    for outDir in [elementDir] + [os.path.join(args.outdir, "pbrt", element, x) for x in instancedPrimitives]:
        if not os.path.exists(outDir):
            os.makedirs(outDir)

    # relative path and full path names of the output files for pbrt.
    pbrtRelOutElement = os.path.join(element, element + ".pbrt")
    pbrtRelOutGeometry = os.path.join(element, element + "_geometry.pbrt")
    pbrtRelOutMaterials = os.path.join(element, "materials.pbrt")
    pbrtOutElement = os.path.join(elementDir, element + ".pbrt")
    pbrtOutGeometry = os.path.join(elementDir, element + "_geometry.pbrt")
    pbrtOutMaterials = os.path.join(elementDir, "materials.pbrt")

    ##############################################
    # Convert materials for the element --> element/materials.pbrt
    # -- we convert these first because we need the information to insert it into the OBJs.
    if not args.skipMat:
        print "################################################# Converting Materials: "
        pu.matToPbrt(element, matFile, pbrtOutMaterials)

    ##############################################
    # Convert geometry -- main file--> element/element_geometry.pbrt  (to be consistent w/ instances)
    if not args.skipObj:
        print "################################################# Converting main element OBJ file: "
        print "Exporting: ", pbrtOutGeometry
        pu.geomToPbrt(objFile, pbrtOutGeometry, matFile)

        ##############################################
        # Convert geometry #2 -- instanced copies with their own geometry (isPalmRig special case code....)
        if instancedCopies:
            print "################################################# Converting Instanced copies OBJ files: "
            for instance, icDict in instancedCopies.iteritems():
                if "geomObjFile" in icDict:
                    pbrtOutInstGeometry = os.path.join(elementDir, instance + "_geometry.pbrt")
                    print "Exporting: ", pbrtOutInstGeometry
                    pu.geomToPbrt(icDict['geomObjFile'], pbrtOutInstGeometry , matFile)

    ##############################################
    # convert geometry #3 - primitives json files and associated archives --> element/primitive/archive_geometry.pbrt
    # note: we don't need to run this for instanced copies -- we're assuming they use the same primitives.
    if not args.skipArchives or not instancedPrimitives:
        print "################################################# Converting Instanced JSON files and OBJ archives: "
        for primitiveName, ipDict in instancedPrimitives.iteritems():
            # this file is the instance name and transforms of each archive, keyed by archive (obj)
            if "archives" in ipDict:
                print "Create geometry data for all archives associated with ", primitiveName
                archiveFiles[primitiveName] = []
                for archive in ipDict["archives"]:
                    # this is a list of all the archives referenced above.
                    base = os.path.basename(archive)[:-4]
                    pbrtOutArchiveGeometry = os.path.join(elementDir, primitiveName, base + "_geometry.pbrt")
                    if not os.path.exists(pbrtOutArchiveGeometry):
                        print "Exporting: ", pbrtOutArchiveGeometry
                        if ipDict["type"] == "archive":
                            pu.geomToPbrt(archive, pbrtOutArchiveGeometry, matFile)
                        else:
                            variantMatFile = os.path.join(args.outdir, "json", ipDict["element"],"materials.json" )
                            print "...Exporting variant with alternate material file: %s" % variantMatFile
                            pu.geomToPbrt(archive, pbrtOutArchiveGeometry, variantMatFile)
                        archiveFiles[primitiveName].append(pbrtOutArchiveGeometry)

    else:
        # If we skipped creating archives, re-build the list here here.
        if instancedPrimitives:
            for primitiveName, iDict in instancedPrimitives.iteritems():
                archiveFiles[primitiveName] = [os.path.join(element, primitiveName, x)
                                               for x in os.listdir(os.path.join(elementDir, primitiveName))
                                               if x.endswith("geometry.pbrt")]

    if not args.skipPrims:
        ##############################################
        # Create primitives #1: element/primitive/instance_primitive.pbrt. This creates either
        # an archive type or curve type of file. All created files are added to a list for later
        # referencing in the main element pbrt file.
        print "################################################# Creating Primitive reference files: "
        for primitiveName, ipDict in instancedPrimitives.iteritems():
            if primDesc and primDesc != primitiveName:
                continue
            # this file is the instance name and transforms of each archive, keyed by archive (obj)
            print "Create main primitive file for:", primitiveName, "using", ipDict["jsonFile"]
            primitiveFiles.append(pu.jsonPrimitiveToPbrt(ipDict["type"] if "type" in ipDict else "", element, primitiveName, ipDict))

            #############################################
            # Create primitve #3 -- for all 'element' archive types, create any included primitives associated with
            # them, i.e. - isBayCedarA1_bonsaiA_xgBonsai
            if "type" in ipDict and ipDict["type"] == "element":
                sourceDict = pu.readJsonFile(os.path.join("json", ipDict["element"], ipDict["element"]+".json"))
                for variant, vDict in sourceDict["variants"].iteritems():
                    if "instancedPrimitiveJsonFiles" not in vDict:
                        continue
                    for primitiveName, vipDict in vDict["instancedPrimitiveJsonFiles"].iteritems():
                        primitiveFiles.append(pu.jsonPrimitiveToPbrt(vipDict["type"], ipDict["element"], primitiveName, vipDict))

            ##############################################
            # Create primitives #3 -- instanced copies with their own primitives (i,e, - isPalmRig3 w xgFronds3)
            if instancedCopies:
                print "################################################# Creating instanced copies primitives files: "
                for instance, icDict in instancedCopies.iteritems():
                    if "instancedPrimitiveJsonFiles" in icDict:
                        for primitiveName, ipDict in icDict["instancedPrimitiveJsonFiles"].iteritems():
                            print "Create primitive file for:", primitiveName, "using", ipDict["jsonFile"]
                            primitiveFiles.append(pu.jsonPrimitiveToPbrt(ipDict["type"] if "type" in ipDict else "", element, primitiveName, ipDict))

    else:
        # If we skipped the conversion, build primitiveFiles list here.
        for primitiveName in instancedPrimitives:
            primitiveFiles += [os.path.join(element, primitiveName, x)
                                           for x in os.listdir(os.path.join(elementDir, primitiveName))
                                           if x.endswith(".pbrt") and
                                              not x.endswith("geometry.pbrt")]

    if not args.skipMain:
        ##############################################
        # Create the main pbrt reference file with builds the hierarchy via includes and transform.
        # --> element/element.pbrt
        print "################################################# Creating main element reference: "
        print "Exporting:", pbrtOutElement
        pu.writeMainElement(pbrtOutElement, pbrtRelOutMaterials, pbrtRelOutGeometry, primitiveFiles,
                            sourceJsonFile=jsonFile)


if __name__ == "__main__":
    if main():
        sys.exit(0)
    sys.exit(1)