# Embeds a .glb's externally referenced images into the file itself, which is what Singe's loader
# needs:  it reads images from buffer views and will not go looking for a file beside the model.
# Downloaded packs often ship a .glb that still points at a shared texture atlas (Kenney's arcade
# models name "Textures/colormap.png"), and those load untextured until this is run over them.
# packGlb.py does the same job for a .gltf with external buffers;  this one leaves the buffer alone
# and only pulls the images in.
# Usage:  python3 util/embedGlbImages.py in.glb out.glb [--base DIRECTORY]
import argparse
import json
import os
import struct
import sys

mimeTypes = {'png': 'image/png', 'jpg': 'image/jpeg', 'jpeg': 'image/jpeg', 'ktx2': 'image/ktx2', 'webp': 'image/webp', 'basis': 'image/basis'}

parser = argparse.ArgumentParser()
parser.add_argument('source')
parser.add_argument('output')
parser.add_argument('--base', default=None, help='where the external images live (default: beside the source)')
args = parser.parse_args()
base = args.base if args.base else os.path.dirname(args.source)

with open(args.source, 'rb') as f:
    magic, version, length = struct.unpack('<III', f.read(12))
    if magic != 0x46546C67:
        sys.exit('%s is not a .glb' % args.source)
    jsonLength, jsonType = struct.unpack('<II', f.read(8))
    gltf = json.loads(f.read(jsonLength).decode('utf-8'))
    binary = bytearray()
    rest = f.read()
    if len(rest) >= 8:
        binLength, binType = struct.unpack('<II', rest[:8])
        binary = bytearray(rest[8:8 + binLength])

views = gltf.setdefault('bufferViews', [])
embedded = 0
for image in gltf.get('images', []):
    uri = image.get('uri')
    if uri is None or uri.startswith('data:'):
        continue
    path = os.path.join(base, uri)
    if not os.path.exists(path):
        sys.exit('%s references %s, which is not beside it' % (args.source, uri))
    with open(path, 'rb') as f:
        data = f.read()
    while len(binary) % 4:
        binary.append(0)
    views.append({'buffer': 0, 'byteOffset': len(binary), 'byteLength': len(data)})
    binary.extend(data)
    image['bufferView'] = len(views) - 1
    image['mimeType'] = mimeTypes.get(uri.rsplit('.', 1)[-1].lower(), 'image/png')
    del image['uri']
    embedded += 1

gltf['buffers'] = [{'byteLength': len(binary)}]
jsonChunk = json.dumps(gltf, separators=(',', ':')).encode('utf-8')
while len(jsonChunk) % 4:
    jsonChunk += b' '
while len(binary) % 4:
    binary.append(0)
with open(args.output, 'wb') as f:
    f.write(struct.pack('<III', 0x46546C67, 2, 12 + 8 + len(jsonChunk) + 8 + len(binary)))
    f.write(struct.pack('<II', len(jsonChunk), 0x4E4F534A))
    f.write(jsonChunk)
    f.write(struct.pack('<II', len(binary), 0x004E4942))
    f.write(binary)
print('%s:  %d image(s) embedded, %d bytes' % (args.output, embedded, os.path.getsize(args.output)))