#!/usr/bin/env python3
"""Convert Ogre3D .mesh.xml + .skeleton.xml to glTF 2.0 (.glb) with all animations."""

import xml.etree.ElementTree as ET
import json
import struct
import math
import os

# ── Math ────────────────────────────────────────────────────────────────────

def axis_angle_to_quat(ax, ay, az, angle):
    s = math.sin(angle / 2)
    return (ax * s, ay * s, az * s, math.cos(angle / 2))

def mat4_from_trs(tx, ty, tz, qx, qy, qz, qw, sx, sy, sz):
    x, y, z, w = qx, qy, qz, qw
    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))*sx, (2*(xy + wz))*sx,     (2*(xz - wy))*sx,     0,
        (2*(xy - wz))*sy,     (1 - 2*(xx + zz))*sy, (2*(yz + wx))*sy,     0,
        (2*(xz + wy))*sz,     (2*(yz - wx))*sz,     (1 - 2*(xx + yy))*sz, 0,
        tx,                   ty,                   tz,                    1,
    ]

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 lerp_keyframes(kf_times, kf_vals, query_times):
    """Linearly interpolate keyframe values to query times."""
    if len(kf_times) == 1:
        v = kf_vals[0]
        result = []
        for _ in query_times:
            result.extend(v)
        return result
    result = []
    for qt in query_times:
        if qt <= kf_times[0]:
            result.extend(kf_vals[0])
        elif qt >= kf_times[-1]:
            result.extend(kf_vals[-1])
        else:
            for i in range(len(kf_times) - 1):
                if kf_times[i] <= qt <= kf_times[i+1]:
                    a = (qt - kf_times[i]) / (kf_times[i+1] - kf_times[i])
                    v = [kf_vals[i][j] + a*(kf_vals[i+1][j] - kf_vals[i][j])
                         for j in range(len(kf_vals[0]))]
                    result.extend(v)
                    break
            else:
                result.extend(kf_vals[-1])
    return result

def pad(data, alignment=4):
    """Pad binary data with null bytes to alignment (for BIN chunk)."""
    while len(data) % alignment:
        data += b'\x00'
    return data

# ── Parsing ─────────────────────────────────────────────────────────────────

def parse_skeleton(filepath):
    tree = ET.parse(filepath)
    root = tree.getroot()
    bones = []
    for b in root.find('bones').findall('bone'):
        bid = int(b.get('id', '0'))
        name = b.get('name', '')
        pos = b.find('position')
        px, py, pz = float(pos.get('x','0')), float(pos.get('y','0')), float(pos.get('z','0'))
        rot = b.find('rotation')
        angle = float(rot.get('angle','0'))
        axis = rot.find('axis')
        ax, ay, az = float(axis.get('x','1')), float(axis.get('y','0')), float(axis.get('z','0'))
        bones.append({'id': bid, 'name': name, 'pos': (px, py, pz),
                      'rot_axis': (ax, ay, az), 'rot_angle': angle, 'scale': (1,1,1)})
    hierarchy = {}
    hier = root.find('bonehierarchy')
    if hier is not None:
        for bp in hier.findall('boneparent'):
            hierarchy[bp.get('bone')] = bp.get('parent')
    animations = {}
    anims = root.find('animations')
    if anims is not None:
        for anim in anims.findall('animation'):
            name = anim.get('name')
            tracks = {}
            for track in anim.find('tracks').findall('track'):
                bn = track.get('bone')
                kfs = []
                for kf in track.find('keyframes').findall('keyframe'):
                    t = float(kf.get('time'))
                    tr = kf.find('translate')
                    tx, ty, tz = float(tr.get('x','0')), float(tr.get('y','0')), float(tr.get('z','0'))
                    ro = kf.find('rotate')
                    ang = float(ro.get('angle','0'))
                    ax_e = ro.find('axis')
                    ax, ay, az = float(ax_e.get('x','1')), float(ax_e.get('y','0')), float(ax_e.get('z','0'))
                    q = axis_angle_to_quat(ax, ay, az, ang)
                    sc = kf.find('scale')
                    sx = float(sc.get('x','1')) if sc is not None else 1.0
                    sy = float(sc.get('y','1')) if sc is not None else 1.0
                    sz = float(sc.get('z','1')) if sc is not None else 1.0
                    kfs.append((t, tx, ty, tz, q[0], q[1], q[2], q[3], sx, sy, sz))
                tracks[bn] = kfs
            animations[name] = tracks
    return bones, hierarchy, animations

