#!/usr/bin/env python3
"""Convert JME3-extracted JSON (from J3OConverter.java) to glTF 2.0 GLB with animations."""

import json
import struct
import math
import os
import sys

def quat_mult(qa, qb):
    ax, ay, az, aw = qa
    bx, by, bz, bw = qb
    return (
        aw*bx + ax*bw + ay*bz - az*by,
        aw*by - ax*bz + ay*bw + az*bx,
        aw*bz + ax*by - ay*bx + az*bw,
        aw*bw - ax*bx - ay*by - az*bz,
    )

def pad4(data):
    while len(data) % 4:
        data += b'\x00'
    return data

def mat4_mult(a, b):
    r = [0]*16
    for i in range(4):
        for j in range(4):
            for k in range(4):
                r[j*4 + i] += a[k*4 + i] * b[j*4 + k]
    return r

def mat4_from_trs(t, r, s):
    x, y, z, w = r
    xx, yy, zz = x*x, y*y, z*z
    xy, xz, yz = x*y, x*z, y*z
    wx, wy, wz = w*x, w*y, w*z
    return [
        (1-2*(yy+zz))*s[0], (2*(xy+wz))*s[0],   (2*(xz-wy))*s[0],   0,
        (2*(xy-wz))*s[1],   (1-2*(xx+zz))*s[1], (2*(yz+wx))*s[1],   0,
        (2*(xz+wy))*s[2],   (2*(yz-wx))*s[2],   (1-2*(xx+yy))*s[2], 0,
        t[0], t[1], t[2], 1,
    ]

def invert_4x4_row(m):
    a,b,c,d, e,f,g,h, i,j,k,l, m_,n,o,p = m
    det = (a*(f*k*p+g*l*n+h*j*o-h*k*n-f*l*o-g*j*p)
           - b*(e*k*p+g*l*m_+h*i*o-h*k*m_-e*l*o-g*i*p)
           + c*(e*j*p+f*l*m_+h*i*n-h*j*m_-e*l*n-f*i*p)
           - d*(e*j*o+f*k*m_+g*i*n-g*j*m_-e*k*n-f*i*o))
    if abs(det) < 1e-12:
        return [1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1]
    idet = 1.0/det
    return [
        (f*k*p+g*l*n+h*j*o-h*k*n-f*l*o-g*j*p)*idet,
        -(b*k*p+c*l*n+d*j*o-d*k*n-b*l*o-c*j*p)*idet,
        (b*g*p+c*h*n+d*f*o-d*g*n-b*h*o-c*f*p)*idet,
        -(b*g*l+c*h*j+d*f*k-d*g*j-b*h*k-c*f*l)*idet,
        -(e*k*p+g*l*m_+h*i*o-h*k*m_-e*l*o-g*i*p)*idet,
        (a*k*p+c*l*m_+d*i*o-d*k*m_-a*l*o-c*i*p)*idet,
        -(a*g*p+c*h*m_+d*e*o-d*g*m_-a*h*o-c*e*p)*idet,
        (a*g*l+c*h*i+d*e*k-d*g*i-a*h*k-c*e*l)*idet,
        (e*j*p+f*l*m_+h*i*n-h*j*m_-e*l*n-f*i*p)*idet,
        -(a*j*p+b*l*m_+d*i*n-d*j*m_-a*l*n-b*i*p)*idet,
        (a*f*p+b*h*m_+d*e*n-d*f*m_-a*h*n-b*e*p)*idet,
        -(a*f*l+b*h*i+d*e*j-d*f*i-a*h*j-b*e*l)*idet,
        -(e*j*o+f*k*m_+g*i*n-g*j*m_-e*k*n-f*i*o)*idet,
        (a*j*o+b*k*m_+c*i*n-c*j*m_-a*k*n-b*i*o)*idet,
        -(a*f*o+b*g*m_+c*e*n-c*f*m_-a*g*n-b*e*o)*idet,
        (a*f*k+b*g*i+c*e*j-c*f*i-a*g*j-b*e*k)*idet,
    ]

