# Compresses a game's textures to KTX2 (Basis Universal UASTC with Zstandard, full mipmap chains)
# with the basisu command-line tool, which the Basis Universal repository builds; toktx from the
# KTX-Software tools does the same job.  Given a folder, every PNG and JPEG in it gets a .ktx2
# beside it.  Given a .gltf, its images are converted, the file rewritten to reference them through
# KHR_texture_basisu (data maps marked linear), and util/packGlb.py then packs it as usual.
# Usage:  python3 util/compressTextures.py [--basisu PATH] [--etc1s] FOLDER-or-Model.gltf
import json, os, shutil, subprocess, sys

args   = sys.argv[1:]
basisu = shutil.which('basisu')
etc1s  = False
while args and args[0].startswith('--'):
    if args[0] == '--basisu':
        basisu = args[1]
        args   = args[2:]
    elif args[0] == '--etc1s':
        etc1s = True
        args  = args[1:]
    else:
        sys.exit('Unknown option ' + args[0])
if len(args) != 1:
    sys.exit(__doc__ or 'Usage:  python3 util/compressTextures.py [--basisu PATH] [--etc1s] FOLDER-or-Model.gltf')
if basisu is None:
    sys.exit('The basisu tool was not found.  Build it from https://github.com/BinomialLLC/basis_universal (cmake, then the basisu target) and pass its path with --basisu, or install KTX-Software for toktx.')
target = args[0]

def convert(source, linear):
    """One image to KTX2 beside it; returns the new file name."""
    out = os.path.splitext(source)[0] + '.ktx2'
    cmd = [basisu, '-ktx2', '-mipmap', '-output_file', out]
    if not etc1s:
        cmd.append('-uastc')
    if linear:
        cmd.append('-linear')
    cmd.append(source)
    print('  ' + os.path.basename(source) + ' -> ' + os.path.basename(out) + (' (linear)' if linear else ''))
    subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL)
    return out

if target.lower().endswith('.gltf'):
    folder = os.path.dirname(target)
    with open(target) as f:
        gltf = json.load(f)
    # Which images carry colour and which carry data, from the materials that use them.
    linearImages = set()
    for material in gltf.get('materials', []):
        for key in ('normalTexture', 'occlusionTexture'):
            if key in material:
                linearImages.add(gltf['textures'][material[key]['index']].get('source'))
        pbr = material.get('pbrMetallicRoughness', {})
        if 'metallicRoughnessTexture' in pbr:
            linearImages.add(gltf['textures'][pbr['metallicRoughnessTexture']['index']].get('source'))
    for index, image in enumerate(gltf.get('images', [])):
        if 'uri' not in image or image['uri'].startswith('data:'):
            print('  image %d is embedded; skipped' % index)
            continue
        source = os.path.join(folder, image['uri'])
        if source.lower().endswith('.ktx2'):
            continue
        out = convert(source, index in linearImages)
        image['uri']      = os.path.relpath(out, folder).replace(os.sep, '/')
        image['mimeType'] = 'image/ktx2'
    for texture in gltf.get('textures', []):
        if 'source' in texture and 'extensions' not in texture:
            texture['extensions'] = {'KHR_texture_basisu': {'source': texture.pop('source')}}
    for key in ('extensionsUsed', 'extensionsRequired'):
        gltf.setdefault(key, [])
        if 'KHR_texture_basisu' not in gltf[key]:
            gltf[key].append('KHR_texture_basisu')
    with open(target, 'w') as f:
        json.dump(gltf, f, indent=1)
    print('Rewrote ' + target)
else:
    for name in sorted(os.listdir(target)):
        if name.lower().endswith(('.png', '.jpg', '.jpeg')):
            convert(os.path.join(target, name), False)