def parse_mesh(filepath):
    tree = ET.parse(filepath)
    root = tree.getroot()
    verts = []
    sg = root.find('sharedgeometry')
    if sg is not None:
        for vb in sg.findall('vertexbuffer'):
            for v in vb.findall('vertex'):
                pos = v.find('position')
                norm = v.find('normal')
                vt = {'x': float(pos.get('x','0')), 'y': float(pos.get('y','0')), 'z': float(pos.get('z','0')),
                      'nx': 0, 'ny': 0, 'nz': 1}
                if norm is not None:
                    vt['nx'] = float(norm.get('x','0')); vt['ny'] = float(norm.get('y','0')); vt['nz'] = float(norm.get('z','0'))
                verts.append(vt)
    faces = []
    for sm in root.findall('.//submesh'):
        for fe in sm.findall('faces'):
            for f in fe.findall('face'):
                faces.append((int(f.get('v1','0')), int(f.get('v2','0')), int(f.get('v3','0'))))
    bw = [{} for _ in range(len(verts))]
    ba = root.find('boneassignments')
    if ba is not None:
        for a in ba.findall('vertexboneassignment'):
            vi = int(a.get('vertexindex','0'))
            bi = int(a.get('boneindex','0'))
            w = float(a.get('weight','1.0'))
            if vi < len(bw):
                bw[vi][bi] = w
    return verts, faces, bw

# ── glTF builder ────────────────────────────────────────────────────────────

