import base64
import json
import re
import os
import shutil
import sys
import zipfile

try:
    import zlib

    compression = zipfile.ZIP_DEFLATED
except:
    compression = zipfile.ZIP_STORED

# -----------------------------------------------------------------------------


def convertDirectoryToZipFile(directoryPath):
    if os.path.isfile(directoryPath):
        return

    zipFilePath = "%s.zip" % directoryPath
    zf = zipfile.ZipFile(zipFilePath, mode="w")

    try:
        for dirName, subdirList, fileList in os.walk(directoryPath):
            for fname in fileList:
                fullPath = os.path.join(dirName, fname)
                relPath = "%s" % (os.path.relpath(fullPath, directoryPath))
                zf.write(fullPath, arcname=relPath, compress_type=compression)
    finally:
        zf.close()

    shutil.rmtree(directoryPath)
    shutil.move(zipFilePath, directoryPath)


# -----------------------------------------------------------------------------


def addDataToViewer(dataPath, srcHtmlPath):
    if os.path.isfile(dataPath) and os.path.exists(srcHtmlPath):
        dstDir = os.path.dirname(dataPath)
        dstHtmlPath = os.path.join(dstDir, "%s.html" % os.path.basename(dataPath)[:-6])

        # Extract data as base64
        with open(dataPath, "rb") as data:
            dataContent = data.read()
            base64Content = base64.b64encode(dataContent)
            base64Content = base64Content.decode().replace("\n", "")

        # Create new output file
        with open(srcHtmlPath, mode="r", encoding="utf-8") as srcHtml:
            with open(dstHtmlPath, mode="w", encoding="utf-8") as dstHtml:
                for line in srcHtml:
                    if "</body>" in line:
                        dstHtml.write("<script>\n")
                        dstHtml.write('var contentToLoad = "%s";\n\n' % base64Content)
                        dstHtml.write(
                            'Glance.importBase64Dataset("%s" , contentToLoad, glanceInstance.proxyManager);\n'
                            % os.path.basename(dataPath)
                        )
                        dstHtml.write("glanceInstance.showApp();\n")
                        dstHtml.write("</script>\n")

                    dstHtml.write(line)


# -----------------------------------------------------------------------------


