import json, sys, struct, numpy as np
from pathlib import Path

def load_glb(path):
    with open(path, 'rb') as f:
        # Header
        magic = f.read(4)
        if magic != b'glTF':
            raise ValueError("Not a GLB file")
        version = struct.unpack('I', f.read(4))[0]
        size = struct.unpack('I', f.read(4))[0]
        
        # JSON chunk
        json_len = struct.unpack('I', f.read(4))[0]
        json_type = f.read(4)
        if json_type != b'JSON':
            raise ValueError("No JSON chunk")
        gltf = json.loads(f.read(json_len).decode())
        
        # Binary chunk (if exists)
        binary = bytearray()
        if f.tell() < size:
            bin_len = struct.unpack('I', f.read(4))[0]
            bin_type = f.read(4)
            if bin_type == b'BIN\x00':
                binary = bytearray(f.read(bin_len))
        
        return gltf, binary

def save_glb(gltf, binary, path):
    json_str = json.dumps(gltf, indent=2)
    json_bytes = json_str.encode('utf-8')
    # Pad to 4 bytes
    json_pad = (4 - len(json_bytes) % 4) % 4
    json_bytes += b' ' * json_pad
    
    bin_pad = (4 - len(binary) % 4) % 4
    binary += b'\x00' * bin_pad
    
    total = 12 + 8 + len(json_bytes) + 8 + len(binary)
    
    with open(path, 'wb') as f:
        f.write(b'glTF')
        f.write(struct.pack('I', 2))
        f.write(struct.pack('I', total))
        f.write(struct.pack('I', len(json_bytes)))
        f.write(b'JSON')
        f.write(json_bytes)
        f.write(struct.pack('I', len(binary)))
        f.write(b'BIN\x00')
        f.write(binary)

def convert(input_path):
    gltf, buf = load_glb(input_path)
    out_path = Path(input_path).stem + "_godot.glb"
    
    for mesh in gltf.get('meshes', []):
        for prim in mesh.get('primitives', []):
            attrs = prim.get('attributes', {})
            
            # Find color attributes
            colors = {}
            for name, idx in attrs.items():
                if '_COLOR' in name.upper():
                    num = int(name.split('_COLOR')[1])
                    colors[num] = idx
            
            if not colors:
                continue
            
            print(f"Found colors: {sorted(colors.keys())}")
            
            # Remove old color attributes
            for name in list(attrs.keys()):
                if '_COLOR' in name.upper():
                    del attrs[name]
            
            tex_idx = 2
            for num in sorted(colors.keys()):
                # Get accessor info
                acc = gltf['accessors'][colors[num]]
                bv = gltf['bufferViews'][acc['bufferView']]
                
                # Calculate offset in buffer
                offset = bv.get('byteOffset', 0) + acc.get('byteOffset', 0)
                
                # Read the color data
                if acc['type'] == 'VEC4':
                    count = acc['count']
                    data = np.frombuffer(buf[offset:offset + count*16], dtype=np.float32)
                    data = data.reshape(count, 4)
                else:  # VEC3
                    count = acc['count']
                    data = np.frombuffer(buf[offset:offset + count*12], dtype=np.float32)
                    data = data.reshape(count, 3)
                    data = np.column_stack([data, np.ones(count)])
                
                # Split into RG and BA
                rg = data[:, :2]
                ba = data[:, 2:]
                
                # Create RG TEXCOORD
                rg_acc = {
                    "bufferView": len(gltf.get('bufferViews', [])),
                    "componentType": 5126,
                    "count": count,
                    "type": "VEC2",
                    "max": rg.max(axis=0).tolist(),
                    "min": rg.min(axis=0).tolist()
                }
                rg_bytes = rg.astype(np.float32).tobytes()
                gltf.setdefault('bufferViews', []).append({
                    "buffer": 0,
                    "byteLength": len(rg_bytes),
                    "byteOffset": len(buf)
                })
                gltf['accessors'].append(rg_acc)
                buf.extend(rg_bytes)
                attrs[f'TEXCOORD_{tex_idx}'] = len(gltf['accessors']) - 1
                
                # Create BA TEXCOORD
                ba_acc = {
                    "bufferView": len(gltf['bufferViews']),
                    "componentType": 5126,
                    "count": count,
                    "type": "VEC2",
                    "max": ba.max(axis=0).tolist(),
                    "min": ba.min(axis=0).tolist()
                }
                ba_bytes = ba.astype(np.float32).tobytes()
                gltf['bufferViews'].append({
                    "buffer": 0,
                    "byteLength": len(ba_bytes),
                    "byteOffset": len(buf)
                })
                gltf['accessors'].append(ba_acc)
                buf.extend(ba_bytes)
                attrs[f'TEXCOORD_{tex_idx+1}'] = len(gltf['accessors']) - 1
                
                print(f"  _COLOR{num} → TEXCOORD_{tex_idx}+{tex_idx+1} → CUSTOM{(tex_idx-2)//2}")
                tex_idx += 2
            
            # Update buffer size
            gltf['buffers'][0]['byteLength'] = len(buf)
    
    save_glb(gltf, buf, out_path)
    print(f"\nSaved: {out_path}")
    print("In Godot shader: CUSTOM0, CUSTOM1, CUSTOM2, CUSTOM3 = your colors")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python script.py model.glb")
        sys.exit(1)
    convert(sys.argv[1])