#!/bin/python3

import json
import sys
import os
import zipfile
import re
import glob
import traceback
import tqdm

out_dir = "Custom/Atom/Person/Appearance/extracted/"
os.makedirs(out_dir, exist_ok=True)

dupes = set()


def getCleanedPersonAtoms(scenejson):
    for atom in scenejson['atoms']:
        if atom['type'] != "Person":
            continue
        storables = atom['storables']
        for storable in [*storables]:
            storable.pop("position", None)
            storable.pop("rotation", None)
            # Remove all keys containing these substrings
            for key in [*storable.keys()]:
                if any(ss in key.lower() for ss in ['position', 'rotation']):
                    storable.pop(key)
            # Remove all storables whose ids contain these substrings
            if any(ss in storable['id'].lower() for ss in ["control", "trigger", "plugin", "preset", "animation"]):
                # print("pop", storable['id'])
                storables.remove(storable)
                continue
            # Remove transient morphs
            if storable['id'] == "geometry" and 'morphs' in storable:
                for morph in [*storable['morphs']]:
                    if any(re.match(ss, morph.get('uid', '')) for ss in [
                        r'Breast Impact',
                        r'^Eyelids (Top|Bottom) (Down|Up) (Left|Right)$', r'Brow .*(Up|Down)', r'^(Left|Right) Fingers'
                        r'^Mouth Open', r'Tongue In-Out', 'Shock', 'Surprise', 'Fear', 'Pain', 'Concentrate', 'Eyes Closed'
                    ]):
                        # print(morph)
                        storable['morphs'].remove(morph)
            if len(storable.keys()) == 1:
                # print("pop", storable['id'])
                storables.remove(storable)
                continue
            yield (atom, storables)


def saveAppearance(appearance, outname, readthumb):
    appearance_hash = json.dumps(appearance, sort_keys=True)
    if appearance_hash not in dupes:
        outpath = os.path.join(out_dir, outname + '.vap')
        with open(outpath, 'w') as fp:
            print("->", outpath)
            fp.write(json.dumps(appearance, indent=4))
        try:
            with open(os.path.join(out_dir, outname + '.jpg'), 'wb') as fp:
                with readthumb() as fp2:
                    fp.write(fp2.read())
        except (KeyError, FileNotFoundError):
            # No image in archive
            pass
        dupes.add(appearance_hash)
    else:
        print("Skip", outname)
        pass


def extractFromVar(var):
    __, filename = os.path.split(var)
    self_name, __ = os.path.splitext(filename)

    author = self_name.split('.')[0]

    try:
        with zipfile.ZipFile(var) as varzip:
            infodict = {
                zi.filename: zi
                for zi in varzip.infolist()
            }
            zip_scenes = {
                k: v for k, v in infodict.items()
                if k.startswith("Saves/scene") and k.endswith(".json")
            }
            for scenejsonpath in zip_scenes.keys():
                __, sjp_filename = os.path.split(scenejsonpath)
                sjp_plain, sjp_ext = os.path.splitext(sjp_filename)

                with varzip.open(scenejsonpath, 'r') as fp:
                    scenejson = json.load(fp)

                def readthumb():
                    return varzip.open(scenejsonpath.replace('.json', '.jpg'), 'r')

                def outnameFn(atom):
                    return f"Preset_{author}.{sjp_plain}{atom['id'].replace('Person', '')}".replace('#', '').replace('/', '-')

                extractFromSceneJson(scenejson, outnameFn, readthumb)

    except zipfile.BadZipFile:
        traceback.print_exc()
        os.rename(var, var + '.invalid')


def extractFromSceneJsonPath(scenejsonpath):
    __, sjp_filename = os.path.split(scenejsonpath)
    sjp_plain, sjp_ext = os.path.splitext(sjp_filename)

    def readthumb():
        return open(scenejsonpath.replace('.json', '.jpg'), 'rb')

    def outnameFn(atom):
        return f"Preset_!local.{sjp_plain}{atom['id'].replace('Person', '')}".replace('#', '').replace('/', '-')

    with open(scenejsonpath, 'r') as fp:
        extractFromSceneJson(json.load(fp), outnameFn, readthumb)


def extractFromSceneJson(scenejson, outnameFn, readthumb):
    for atom, storables in getCleanedPersonAtoms(scenejson):
        appearance = {
            "setUnlistedParamsToDefault": "true",
            'storables': storables
        }

        saveAppearance(appearance, outnameFn(atom), readthumb)


def main():
    try:
        for filepath in tqdm.tqdm(sum((glob.glob(a, recursive=True) for a in sys.argv[1:]), [])):
            __, filename = os.path.split(filepath)
            self_name, ext = os.path.splitext(filename)

            if ext == ".var":
                extractFromVar(filepath)
            elif ext == ".json":
                extractFromSceneJsonPath(filepath)
            else:
                raise NotImplementedError(ext)

    except KeyboardInterrupt:
        return


main()