#!/usr/bin/env python3
"""
glb_strip_utility_nodes.py
Usage: python3 tooling/glb_strip_utility_nodes.py <input.glb> [<output.glb>]

Removes utility mesh nodes from a GLB's GLTF scene graph.
A "utility node" is any mesh node where:
  - The mesh has 0 primitives, OR
  - The mesh name matches a known utility pattern (Icosphere, WGT-*, etc.)
  - The mesh has no skinning joint data (vgroups=0 in Blender terms → no skin references)

Strips the node from the scene graph (and its parent's children list) without
touching the underlying mesh/accessor data. The mesh data becomes orphaned
(not referenced by any node) and most GLTF consumers (Godot, Blender) ignore it.

GLB format:
  - 12-byte header: magic (0x46546C67), version (2), total length
  - Chunk 0: JSON (type 0x4E4F534A)
  - Chunk 1: BIN  (type 0x004E4942)

If no output path is given, overwrites the input file.
"""

import json
import struct
import sys
import os
import re

GLB_MAGIC   = 0x46546C67
CHUNK_JSON  = 0x4E4F534A
CHUNK_BIN   = 0x004E4942

# Mesh names considered utility/non-geometry
UTILITY_PATTERNS = [
    re.compile(r'^Icosphere', re.IGNORECASE),
    re.compile(r'^WGT-', re.IGNORECASE),
    re.compile(r'^DEF-', re.IGNORECASE),
    re.compile(r'^ORG-', re.IGNORECASE),
]


def is_utility_mesh(mesh: dict) -> bool:
    name = mesh.get('name', '')
    for pat in UTILITY_PATTERNS:
        if pat.match(name):
            return True
    # Also treat meshes with no primitives as utility
    if not mesh.get('primitives'):
        return True
    return False


def read_glb(path: str):
    """Read a GLB file and return (gltf_dict, bin_data)."""
    with open(path, 'rb') as f:
        data = f.read()

    magic, version, length = struct.unpack_from('<III', data, 0)
    assert magic == GLB_MAGIC, f"Not a GLB file (magic={magic:#010x})"
    assert version == 2, f"Unexpected GLB version {version}"

    offset = 12
    gltf_json = None
    bin_data  = b''

    while offset < length:
        chunk_len, chunk_type = struct.unpack_from('<II', data, offset)
        offset += 8
        chunk_data = data[offset: offset + chunk_len]
        offset += chunk_len

        if chunk_type == CHUNK_JSON:
            gltf_json = json.loads(chunk_data.decode('utf-8'))
        elif chunk_type == CHUNK_BIN:
            bin_data = chunk_data

    assert gltf_json is not None, "No JSON chunk found in GLB"
    return gltf_json, bin_data


def write_glb(gltf_dict: dict, bin_data: bytes, path: str):
    """Write a GLB file from a GLTF dict and binary data."""
    json_bytes = json.dumps(gltf_dict, separators=(',', ':'), ensure_ascii=False).encode('utf-8')
    # Pad JSON to 4-byte alignment with spaces
    pad = (4 - len(json_bytes) % 4) % 4
    json_bytes += b' ' * pad

    # Pad binary to 4-byte alignment with zeros
    bin_pad = (4 - len(bin_data) % 4) % 4 if bin_data else 0
    bin_padded = bin_data + b'\x00' * bin_pad

    json_chunk = struct.pack('<II', len(json_bytes), CHUNK_JSON) + json_bytes
    bin_chunk  = (struct.pack('<II', len(bin_padded), CHUNK_BIN) + bin_padded) if bin_data else b''

    total_length = 12 + len(json_chunk) + len(bin_chunk)
    header = struct.pack('<III', GLB_MAGIC, 2, total_length)

    with open(path, 'wb') as f:
        f.write(header + json_chunk + bin_chunk)


def strip_utility_nodes(gltf: dict) -> tuple[int, list[str]]:
    """
    Remove utility mesh nodes from the scene graph.
    Returns (count_removed, names_removed).
    """
    meshes = gltf.get('meshes', [])
    nodes  = gltf.get('nodes', [])
    scenes = gltf.get('scenes', [])

    # Find utility mesh indices
    utility_mesh_indices = {
        i for i, m in enumerate(meshes)
        if is_utility_mesh(m)
    }

    if not utility_mesh_indices:
        return 0, []

    # Find node indices that reference utility meshes
    utility_node_indices = {
        i for i, n in enumerate(nodes)
        if n.get('mesh') in utility_mesh_indices
    }

    if not utility_node_indices:
        return 0, []

    removed_names = [nodes[i].get('name', f'node_{i}') for i in sorted(utility_node_indices)]

    # Remove utility nodes from scene top-level node lists
    for scene in scenes:
        scene['nodes'] = [n for n in scene.get('nodes', []) if n not in utility_node_indices]

    # Remove utility nodes from every other node's children list
    for node in nodes:
        if 'children' in node:
            node['children'] = [c for c in node['children'] if c not in utility_node_indices]
            if not node['children']:
                del node['children']  # clean up empty children array

    # Note: we intentionally leave the mesh data and nodes in place as orphaned entries.
    # Reindexing would require updating all skin joint arrays and is complex.
    # Orphaned nodes/meshes are ignored by Godot and Blender on import.

    return len(utility_node_indices), removed_names


def process_file(input_path: str, output_path: str):
    print(f"Processing: {input_path}")
    gltf, bin_data = read_glb(input_path)

    count, names = strip_utility_nodes(gltf)

    if count == 0:
        print(f"  No utility nodes found — file unchanged")
        if output_path != input_path:
            import shutil
            shutil.copy2(input_path, output_path)
        return

    print(f"  Stripped {count} utility node(s): {names}")
    write_glb(gltf, bin_data, output_path)
    print(f"  Written: {output_path} ({os.path.getsize(output_path):,} bytes)")


def process_directory(input_dir: str):
    """Process all .glb files in a directory tree."""
    total = 0
    for root, dirs, files in os.walk(input_dir):
        for fname in files:
            if fname.lower().endswith('.glb'):
                fpath = os.path.join(root, fname)
                gltf, bin_data = read_glb(fpath)
                count, names = strip_utility_nodes(gltf)
                if count > 0:
                    write_glb(gltf, bin_data, fpath)
                    print(f"  {os.path.relpath(fpath, input_dir)}: stripped {count} nodes {names}")
                    total += count
    print(f"\nTotal utility nodes stripped: {total}")


if __name__ == '__main__':
    args = sys.argv[1:]
    if not args:
        print("Usage: python3 glb_strip_utility_nodes.py <input.glb> [<output.glb>]")
        print("       python3 glb_strip_utility_nodes.py --dir <directory>")
        sys.exit(1)

    if args[0] == '--dir':
        if len(args) < 2:
            print("ERROR: --dir requires a directory path")
            sys.exit(1)
        process_directory(args[1])
    else:
        input_path  = args[0]
        output_path = args[1] if len(args) > 1 else args[0]
        process_file(input_path, output_path)