def convert(json_path, output_path):
    with open(json_path) as f:
        data = json.load(f)

    geos = data.get('geometries', [])
    skels = data.get('skeletons', [])
    anims = data.get('animations', {})

    if not geos:
        print(f'  No geometry in {json_path}')
        return

    geo = geos[0]
    positions = geo['positions']
    normals = geo.get('normals', [])
    indices = geo.get('indices', [])
    texcoords = geo.get('texcoords', [])
    vertex_count = geo.get('vertexCount', len(positions) // 3)

    bone_weights = geo.get('boneWeights', [])
    bone_indices = geo.get('boneIndices', [])

    skel = skels[0] if skels else None

    nodes = []
    skin = None
    animations_out = []

    # ── Skeleton → glTF nodes ──────────────────────────────────────────────
    if skel:
        bones = skel['bones']
        roots = skel['roots']
        name_to_bone = {b['name']: b for b in bones}
        name_to_idx = {}

        # IMPORTANT: Keep the ORIGINAL bone list order (from JME3 skeleton.getBone(i)).
        # BoneTrack.getTargetBoneIndex() and the mesh JOINTS_0 data reference
        # this exact order.
        order = [b['name'] for b in bones]
        for i, name in enumerate(order):
            name_to_idx[name] = i
            b = name_to_bone[name]
            nodes.append({
                'name': name,
                'translation': b['position'],
                'rotation': b['rotation'],  # x,y,z,w
                'scale': b['scale'],
            })

        # Hierarchy (by name)
        for name in order:
            b = name_to_bone[name]
            for child_name in b['children']:
                if child_name in name_to_idx and name in name_to_idx:
                    pi = name_to_idx[name]
                    ci = name_to_idx[child_name]
                    nodes[pi].setdefault('children', []).append(ci)

        # Inverse bind matrices
        def world_transform(name, cache={}):
            if name in cache:
                return cache[name]
            b = name_to_bone[name]
            local = mat4_from_trs(b['position'], b['rotation'], b['scale'])
            # Find parent
            parent = None
            for pn, pb in name_to_bone.items():
                if name in pb['children']:
                    parent = pn
                    break
            if parent and parent in name_to_bone:
                w = mat4_mult(world_transform(parent, cache), local)
            else:
                w = local
            cache[name] = w
            return w

        ibms = []
        for name in order:
            w = world_transform(name)
            # Convert col-major to row-major, invert, back to col-major
            row = [w[i + j*4] for j in range(4) for i in range(4)]
            inv_row = invert_4x4_row(row)
            inv_col = [inv_row[j + i*4] for i in range(4) for j in range(4)]
            ibms.extend(inv_col)

        # Joint indices (identity: bone list order == glTF node order)
        joints = list(range(len(order)))
        skin = {'joints': joints}

        # Mesh node with skin
        mesh_node_idx = len(nodes)
        nodes.append({'name': 'mesh', 'mesh': 0, 'skin': 0})

        # Skeleton root node
        root_node_idx = len(nodes)
        root_children = [name_to_idx[r] for r in roots if r in name_to_idx]
        root_children.append(mesh_node_idx)
        nodes.append({'name': 'skeleton_root', 'children': root_children})

        # Scene root
        scene_node_idx = len(nodes)
        nodes.append({'name': 'scene_root', 'children': [root_node_idx]})

        # ── Animations ──────────────────────────────────────────────────────
        anim_bin = bytearray()
        anim_buf_views = []
        anim_accessors = []

        for anim_name, anim in anims.items():
            tracks = anim.get('tracks', [])
            if not tracks:
                continue

            # Collect all keyframe times
            all_times = set()
            for t in tracks:
                for tm in t.get('times', []):
                    all_times.add(round(tm, 6))
            sorted_times = sorted(all_times)
            if len(sorted_times) < 2:
                continue

            # Per-animation sampler/channel lists
            anim_samplers_this = []
            anim_channels_this = []

            # Time buffer
            time_bytes = pad4(struct.pack(f'<{len(sorted_times)}f', *sorted_times))
            time_off = len(anim_bin)
            anim_bin.extend(time_bytes)
            tv_idx = len(anim_buf_views)
            anim_buf_views.append({'byteOffset': time_off, 'byteLength': len(time_bytes)})
            t_acc = len(anim_accessors)
            anim_accessors.append({'type': 'SCALAR', 'count': len(sorted_times),
                                   'min': [sorted_times[0]], 'max': [sorted_times[-1]]})

            for track in tracks:
                bone_idx = track.get('boneIndex', 0)
                if bone_idx >= len(order):
                    continue
                node_idx = name_to_idx.get(order[bone_idx], None)
                if node_idx is None:
                    node_idx = bone_idx

                # Get bind pose for this bone (JME3 tracks are relative to bind)
                bone_name = order[bone_idx]
                bind = name_to_bone[bone_name]
                bind_pos = bind['position']
                bind_rot = bind['rotation']
                bind_scale = bind['scale']

                kf_times = list(track.get('times', []))
                kf_trans = list(track.get('translations', []))
                kf_rot = list(track.get('rotations', []))
                kf_scale = list(track.get('scales', []))

                # Combine with bind pose: final = bind * relative
                # rotation: bindRot * trackRot
                combined_rot = []
                for i in range(0, len(kf_rot), 4):
                    q = quat_mult(bind_rot, kf_rot[i:i+4])
                    combined_rot.extend(q)
                kf_rot = combined_rot

                # translation: bindPos + trackTrans
                combined_trans = []
                for i in range(0, len(kf_trans), 3):
                    combined_trans.extend([
                        bind_pos[0] + kf_trans[i],
                        bind_pos[1] + kf_trans[i+1],
                        bind_pos[2] + kf_trans[i+2],
                    ])
                kf_trans = combined_trans

                # scale: bindScale * trackScale
                combined_scale = []
                for i in range(0, len(kf_scale), 3):
                    combined_scale.extend([
                        bind_scale[0] * kf_scale[i],
                        bind_scale[1] * kf_scale[i+1],
                        bind_scale[2] * kf_scale[i+2],
                    ])
                kf_scale = combined_scale

                def lerp_to(times, values, val_size, targets):
                    if not values:
                        return []
                    result = []
                    n_kf = len(times)
                    for qt in targets:
                        if qt <= times[0]:
                            result.extend(values[0:val_size])
                        elif qt >= times[-1]:
                            result.extend(values[(n_kf-1)*val_size:(n_kf-1)*val_size+val_size])
                        else:
                            for i in range(n_kf - 1):
                                if times[i] <= qt <= times[i+1]:
                                    a = (qt - times[i]) / (times[i+1] - times[i]) if times[i+1] != times[i] else 0
                                    base_i = i * val_size
                                    base_j = (i+1) * val_size
                                    result.extend([
                                        values[base_i + j] + a * (values[base_j + j] - values[base_i + j])
                                        for j in range(val_size)
                                    ])
                                    break
                            else:
                                result.extend(values[(n_kf-1)*val_size:(n_kf-1)*val_size+val_size])
                    return result

                trans_interp = lerp_to(kf_times, kf_trans, 3, sorted_times) if kf_trans else []
                rot_interp = lerp_to(kf_times, kf_rot, 4, sorted_times) if kf_rot else []
                scale_interp = lerp_to(kf_times, kf_scale, 3, sorted_times) if kf_scale else []

                for vals, count, acc_type, path in [
                    (trans_interp, 3, 'VEC3', 'translation'),
                    (rot_interp, 4, 'VEC4', 'rotation'),
                    (scale_interp, 3, 'VEC3', 'scale'),
                ]:
                    if not vals:
                        continue
                    raw = pad4(struct.pack(f'<{len(vals)}f', *vals))
                    off2 = len(anim_bin)
                    anim_bin.extend(raw)
                    dv_idx = len(anim_buf_views)
                    anim_buf_views.append({'byteOffset': off2, 'byteLength': len(raw)})
                    acc_idx = len(anim_accessors)
                    anim_accessors.append({'type': acc_type, 'count': len(sorted_times)})
                    samp_idx = len(anim_samplers_this)
                    anim_samplers_this.append({'input': t_acc, 'output': acc_idx})
                    anim_channels_this.append({'sampler': samp_idx,
                                               'target': {'node': node_idx, 'path': path}})

            if anim_samplers_this:
                animations_out.append({
                    'name': anim_name,
                    'samplers': anim_samplers_this,
                    'channels': anim_channels_this,
                })
    else:
        # No skeleton - simple mesh node
        mesh_node_idx = len(nodes)
        nodes.append({'name': 'mesh', 'mesh': 0})
        scene_node_idx = len(nodes)
        nodes.append({'name': 'scene_root', 'children': [mesh_node_idx]})
        anim_bin = bytearray()
        anim_buf_views = []
        anim_accessors = []
        anim_samplers = []
        anim_channels = []

    # ── Build binary buffer ─────────────────────────────────────────────────
    pos_bytes = pad4(struct.pack(f'<{len(positions)}f', *positions))
    norm_bytes = pad4(struct.pack(f'<{len(normals)}f', *normals))
    idx_bytes = pad4(struct.pack(f'<{len(indices)}H', *indices))
    uv_bytes = pad4(struct.pack(f'<{len(texcoords)}f', *texcoords)) if texcoords else b''

    joint_bytes = b''
    weight_bytes = b''
    if skel and bone_weights:
        # Joints: use existing bone_indices if available
        joint_vals = [i & 0xFFFF for i in bone_indices] if bone_indices else [0] * (vertex_count * 4)
        joint_bytes = pad4(struct.pack(f'<{len(joint_vals)}H', *joint_vals))
        weight_bytes = pad4(struct.pack(f'<{len(bone_weights)}f', *bone_weights))

    ibm_bytes = pad4(struct.pack(f'<{len(ibms)}f', *ibms)) if skel else b''

    all_bin = pos_bytes + norm_bytes + idx_bytes + uv_bytes + joint_bytes + weight_bytes + ibm_bytes + bytes(anim_bin)

    # ── Buffer views and accessors ──────────────────────────────────────────
    off = 0
    buffer_views = []
    bv_pos = {'buffer': 0, 'byteOffset': off, 'byteLength': len(pos_bytes), 'target': 34962}; off += len(pos_bytes)
    bv_nrm = {'buffer': 0, 'byteOffset': off, 'byteLength': len(norm_bytes), 'target': 34962}; off += len(norm_bytes)
    bv_idx = {'buffer': 0, 'byteOffset': off, 'byteLength': len(idx_bytes), 'target': 34963}; off += len(idx_bytes)
    buffer_views = [bv_pos, bv_nrm, bv_idx]

    accessors = [
        {'bufferView': 0, 'componentType': 5126, 'count': vertex_count, 'type': 'VEC3'},
        {'bufferView': 1, 'componentType': 5126, 'count': vertex_count, 'type': 'VEC3'},
        {'bufferView': 2, 'componentType': 5123, 'count': len(indices), 'type': 'SCALAR'},
    ]

    attributes = {'POSITION': 0, 'NORMAL': 1}

    if texcoords:
        bv_uv = {'buffer': 0, 'byteOffset': off, 'byteLength': len(uv_bytes), 'target': 34962}; off += len(uv_bytes)
        buffer_views.append(bv_uv)
        accessors.append({'bufferView': 3, 'componentType': 5126, 'count': vertex_count, 'type': 'VEC2'})
        attributes['TEXCOORD_0'] = 3

    # Joint/weight accessor indices shift if UV present
    jnt_acc = 4 if texcoords else 3
    wgt_acc = 5 if texcoords else 4

    if skel and bone_weights:
        bv_jnt = {'buffer': 0, 'byteOffset': off, 'byteLength': len(joint_bytes), 'target': 34962}; off += len(joint_bytes)
        bv_wgt = {'buffer': 0, 'byteOffset': off, 'byteLength': len(weight_bytes), 'target': 34962}; off += len(weight_bytes)
        buffer_views.extend([bv_jnt, bv_wgt])
        accessors.append({'bufferView': 3 if not texcoords else 4, 'componentType': 5123, 'count': vertex_count, 'type': 'VEC4'})
        accessors.append({'bufferView': 4 if not texcoords else 5, 'componentType': 5126, 'count': vertex_count, 'type': 'VEC4'})
        attributes['JOINTS_0'] = jnt_acc
        attributes['WEIGHTS_0'] = wgt_acc

    ibm_accessor_idx = None
    if skel:
        bv_ibm = {'buffer': 0, 'byteOffset': off, 'byteLength': len(ibm_bytes)}; off += len(ibm_bytes)
        buffer_views.append(bv_ibm)
        ibm_accessor_idx = len(accessors)
        accessors.append({'bufferView': len(buffer_views)-1, 'componentType': 5126,
                          'count': len(skin['joints']), 'type': 'MAT4'})
        skin['inverseBindMatrices'] = ibm_accessor_idx

    anim_base = off
    for bv in anim_buf_views:
        buffer_views.append({'buffer': 0, 'byteOffset': anim_base + bv['byteOffset'], 'byteLength': bv['byteLength']})

    acc_offset = len(accessors)
    for acc in anim_accessors:
        a = {'componentType': 5126, 'type': acc['type'], 'count': acc['count'],
             'bufferView': len(buffer_views) - len(anim_accessors) + 0}
        accessors.append(a)

    # Fix animation accessor bufferView indices
    for i, acc in enumerate(anim_accessors):
        accessors[acc_offset + i]['bufferView'] = len(buffer_views) - len(anim_accessors) + i

    # Fix sampler input/output indices (they reference anim_accessors relative)
    for anim in animations_out:
        for s in anim['samplers']:
            s['input'] += acc_offset
            s['output'] += acc_offset

    # ── Assemble glTF JSON ──────────────────────────────────────────────────
    meshes = [{'name': geo.get('name', 'mesh'), 'primitives': [{
        'attributes': attributes,
        'indices': 2,
        'material': 0,
    }]}]

    gltf = {
        'asset': {'version': '2.0', 'generator': 'jme3-json2gltf.py'},
        'scene': 0,
        'scenes': [{'nodes': [scene_node_idx]}],
        'nodes': nodes,
        'meshes': meshes,
        'materials': [{'name': 'default', 'pbrMetallicRoughness': {
            'baseColorFactor': [1, 1, 1, 1], 'metallicFactor': 0, 'roughnessFactor': 0.5}}],
        'animations': animations_out,
        'accessors': accessors,
        'bufferViews': buffer_views,
        'buffers': [{'byteLength': len(all_bin)}],
    }

    if skin:
        gltf['skins'] = [skin]

    gltf_json = json.dumps(gltf, separators=(',', ':')).encode('utf-8')
    while len(gltf_json) % 4:
        gltf_json += b' '
    while len(all_bin) % 4:
        all_bin += b'\x00'

    header = struct.pack('<III', 0x46546C67, 2, 12 + 8 + len(gltf_json) + 8 + len(all_bin))
    json_chunk = struct.pack('<II', len(gltf_json), 0x4E4F534A)
    bin_chunk = struct.pack('<II', len(all_bin), 0x004E4942)

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

    print(f'  -> {output_path} (verts={vertex_count}, bones={len(skin["joints"]) if skin else 0}, '
          f'anims={len(animations_out)})')

def main():
    in_dir = '/var/home/nico/Gamedev/Wackelpeter/web/public/models'
    models = ['spider', 'staff', 'cross', 'helmet', 'wall1', 'tile1', 'tile2', 'blobknochen']
    for m in models:
        jp = os.path.join(in_dir, f'{m}.json')
        gp = os.path.join(in_dir, f'{m}.glb')
        if os.path.exists(jp):
            print(f'Converting {m}...')
            try:
                convert(jp, gp)
            except Exception as e:
                print(f'  ERROR: {e}')
                import traceback
                traceback.print_exc()

if __name__ == '__main__':
    main()