#!/bin/python3

SCRIPT_VERSION = '1.3.1'

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

logging.getLogger().setLevel(logging.WARNING)
logging.getLogger().setLevel(logging.INFO)

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

out_dir = None # Set in main()

dupes = set()

def update_dict(d, fn, key=None):
    if isinstance(d, dict):
        for k, v in d.items():
            d[k] = update_dict(v, fn, k)
        # return not needed to update dicts by reference, but used
        # here to assist in assigning subtrees
        return d
    elif isinstance(d, list):
        return [update_dict(i, fn) for i in d]
    else:
        return fn(d, key)

def update_dict_selfname(d, self_name):
    def update_fn(v, key=None):
        if key == 'id' and isinstance(v, str):
            return v.replace("SELF:/", self_name + ":/")
        else:
            return v
            
    update_dict(d, update_fn)

def getPersonAtoms(scenejson):
    for atom in scenejson['atoms']:
        if atom['type'] == "Person":
            yield atom
            
def cleanPersonStorables(atom):
    storables = atom['storables']
    storables_new = []
    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:
            storables_new.append(storable)
            
    return storables_new


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', encoding='utf-8') 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)
                    
                if False:
                    update_dict_selfname(scenejson, self_name)
                else:
                    # This method is potentially inaccurate but MUCH faster:
                    scenejson = json.loads(json.dumps(scenejson).replace("SELF:/", self_name + ":/"))

                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, noclobber=True)
                
    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', encoding='utf-8') as fp:
        extractFromSceneJson(json.load(fp), outnameFn, readthumb)


def extractFromSceneJson(scenejson, outnameFn, readthumb, noclobber=False):
    # print([atom['id'] for atom in scenejson['atoms'] if atom['type'] == 'Person'])

    for atom in getPersonAtoms(scenejson):
        outname = outnameFn(atom)
        if noclobber and os.path.isfile(os.path.join(out_dir, outname + '.vap')):
            logging.info(f"Skipping {outname!r}")
            continue
    
        # print(atom['id'])
        appearance = {
            "setUnlistedParamsToDefault": "true",
            'storables': cleanPersonStorables(atom)
        }

        saveAppearance(appearance, outname, 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.warning("Can't cleanly merge into Person/Appearance presets; making ExtractedAppearance folder instead.")
        out_dir = "./ExtractedAppearance/"
    else:
        out_dir = "./Custom/Atom/Person/Appearance/extracted/"
        
    logging.info(f"Writing output to {out_dir!r}")
    
    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()