def build_gltf(mesh_path, skeleton_path, output_path):
    verts, faces, bw = parse_mesh(mesh_path)
    bones, hierarchy, animations = parse_skeleton(skeleton_path)

    # Topological sort of bones
    bone_by_name = {b['name']: b for b in bones}
    bone_by_id = {b['id']: b for b in bones}

    roots = [b for b in bones if b['name'] not in hierarchy]
    order = []
    visited = set()
    def dfs(name):
        if name in visited: return
        visited.add(name)
        order.append(name)
        for child, parent in hierarchy.items():
            if parent == name:
                dfs(child)
    for root in roots:
        dfs(root['name'])

    name_to_idx = {name: i for i, name in enumerate(order)}

    # glTF nodes (bones only for now)
    nodes = []
    for name in order:
        b = bone_by_name[name]
        q = axis_angle_to_quat(*b['rot_axis'], b['rot_angle'])
        nodes.append({
            'name': name,
            'translation': list(b['pos']),
            'rotation': [q[0], q[1], q[2], q[3]],
            'scale': list(b['scale']),
        })
    for child, parent in hierarchy.items():
        if child in name_to_idx and parent in name_to_idx:
            ci, pi = name_to_idx[child], name_to_idx[parent]
            nodes[pi].setdefault('children', []).append(ci)

    # Add a skeleton root node (identity matrix, all bones as children)
    skeleton_root_idx = len(nodes)
    nodes.append({'name': '__skeleton_root__', 'children': []})
    # Only add bones that are NOT children of other bones in hierarchy
    for i, name in enumerate(order):
        is_child = any(i in nodes[idx].get('children', []) for idx in range(len(nodes)))
        if not is_child:
            nodes[skeleton_root_idx]['children'].append(i)

    # Inverse bind matrices
    def world_transform(name):
        b = bone_by_name[name]
        q = axis_angle_to_quat(*b['rot_axis'], b['rot_angle'])
        local = mat4_from_trs(b['pos'][0], b['pos'][1], b['pos'][2],
                              q[0], q[1], q[2], q[3], 1, 1, 1)
        parent = hierarchy.get(name)
        if parent and parent in bone_by_name:
            return mat4_mult(world_transform(parent), local)
        return local

    ibm_flat = []
    for name in order:
        w = world_transform(name)
        # Invert 4x4 matrix
        m = [w[i + j*4] for j in range(4) for i in range(4)]  # transpose to row-major for inversion
        inv = invert_4x4(m)
        ibm_flat.extend(inv)  # back in row-major? No - glTF uses column-major

    # Actually glTF stores matrices column-major in the linear buffer
    # The inverse bind matrix in column-major form: each 4 floats = 1 column
    ibm_col_major = []
    for name in order:
        w = world_transform(name)
        # w is column-major: [c0x, c0y, c0z, c0w, c1x, ..., c3w]
        # Invert it
        inv = invert_4x4_colmajor(w)
        ibm_col_major.extend(inv)

    # Joint indices for skin
    skin_joints = [name_to_idx[name] for name in order]

    # Vertex data
    positions = []
    nrm = []
    joints0 = []
    weights0 = []
    for i, v in enumerate(verts):
        positions.extend([v['x'], v['y'], v['z']])
        nrm.extend([v['nx'], v['ny'], v['nz']])
        wmap = bw[i] if i < len(bw) else {}
        sw = sorted(wmap.items(), key=lambda x: x[1], reverse=True)[:4]
        j = [0, 0, 0, 0]
        wgt = [0.0, 0.0, 0.0, 0.0]
        for k, (bid, weight) in enumerate(sw):
            if bid in bone_by_id:
                jname = bone_by_id[bid]['name']
                j[k] = name_to_idx.get(jname, 0)
            wgt[k] = weight
        total = sum(wgt)
        if total > 0:
            wgt = [x / total for x in wgt]
        joints0.extend(j)
        weights0.extend(wgt)

    indices = []
    for f in faces:
        indices.extend(list(f))

    # Bounding box
    all_x = [v['x'] for v in verts]
    all_y = [v['y'] for v in verts]
    all_z = [v['z'] for v in verts]
    bbox_min = [min(all_x), min(all_y), min(all_z)]
    bbox_max = [max(all_x), max(all_y), max(all_z)]

    # Model node (has mesh and skin)
    model_node_idx = len(nodes)
    nodes.append({'name': os.path.basename(mesh_path).replace('.mesh.xml', '') + '_mesh',
                  'mesh': 0, 'skin': 0})

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

    # ── Build buffer ────────────────────────────────────────────────────────
    pos_bytes = pad(struct.pack(f'<{len(positions)}f', *positions))
    nrm_bytes = pad(struct.pack(f'<{len(nrm)}f', *nrm))
    idx_bytes = pad(struct.pack(f'<{len(indices)}H', *indices))
    jnt_bytes = pad(struct.pack(f'<{len(joints0)}H', *joints0))
    wgt_bytes = pad(struct.pack(f'<{len(weights0)}f', *weights0))
    ibm_bytes = pad(struct.pack(f'<{len(ibm_col_major)}f', *ibm_col_major))

    # Base data end
    base_end = len(pos_bytes) + len(nrm_bytes) + len(idx_bytes) + len(jnt_bytes) + len(wgt_bytes) + len(ibm_bytes)

    # ── Animations ──────────────────────────────────────────────────────────
    anim_data = bytearray()
    gltf_animations = []
    buf_views = []
    accessors_list = []

    # Fill static buffer views and accessors first
    off = 0
    bv_pos = {'buffer': 0, 'byteOffset': off, 'byteLength': len(pos_bytes), 'target': 34962}; off += len(pos_bytes)
    bv_nrm = {'buffer': 0, 'byteOffset': off, 'byteLength': len(nrm_bytes), 'target': 34962}; off += len(nrm_bytes)
    bv_idx = {'buffer': 0, 'byteOffset': off, 'byteLength': len(idx_bytes), 'target': 34963}; off += len(idx_bytes)
    bv_jnt = {'buffer': 0, 'byteOffset': off, 'byteLength': len(jnt_bytes), 'target': 34962}; off += len(jnt_bytes)
    bv_wgt = {'buffer': 0, 'byteOffset': off, 'byteLength': len(wgt_bytes), 'target': 34962}; off += len(wgt_bytes)
    bv_ibm = {'buffer': 0, 'byteOffset': off, 'byteLength': len(ibm_bytes)}; off += len(ibm_bytes)

    buffer_views = [bv_pos, bv_nrm, bv_idx, bv_jnt, bv_wgt, bv_ibm]

    accessors = [
        {'bufferView': 0, 'componentType': 5126, 'count': len(verts), 'type': 'VEC3',
         'min': bbox_min, 'max': bbox_max},
        {'bufferView': 1, 'componentType': 5126, 'count': len(verts), 'type': 'VEC3'},
        {'bufferView': 2, 'componentType': 5123, 'count': len(indices), 'type': 'SCALAR'},
        {'bufferView': 3, 'componentType': 5123, 'count': len(verts), 'type': 'VEC4'},
        {'bufferView': 4, 'componentType': 5126, 'count': len(verts), 'type': 'VEC4'},
        {'bufferView': 5, 'componentType': 5126, 'count': len(order), 'type': 'MAT4'},
    ]

    anim_base = off  # animations start here

    for anim_name, tracks in animations.items():
        all_times = set()
        for bn, kfs in tracks.items():
            for t, *_ in kfs:
                all_times.add(t)
        sorted_times = sorted(all_times)
        if len(sorted_times) < 2:
            continue

        time_data = pad(struct.pack(f'<{len(sorted_times)}f', *sorted_times))
        time_off = len(anim_data)
        anim_data.extend(time_data)

        # Buffer view for time
        tv_idx = len(buffer_views)
        buffer_views.append({'buffer': 0, 'byteOffset': anim_base + time_off, 'byteLength': len(time_data)})

        # Accessor for time (shared per animation)
        t_acc_idx = len(accessors)
        accessors.append({'bufferView': tv_idx, 'componentType': 5126, 'count': len(sorted_times),
                          'type': 'SCALAR', 'min': [sorted_times[0]], 'max': [sorted_times[-1]]})

        samplers = []
        channels = []

        for bone_name, keyframes in tracks.items():
            if bone_name not in name_to_idx:
                continue
            node_idx = name_to_idx[bone_name]

            kf_times = [t for t, *_ in keyframes]
            kf_trans = [list(kf[1:4]) for kf in keyframes]
            kf_rot = [list(kf[4:8]) for kf in keyframes]
            kf_scale = [list(kf[8:11]) for kf in keyframes]

            for data, acc_type, path in [
                (lerp_keyframes(kf_times, kf_trans, sorted_times), 'VEC3', 'translation'),
                (lerp_keyframes(kf_times, kf_rot, sorted_times), 'VEC4', 'rotation'),
                (lerp_keyframes(kf_times, kf_scale, sorted_times), 'VEC3', 'scale'),
            ]:
                raw = pad(struct.pack(f'<{len(data)}f', *data))
                off2 = len(anim_data)
                anim_data.extend(raw)

                dv_idx = len(buffer_views)
                buffer_views.append({'buffer': 0, 'byteOffset': anim_base + off2, 'byteLength': len(raw)})

                acc_idx = len(accessors)
                accessors.append({'bufferView': dv_idx, 'componentType': 5126,
                                  'count': len(sorted_times), 'type': acc_type})

                samp_idx = len(samplers)
                samplers.append({'input': t_acc_idx, 'interpolation': 'LINEAR', 'output': acc_idx})

                channels.append({'sampler': samp_idx, 'target': {'node': node_idx, 'path': path}})

        if samplers and channels:
            gltf_animations.append({'name': anim_name, 'samplers': samplers, 'channels': channels})

    anim_bin = bytes(anim_data)

    # Combine
    all_bin = pos_bytes + nrm_bytes + idx_bytes + jnt_bytes + wgt_bytes + ibm_bytes + anim_bin

    # ── Build JSON ──────────────────────────────────────────────────────────
    meshes = [{'name': 'mesh', 'primitives': [{
        'attributes': {'POSITION': 0, 'NORMAL': 1, 'JOINTS_0': 3, 'WEIGHTS_0': 4},
        'indices': 2,
        'material': 0,
    }]}]

    materials = [{
        'name': 'default',
        'pbrMetallicRoughness': {
            'baseColorFactor': [1.0, 1.0, 1.0, 1.0],
            'metallicFactor': 0.0,
            'roughnessFactor': 0.5,
        },
    }]

    skin = [{'inverseBindMatrices': 5, 'joints': skin_joints, 'name': 'skin'}]

    gltf = {
        'asset': {'version': '2.0', 'generator': 'ogre2gltf.py'},
        'scene': 0,
        'scenes': [{'nodes': [scene_node_idx]}],
        'nodes': nodes,
        'meshes': meshes,
        'materials': materials,
        'skins': skin,
        'animations': gltf_animations,
        'accessors': accessors,
        'bufferViews': buffer_views,
        'buffers': [{'byteLength': len(all_bin)}],
    }

    # ── Write GLB ───────────────────────────────────────────────────────────
    gltf_json = json.dumps(gltf, separators=(',', ':'), allow_nan=False).encode('utf-8')
    # GLB spec: JSON chunk padded with spaces (0x20), BIN chunk padded with nulls (0x00)
    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)

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    with open(output_path, 'wb') as f:
        f.write(header + json_chunk + gltf_json + bin_chunk + all_bin)

    print(f'  -> {output_path} ({len(verts)}v {len(faces)}f {len(bones)}b {len(gltf_animations)}a)')

