#!/bin/python3

SCRIPT_VERSION = '1.1'

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

try:
    from tqdm import tqdm as tqdmfunc
except ImportError:
    print("tqdm not installed, not showing progress bar")
    def tqdmfunc(iterator):
        yield from iterator

out_dir = None # Set in main()

dupes = set()

# logging.getLogger().setLevel(logging.INFO)

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)
            
            # Remove storables with just an id and no properties
            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:
            logging.info("-> %s", 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:
        # logging.warn("Skip duplicate %s", 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():
                logging.debug("%s:%s", var, scenejsonpath)
                
                __, 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_{self_name}.{sjp_plain}{atom['id'].replace('Person', '')}".replace('#', '').replace('/', '-')

                extractFromSceneJson(scenejson, outnameFn, readthumb)
                
    except json.JSONDecodeError:
        print(var)
        traceback.print_exc()
        os.rename(var, var + '.invalid')
    except zipfile.BadZipFile:
        print(var)
        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):
    # print([atom['id'] for atom in scenejson['atoms'] if atom['type'] == 'Person'])

    for atom, storables in getCleanedPersonAtoms(scenejson):
        # print(atom['id'])
        appearance = {
            "setUnlistedParamsToDefault": "true",
            'storables': storables
        }

        saveAppearance(appearance, outnameFn(atom), readthumb)


def main():
    if not sys.argv[1:]:
        print('Appearance Preset Extractor', SCRIPT_VERSION, 'by cartitolaire')
        print('')
        print('This tool automatically export presets from existing scenes and collects them in an "extracted" folder under presets. The presets will use the thumbnail for the source scene.')
        print('')
        print("Example usages:")
        print(' ', sys.argv[0], 'AddonPackages/Anon.SceneName.1.var')
        print(' ', sys.argv[0], '"AddonPackages/Anon*.var"')
        print(' ', sys.argv[0], '"Saves/scene/*.json" "AddonPackages/*.var"')
        return
    
    global out_dir
    if not os.path.isdir("./Custom/Atom/Person/"):
        logging.error("Script is not being run from the VaM root directory!")
        logging.warn("Can't cleanly merge into Person/Appearance presets; making ExtractedAppearance folder instead.")
        out_dir = "./ExtractedAppearance/"
    else:
        out_dir = "./Custom/Atom/Person/Appearance/extracted/"
        
    os.makedirs(out_dir, exist_ok=True)
    
    try:
        iterable = sum((glob.glob(a, recursive=True) for a in sys.argv[1:]), [])
        for filepath in tqdmfunc(iterable):
            __, 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

if __name__ == "__main__":
    main()