def zipAllTimeSteps(directoryPath):
    if os.path.isfile(directoryPath):
        return

    class UrlCounterDict(dict):
        Counter = 0

        def GetUrlName(self, name):
            if name not in self.keys():
                self[name] = str(objNameToUrls.Counter)
                self.Counter = self.Counter + 1
            return self[name]

    def InitIndex(sourcePath, destObj):
        with open(sourcePath, "r") as sourceFile:
            sourceData = sourceFile.read()
            sourceObj = json.loads(sourceData)
            for key in sourceObj:
                destObj[key] = sourceObj[key]
            # remove vtkHttpDataSetReader information
            for obj in destObj["scene"]:
                obj.pop(obj["type"])
                obj.pop("type")

    def getUrlToNameDictionary(indexObj):
        urls = {}
        for obj in indexObj["scene"]:
            urls[obj[obj["type"]]["url"]] = obj["name"]
        return urls

    def addDirectoryToZip(
        dirname, zipobj, storedData, rootIdx, timeStep, objNameToUrls
    ):
        # Update root index.json file from index.json of this timestep
        with open(os.path.join(dirname, "index.json"), "r") as currentIdxFile:
            currentIdx = json.loads(currentIdxFile.read())
            urlToName = getUrlToNameDictionary(currentIdx)
            rootTimeStepSection = rootIdx["animation"]["timeSteps"][timeStep]
            for key in currentIdx:
                if key == "scene" or key == "version":
                    continue
                rootTimeStepSection[key] = currentIdx[key]
            for obj in currentIdx["scene"]:
                objName = obj["name"]
                rootTimeStepSection[objName] = {}
                rootTimeStepSection[objName]["actor"] = obj["actor"]
                rootTimeStepSection[objName]["actorRotation"] = obj["actorRotation"]
                rootTimeStepSection[objName]["mapper"] = obj["mapper"]
                rootTimeStepSection[objName]["property"] = obj["property"]

        # For every object in the current timestep
        for folder in sorted(os.listdir(dirname)):
            currentItem = os.path.join(dirname, folder)
            if os.path.isdir(currentItem) is False:
                continue
            # Write all data array of the current timestep in the archive
            for filename in os.listdir(os.path.join(currentItem, "data")):
                fullpath = os.path.join(currentItem, "data", filename)
                if os.path.isfile(fullpath) and filename not in storedData:
                    storedData.add(filename)
                    relPath = os.path.join("data", filename)
                    zipobj.write(fullpath, arcname=relPath, compress_type=compression)
            # Write the index.json containing pointers to these data arrays
            # while replacing every basepath as '../../data'
            objIndexFilePath = os.path.join(dirname, folder, "index.json")
            with open(objIndexFilePath, "r") as objIndexFile:
                objIndexObjData = json.loads(objIndexFile.read())
            for elm in objIndexObjData.keys():
                try:
                    if "ref" in objIndexObjData[elm].keys():
                        objIndexObjData[elm]["ref"]["basepath"] = "../../data"
                    if "arrays" in objIndexObjData[elm].keys():
                        for array in objIndexObjData[elm]["arrays"]:
                            array["data"]["ref"]["basepath"] = "../../data"
                except AttributeError:
                    continue
            currentObjName = urlToName[folder]
            objIndexRelPath = os.path.join(
                objNameToUrls.GetUrlName(currentObjName), str(timeStep), "index.json"
            )
            zipobj.writestr(
                objIndexRelPath,
                json.dumps(objIndexObjData, indent=2),
                compress_type=compression,
            )

    # ---

    zipFilePath = "%s.zip" % directoryPath
    currentDirectory = os.path.abspath(os.path.join(directoryPath, os.pardir))
    rootIndexPath = os.path.join(currentDirectory, "index.json")
    rootIndexFile = open(rootIndexPath, "r")
    rootIndexObj = json.loads(rootIndexFile.read())

    zf = zipfile.ZipFile(zipFilePath, mode="w")
    try:
        # We copy the scene from an index of a specific timestep to the root index
        # Scenes should all have the same objects so only do it for the first one
        isSceneInitialized = False
        # currentlyAddedData set stores hashes of every data we already added to the
        # vtkjs archive to prevent data duplication
        currentlyAddedData = set()
        # Regex that folders storing timestep data from paraview should follow
        reg = re.compile(r"^" + os.path.basename(directoryPath) + r"\.[0-9]+$")
        # We assume an object will not be deleted from a timestep to another so we create a generic index.json for each object
        genericIndexObj = {}
        genericIndexObj["series"] = []
        timeStep = 0
        for item in rootIndexObj["animation"]["timeSteps"]:
            genericIndexObj["series"].append({})
            genericIndexObj["series"][timeStep]["url"] = str(timeStep)
            genericIndexObj["series"][timeStep]["timeStep"] = float(item["time"])
            timeStep = timeStep + 1
        # Keep track of the url for every object
        objNameToUrls = UrlCounterDict()

        timeStep = 0
        # zip all timestep directories
        for folder in sorted(os.listdir(currentDirectory)):
            fullPath = os.path.join(currentDirectory, folder)
            if os.path.isdir(fullPath) and reg.match(folder):
                if not isSceneInitialized:
                    InitIndex(os.path.join(fullPath, "index.json"), rootIndexObj)
                    isSceneInitialized = True
                addDirectoryToZip(
                    fullPath,
                    zf,
                    currentlyAddedData,
                    rootIndexObj,
                    timeStep,
                    objNameToUrls,
                )
                shutil.rmtree(fullPath)
                timeStep = timeStep + 1

        # Write every index.json holding time information for each object
        for name in objNameToUrls:
            zf.writestr(
                os.path.join(objNameToUrls[name], "index.json"),
                json.dumps(genericIndexObj, indent=2),
                compress_type=compression,
            )

        # Update root index.json urls and write it in the archive
        for obj in rootIndexObj["scene"]:
            obj["id"] = obj["name"]
            obj["type"] = "vtkHttpDataSetSeriesReader"
            obj["vtkHttpDataSetSeriesReader"] = {}
            obj["vtkHttpDataSetSeriesReader"]["url"] = objNameToUrls[obj["name"]]
        zf.writestr(
            "index.json", json.dumps(rootIndexObj, indent=2), compress_type=compression
        )
        os.remove(rootIndexPath)

    finally:
        zf.close()

    shutil.move(zipFilePath, directoryPath)


# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print(
            "Usage: directoryToFile /path/to/directory.vtkjs [/path/to/ParaViewGlance.html]"
        )
    else:
        fileName = sys.argv[1]
        convertDirectoryToZipFile(fileName)

        if len(sys.argv) == 3:
            addDataToViewer(fileName, sys.argv[2])