# ── Matrix inversion ───────────────────────────────────────────────────────

def invert_4x4(m):
    """Invert 4x4 row-major matrix, return row-major."""
    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]
    inv_det = 1.0 / det
    return [
        (f*k*p + g*l*n + h*j*o - h*k*n - f*l*o - g*j*p) * inv_det,
        (b*k*p + c*l*n + d*j*o - d*k*n - b*l*o - c*j*p) * inv_det,
        (b*g*p + c*h*n + d*f*o - d*g*n - b*h*o - c*f*p) * inv_det,
        (b*g*l + c*h*j + d*f*k - d*g*j - b*h*k - c*f*l) * inv_det,
        (e*k*p + g*l*m_ + h*i*o - h*k*m_ - e*l*o - g*i*p) * inv_det,
        (a*k*p + c*l*m_ + d*i*o - d*k*m_ - a*l*o - c*i*p) * inv_det,
        (a*g*p + c*h*m_ + d*e*o - d*g*m_ - a*h*o - c*e*p) * inv_det,
        (a*g*l + c*h*i + d*e*k - d*g*i - a*h*k - c*e*l) * inv_det,
        (e*j*p + f*l*m_ + h*i*n - h*j*m_ - e*l*n - f*i*p) * inv_det,
        (a*j*p + b*l*m_ + d*i*n - d*j*m_ - a*l*n - b*i*p) * inv_det,
        (a*f*p + b*h*m_ + d*e*n - d*f*m_ - a*h*n - b*e*p) * inv_det,
        (a*f*l + b*h*i + d*e*j - d*f*i - a*h*j - b*e*l) * inv_det,
        (e*j*o + f*k*m_ + g*i*n - g*j*m_ - e*k*n - f*i*o) * inv_det,
        (a*j*o + b*k*m_ + c*i*n - c*j*m_ - a*k*n - b*i*o) * inv_det,
        (a*f*o + b*g*m_ + c*e*n - c*f*m_ - a*g*n - b*e*o) * inv_det,
        (a*f*k + b*g*i + c*e*j - c*f*i - a*g*j - b*e*k) * inv_det,
    ]

def invert_4x4_colmajor(m):
    """Invert 4x4 column-major matrix, return column-major."""
    # Convert col-major to row-major, invert, convert back
    rowm = [m[i + j*4] for j in range(4) for i in range(4)]
    inv_rowm = invert_4x4(rowm)
    return [inv_rowm[j + i*4] for i in range(4) for j in range(4)]

# ── Main ────────────────────────────────────────────────────────────────────

def main():
    base = '/var/home/nico/Gamedev/Wackelpeter/extracted/Models'
    out = '/var/home/nico/Gamedev/Wackelpeter/web/public/models'

    for name, mesh, skel in [
        ('blob', 'blob/Blob', 'blob/Blob'),
        ('sword', 'sword/sword', 'sword/sword'),
        ('shield', 'shield/shield', 'shield/shield'),
    ]:
        mf = os.path.join(base, f'{mesh}.mesh.xml')
        sf = os.path.join(base, f'{skel}.skeleton.xml')
        if os.path.exists(mf) and os.path.exists(sf):
            print(f'Converting {name}...')
            build_gltf(mf, sf, os.path.join(out, f'{name}.glb'))
        else:
            print(f'Skipping {name} (missing XML)')

if __name__ == '__main__':
    main()