#!/usr/bin/env python3
"""Parse JMonkeyEngine .j3o binary files and extract mesh/skeleton/animation data.

Outputs glTF 2.0 .glb for meshes with skeletons, or OBJ-like JSON for static meshes.
"""

import struct
import math
import os
import json
from pathlib import Path

# ── Constants ───────────────────────────────────────────────────────────────

NULL_OBJECT = -1
DEFAULT_OBJECT = -2
SIGNATURE = 0x4A4D4533

FIELD_TYPE = {
    0:  'BYTE', 1: 'BYTE_1D', 2: 'BYTE_2D',
    10: 'INT', 11: 'INT_1D', 12: 'INT_2D',
    20: 'FLOAT', 21: 'FLOAT_1D', 22: 'FLOAT_2D',
    30: 'DOUBLE', 31: 'DOUBLE_1D', 32: 'DOUBLE_2D',
    40: 'LONG', 41: 'LONG_1D', 42: 'LONG_2D',
    50: 'SHORT', 51: 'SHORT_1D', 52: 'SHORT_2D',
    60: 'BOOLEAN', 61: 'BOOLEAN_1D', 62: 'BOOLEAN_2D',
    70: 'STRING', 71: 'STRING_1D', 72: 'STRING_2D',
    80: 'BITSET',
    90: 'SAVABLE', 91: 'SAVABLE_1D', 92: 'SAVABLE_2D',
    100: 'SAVABLE_ARRAYLIST', 101: 'SAVABLE_ARRAYLIST_1D', 102: 'SAVABLE_ARRAYLIST_2D',
    105: 'SAVABLE_MAP', 106: 'STRING_SAVABLE_MAP', 107: 'INT_SAVABLE_MAP',
    110: 'FLOATBUFFER_ARRAYLIST', 111: 'BYTEBUFFER_ARRAYLIST',
    120: 'FLOATBUFFER', 121: 'INTBUFFER', 122: 'BYTEBUFFER', 123: 'SHORTBUFFER',
}

# VertexBuffer usage types (key for IntMap in Mesh.buffers)
VB_USAGE = {
    'Position': 0, 'Normal': 1, 'TexCoord': 2, 'Color': 3,
    'Tangent': 4, 'Binormal': 5,
    'Size': 6, 'InterleavedData': 7, 'Misc': 8,
    'BoneIndex': 9, 'BoneWeight': 10, 'HWBoneIndex': 11, 'HWBoneWeight': 12,
    'BindPosePosition': 13, 'BindPoseNormal': 14, 'BindPoseTangent': 15,
}

VB_FORMAT = {'Float': 0, 'Short': 2, 'UnsignedShort': 4, 'UnsignedByte': 5, 'Byte': 6, 'Half': 7, 'Int': 8, 'UnsignedInt': 9}


# ── Binary reader ───────────────────────────────────────────────────────────

class J3OReader:
    def __init__(self, data):
        self.data = data
        self.offset = 0

    def tell(self):
        return self.offset

    def read_bytes(self, n):
        r = self.data[self.offset:self.offset + n]
        self.offset += n
        return r

    def read_be_int32(self):
        return struct.unpack('>i', self.read_bytes(4))[0]

    def read_be_uint32(self):
        return struct.unpack('>I', self.read_bytes(4))[0]

    def read_be_int16(self):
        return struct.unpack('>h', self.read_bytes(2))[0]

    def read_byte(self):
        b = self.data[self.offset]
        self.offset += 1
        return b

    def read_float(self):
        i = struct.unpack('>i', self.read_bytes(4))[0]
        return struct.unpack('>f', struct.pack('>i', i))[0]

    def read_double(self):
        i = struct.unpack('>q', self.read_bytes(8))[0]
        return struct.unpack('>d', struct.pack('>q', i))[0]

    def read_short(self):
        return struct.unpack('>h', self.read_bytes(2))[0]

    def read_bool(self):
        return self.read_byte() != 0

    def read_compressed_int(self):
        first = self.read_byte()
        if first == 0xFF:
            return NULL_OBJECT
        if first == 0xFE:
            return DEFAULT_OBJECT
        if first == 0x00:
            return 0
        length = first & 0xFF
        raw = self.read_bytes(length)
        # Right-align to 4 bytes
        if length <= 4:
            pad_len = 4 - length
            value_bytes = b'\x00' * pad_len + raw
        else:
            # If longer than 4 bytes, take only the last 4
            value_bytes = raw[length-4:]
        try:
            val = struct.unpack('>i', value_bytes)[0]
        except struct.error:
            return NULL_OBJECT
        if val in (NULL_OBJECT, DEFAULT_OBJECT):
            if length == 4:
                self.offset -= 4
        return val

    def read_compressed_long(self):
        first = self.read_byte()
        if first == 0xFF:
            return NULL_OBJECT
        if first == 0xFE:
            return DEFAULT_OBJECT
        if first == 0x00:
            return 0
        length = first & 0xFF
        raw = self.read_bytes(length)
        pad = b'\x00' * (8 - length)
        return struct.unpack('>q', pad + raw)[0]

    def read_string(self):
        length = self.read_compressed_int()
        if length == NULL_OBJECT:
            return None
        return self.read_bytes(length).decode('utf-8', errors='replace')


# ── Field reading ───────────────────────────────────────────────────────────

class CapsuleReader:
    """Reads field data for a single capsule."""
    def __init__(self, rdr: J3OReader, length):
        self.rdr = rdr
        self.end_offset = rdr.tell() + length
        self.class_fields = None  # set by BinaryClassObject

    def read_field_value(self, field_type, expected_alias=None):
        """Read a value of the given type. Returns (value, field_alias)."""
        if expected_alias is not None:
            alias = self.rdr.read_byte()
        else:
            alias = 0

        val = None

        if field_type in (0,):  # BYTE
            val = self.rdr.read_byte()

        elif field_type in (1,):  # BYTE_1D
            ln = self.rdr.read_compressed_int()
            val = list(self.rdr.read_bytes(ln)) if ln > 0 else []

        elif field_type in (2,):  # BYTE_2D
            outer = self.rdr.read_compressed_int()
            val = []
            for _ in range(outer):
                ln = self.rdr.read_compressed_int()
                val.append(list(self.rdr.read_bytes(ln)) if ln > 0 else [])

        elif field_type in (10,):  # INT
            val = self.rdr.read_compressed_int()

        elif field_type in (11,):  # INT_1D
            ln = self.rdr.read_compressed_int()
            val = [self.rdr.read_compressed_int() for _ in range(ln)]

        elif field_type in (12,):  # INT_2D
            outer = self.rdr.read_compressed_int()
            val = [[self.rdr.read_compressed_int() for _ in range(self.rdr.read_compressed_int())] for _ in range(outer)]

        elif field_type in (20,):  # FLOAT
            val = self.rdr.read_float()

        elif field_type in (21,):  # FLOAT_1D
            ln = self.rdr.read_compressed_int()
            val = [self.rdr.read_float() for _ in range(ln)]

        elif field_type in (22,):  # FLOAT_2D
            outer = self.rdr.read_compressed_int()
            val = [[self.rdr.read_float() for _ in range(self.rdr.read_compressed_int())] for _ in range(outer)]

        elif field_type in (30,):  # DOUBLE
            val = self.rdr.read_double()

        elif field_type in (40,):  # LONG
            val = self.rdr.read_compressed_long()

        elif field_type in (50,):  # SHORT
            val = self.rdr.read_short()

        elif field_type in (51,):  # SHORT_1D
            ln = self.rdr.read_compressed_int()
            val = [self.rdr.read_short() for _ in range(ln)]

        elif field_type in (60,):  # BOOLEAN
            val = self.rdr.read_bool()

        elif field_type in (61,):  # BOOLEAN_1D
            ln = self.rdr.read_compressed_int()
            val = [self.rdr.read_bool() for _ in range(ln)]

        elif field_type in (70,):  # STRING
            val = self.rdr.read_string()

        elif field_type in (71,):  # STRING_1D
            ln = self.rdr.read_compressed_int()
            val = [self.rdr.read_string() for _ in range(ln)]

        elif field_type in (80,):  # BITSET
            ln = self.rdr.read_compressed_int()
            val = [self.rdr.read_bool() for _ in range(ln)]

        elif field_type in (90,):  # SAVABLE
            val = self.rdr.read_compressed_int()

        elif field_type in (91, 100):  # SAVABLE_1D / SAVABLE_ARRAYLIST
            ln = self.rdr.read_compressed_int()
            val = [self.rdr.read_compressed_int() for _ in range(ln)]

        elif field_type in (92, 101):  # SAVABLE_2D
            outer = self.rdr.read_compressed_int()
            val = [[self.rdr.read_compressed_int() for _ in range(self.rdr.read_compressed_int())] for _ in range(outer)]

        elif field_type in (105,):  # SAVABLE_MAP
            ln = self.rdr.read_compressed_int()
            val = [[self.rdr.read_compressed_int(), self.rdr.read_compressed_int()] for _ in range(ln)]

        elif field_type in (106,):  # STRING_SAVABLE_MAP
            ln = self.rdr.read_compressed_int()
            keys = [self.rdr.read_string() for _ in range(ln)]
            _ = self.rdr.read_compressed_int()  # skip length
            vals = [self.rdr.read_compressed_int() for _ in range(ln)]
            val = dict(zip(keys, vals))

        elif field_type in (107,):  # INT_SAVABLE_MAP
            ln = self.rdr.read_compressed_int()
            keys = [self.rdr.read_compressed_int() for _ in range(ln)]
            ln2 = self.rdr.read_compressed_int()
            vals = [self.rdr.read_compressed_int() for _ in range(ln2)]
            val = dict(zip(keys, vals))

        elif field_type in (120, 121, 122, 123):  # Buffer types
            ln = self.rdr.read_compressed_int()
            if ln > 0:
                elem_size = {120: 4, 121: 4, 122: 1, 123: 2}[field_type]
                raw = self.rdr.read_bytes(ln * elem_size)
                val = raw
            else:
                val = b''

        elif field_type in (110,):  # FLOATBUFFER_ARRAYLIST
            ln = self.rdr.read_compressed_int()
            val = []
            for _ in range(ln):
                nb = self.rdr.read_compressed_int()
                if nb > 0:
                    val.append(self.rdr.read_bytes(nb * 4))
                else:
                    val.append(b'')

        else:
            # Unknown type - skip ahead
            val = None

        return val, alias

    def read_all_fields(self, bco):
        """Read all fields of a capsule using its class definition."""
        fields_read = {}
        while self.rdr.tell() < self.end_offset - 1:
            try:
                alias = self.rdr.read_byte()
            except (IndexError, struct.error):
                break
            field = bco.alias_fields.get(alias)
            if field is None:
                # Unknown field alias - skip the rest
                break
            try:
                val, _ = self.read_field_value(field['type'], expected_alias=None)
                fields_read[field['name']] = val
            except (IndexError, struct.error):
                break
        return fields_read


# ── Class table ─────────────────────────────────────────────────────────────

class BinaryClassObject:
    def __init__(self, alias, name, fields):
        self.alias = alias
        self.name = name
        self.fields = fields
        self.alias_fields = {f['alias']: f for f in fields}


def parse_j3o_header(filepath):
    """Parse J3O header, class table, and location table. Return loader function."""
    with open(filepath, 'rb') as f:
        data = f.read()

    rdr = J3OReader(data)

    # Header
    sig = rdr.read_be_uint32()
    if sig != SIGNATURE:
        # Old format (version 0): no signature, first int is num_classes
        rdr.offset = 0
        version = 0
        num_classes = rdr.read_be_int32()
    else:
        version = rdr.read_be_int32()
        num_classes = rdr.read_be_int32()

    alias_width = int(math.log(max(num_classes, 2), 256)) + 1

    # Class table
    classes = {}
    for i in range(num_classes):
        alias = rdr.read_bytes(alias_width)

        if version >= 1:
            hier_size = rdr.read_byte()
            _ = [rdr.read_be_int32() for _ in range(hier_size)]
        else:
            _ = [0]

        name_len = rdr.read_be_int32()
        name = rdr.read_bytes(name_len).decode('ascii', errors='replace')

        num_fields = rdr.read_be_int32()
        fields = []
        for _ in range(num_fields):
            fa = rdr.read_byte()
            ft = rdr.read_byte()
            fn_len = rdr.read_be_int32()
            fn = rdr.read_bytes(fn_len).decode('ascii', errors='replace')
            fields.append({'alias': fa, 'type': ft, 'name': fn})

        classes[alias] = BinaryClassObject(alias, name, fields)

    # Location table
    num_locs = rdr.read_be_int32()
    location_table = {}
    for _ in range(num_locs):
        obj_id = rdr.read_be_int32()
        loc = rdr.read_be_int32()
        location_table[obj_id] = loc

    # Root
    _ = rdr.read_be_int32()
    root_id = rdr.read_be_int32()

    payload_start = rdr.tell()
    payload = data[payload_start:]

    return alias_width, classes, location_table, root_id, payload


def read_object(obj_id, alias_width, classes, location_table, payload, cache=None):
    """Read a single capsule and return its class name and field data."""
    if cache is None:
        cache = {}

    if obj_id in cache:
        return cache[obj_id]

    if obj_id not in location_table:
        cache[obj_id] = ('unknown', {})
        return cache[obj_id]

    loc = location_table[obj_id]
    rdr = J3OReader(payload)
    rdr.offset = loc

    alias = rdr.read_bytes(alias_width)
    bco = classes.get(alias)

    data_len = rdr.read_be_int32()

    if bco is None:
        cache[obj_id] = ('unknown', {})
        return cache[obj_id]

    cap_rdr = CapsuleReader(rdr, data_len)
    fields = cap_rdr.read_all_fields(bco)

    result = (bco.name, fields)
    cache[obj_id] = result
    return result


def resolve_savable(obj_id, alias_width, classes, location_table, payload, cache):
    """Resolve a Savable reference, returning (class_name, fields)."""
    if obj_id <= 0 or obj_id == NULL_OBJECT or obj_id == DEFAULT_OBJECT:
        return None
    return read_object(obj_id, alias_width, classes, location_table, payload, cache)


# ── Vertex buffer decoding ──────────────────────────────────────────────────

def decode_floatbuffer(raw, count):
    """Decode LE float buffer. raw is bytes."""
    if len(raw) != count * 4:
        return [0.0] * count
    return list(struct.unpack(f'<{count}f', raw))

def decode_shortbuffer(raw, count):
    return list(struct.unpack(f'<{count}h', raw))

def decode_bytebuffer(raw, count):
    return list(raw[:count])

def decode_intbuffer(raw, count):
    return list(struct.unpack(f'<{count}i', raw))


# ── Mesh extraction ─────────────────────────────────────────────────────────

def extract_mesh_geometry(obj_id, alias_width, classes, location_table, payload, cache):
    obj = resolve_savable(obj_id, alias_width, classes, location_table, payload, cache)
    if obj is None:
        return None

    cls_name, fields = obj
    if 'Mesh' not in cls_name:
        return None

    vert_count = fields.get('vertCount', 0)
    if vert_count == 0:
        return None

    buffers_map = fields.get('buffers', {})
    element_count = fields.get('elementCount', 0)

    positions = []
    normals = []
    texcoords = []
    indices = []
    bone_weights_raw = b''
    bone_indices_raw = b''
    bone_idx_bytes_per_elem = 2
    weight_comps = 4
    idx_comps = 4

    for usage_code, vb_id in buffers_map.items():
        vb_obj = resolve_savable(vb_id, alias_width, classes, location_table, payload, cache)
        if vb_obj is None:
            continue

        _, vb_fields = vb_obj
        comp = vb_fields.get('components', 0)
        format_code = vb_fields.get('format', -1)

        raw_data = None
        is_float = False
        is_short = False
        is_byte = False
        for dk, fl, sh, bt in [
            ('dataFloat', True, False, False),
            ('dataUnsignedShort', False, True, False),
            ('dataUnsignedByte', False, False, True),
            ('dataShort', False, True, False),
            ('dataByte', False, False, True),
            ('dataInt', False, False, False),
            ('data', False, False, False),
        ]:
            if dk in vb_fields:
                raw_data = vb_fields[dk]
                is_float = fl
                is_short = sh
                is_byte = bt
                break

        if raw_data is None:
            continue

        if usage_code == VB_USAGE['Position']:
            elem_count = len(raw_data) // 4
            positions = decode_floatbuffer(raw_data, elem_count)
            actual_vert_count = elem_count // comp
        elif usage_code == VB_USAGE['Normal']:
            elem_count = len(raw_data) // 4
            normals = decode_floatbuffer(raw_data, elem_count)
        elif usage_code == VB_USAGE['TexCoord']:
            elem_count = len(raw_data) // 4
            texcoords = decode_floatbuffer(raw_data, elem_count)
        elif usage_code == VB_USAGE['BoneWeight']:  # 10
            bone_weights_raw = raw_data
            weight_comps = comp
        elif usage_code == VB_USAGE['BoneIndex']:  # 9
            bone_indices_raw = raw_data
            bone_idx_bytes_per_elem = 2 if is_short else 1
            idx_comps = comp
        elif usage_code == VB_USAGE['Misc']:  # 8
            if is_short:
                indices = list(struct.unpack(f'<{len(raw_data)//2}h', raw_data))
            elif len(raw_data) >= 2:
                indices = list(struct.unpack(f'<{len(raw_data)//2}H', raw_data))
            elif raw_data:
                indices = list(raw_data)

    if not positions:
        return None

    # Decode bone data now that we know actual_vert_count
    actual_vert_count = len(positions) // 3 if positions else vert_count

    bone_indices = []
    if bone_indices_raw:
        if bone_idx_bytes_per_elem == 2:
            count = len(bone_indices_raw) // 2
            total = min(count, actual_vert_count * idx_comps)
            bone_indices = list(struct.unpack(f'<{total}H', bone_indices_raw[:total*2]))
        else:
            total = min(len(bone_indices_raw), actual_vert_count * idx_comps)
            bone_indices = list(bone_indices_raw[:total])

    bone_weights = []
    if bone_weights_raw:
        total = min(len(bone_weights_raw) // 4, actual_vert_count * weight_comps)
        bone_weights = decode_floatbuffer(bone_weights_raw, total)

    if not indices and element_count > 0:
        indices = list(range(element_count))

    return {
        'vertexCount': actual_vert_count,
        'positions': positions,
        'normals': normals,
        'texcoords': texcoords,
        'indices': indices,
        'boneWeights': bone_weights,
        'boneIndices': bone_indices,
    }


# ── Skeleton extraction ─────────────────────────────────────────────────────

def extract_skeleton(skeleton_id, alias_width, classes, location_table, payload, cache):
    """Extract skeleton data."""
    obj = resolve_savable(skeleton_id, alias_width, classes, location_table, payload, cache)
    if obj is None:
        return None

    cls_name, fields = obj
    if 'Skeleton' not in cls_name:
        return None

    root_bones_ids = fields.get('rootBones', [])
    bone_list_ids = fields.get('boneList', [])

    bones = []
    bid_to_idx = {}

    for i, bid in enumerate(bone_list_ids):
        bone_obj = resolve_savable(bid, alias_width, classes, location_table, payload, cache)
        if bone_obj is None:
            continue
        _, bf = bone_obj
        bones.append({
            'id': i,
            'name': bf.get('name', f'bone_{i}'),
            'position': bf.get('bindPos', [0,0,0]) if 'bindPos' in bf else [0,0,0],
            'rotation': bf.get('bindRot', [0,0,0,1]) if 'bindRot' in bf else [0,0,0,1],
            'scale': bf.get('bindScale', [1,1,1]) if 'bindScale' in bf else [1,1,1],
        })
        bid_to_idx[bid] = i

    # Hierarchy
    hierarchy = {}
    for bid in bone_list_ids:
        bone_obj = resolve_savable(bid, alias_width, classes, location_table, payload, cache)
        if bone_obj is None:
            continue
        _, bf = bone_obj
        children_ids = bf.get('children', [])
        parent_idx = bid_to_idx.get(bid)
        for cid in children_ids:
            child_idx = bid_to_idx.get(cid)
            if child_idx is not None and parent_idx is not None:
                hierarchy[bones[child_idx]['name']] = bones[parent_idx]['name']

    roots = [bid_to_idx[rid] for rid in root_bones_ids if rid in bid_to_idx]

    return {
        'bones': bones,
        'hierarchy': hierarchy,
        'roots': roots,
    }


# ── Animation extraction ────────────────────────────────────────────────────

def extract_animations(anim_map, alias_width, classes, location_table, payload, cache):
    """Extract animation data from StringSavableMap of anim_name -> Animation."""
    if anim_map is None:
        return {}

    animations = {}
    for anim_name, anim_id in anim_map.items():
        anim_obj = resolve_savable(anim_id, alias_width, classes, location_table, payload, cache)
        if anim_obj is None:
            continue
        _, af = anim_obj
        tracks_ids = af.get('tracks', [])
        anim_len = af.get('length', 1.0)
        anim_name_str = af.get('name', anim_name) or anim_name

        tracks = []
        for tid in tracks_ids:
            track_obj = resolve_savable(tid, alias_width, classes, location_table, payload, cache)
            if track_obj is None:
                continue
            ct, tf = track_obj
            bone_idx = tf.get('boneIndex', 0)

            # times: float array
            times = tf.get('times', [])

            # translations: savable CompactVector3Array
            trans_id = tf.get('translations', None)
            translations = []
            if trans_id is not None:
                trans_obj = resolve_savable(trans_id, alias_width, classes, location_table, payload, cache)
                if trans_obj:
                    _, trf = trans_obj
                    trans_data = trf.get('array', None)
                    if trans_data:
                        translations = list(struct.unpack(f'<{len(trans_data)//4}f', trans_data)) if trans_data else []

            # rotations: savable CompactQuaternionArray
            rot_id = tf.get('rotations', None)
            rotations = []
            if rot_id is not None:
                rot_obj = resolve_savable(rot_id, alias_width, classes, location_table, payload, cache)
                if rot_obj:
                    _, rf = rot_obj
                    rot_data = rf.get('array', None)
                    if rot_data:
                        rotations = list(struct.unpack(f'<{len(rot_data)//4}f', rot_data)) if rot_data else []

            # scales: savable CompactVector3Array
            scale_id = tf.get('scales', None)
            scales = []
            if scale_id is not None:
                scale_obj = resolve_savable(scale_id, alias_width, classes, location_table, payload, cache)
                if scale_obj:
                    _, sf = scale_obj
                    scale_data = sf.get('array', None)
                    if scale_data:
                        scales = list(struct.unpack(f'<{len(scale_data)//4}f', scale_data)) if scale_data else []

            tracks.append({
                'boneIndex': bone_idx,
                'times': times,
                'translations': translations,
                'rotations': rotations,
                'scales': scales,
            })

        animations[anim_name_str] = {
            'length': anim_len,
            'tracks': tracks,
        }

    return animations


# ── Main extraction ─────────────────────────────────────────────────────────

def extract_j3o(filepath):
    """Extract mesh, skeleton, and animation data from a .j3o file."""
    alias_width, classes, location_table, root_id, payload = parse_j3o_header(filepath)
    cache = {}
    result = {
        'file': filepath,
        'geometries': [],
        'skeletons': [],
        'animations': {},
    }

    # Walk object graph looking for mesh/skeleton/animation data
    def walk_obj(obj_id, depth=0):
        if depth > 50 or obj_id <= 0:
            return
        try:
            obj = resolve_savable(obj_id, alias_width, classes, location_table, payload, cache)
        except Exception:
            return
        if obj is None:
            return
        cls_name, fields = obj

        # Check for Geometry / Mesh data
        if 'Geometry' in cls_name:
            mesh_id = fields.get('mesh', None)
            if mesh_id and mesh_id > 0:
                try:
                    geo = extract_mesh_geometry(mesh_id, alias_width, classes, location_table, payload, cache)
                    if geo:
                        result['geometries'].append({
                            'name': fields.get('name', 'unknown'),
                            'geometry': geo,
                        })
                except Exception:
                    pass

        # Check for AnimControl
        if 'AnimControl' in cls_name:
            skeleton_id = fields.get('skeleton', None)
            if skeleton_id and skeleton_id > 0:
                skel = extract_skeleton(skeleton_id, alias_width, classes, location_table, payload, cache)
                if skel:
                    result['skeletons'].append(skel)

            anim_map = fields.get('animations', None)
            if isinstance(anim_map, dict):
                anims = extract_animations(anim_map, alias_width, classes, location_table, payload, cache)
                result['animations'].update(anims)

        # Check for SkeletonControl
        if 'SkeletonControl' in cls_name:
            skeleton_id = fields.get('skeleton', None)
            if skeleton_id and skeleton_id > 0 and not result['skeletons']:
                skel = extract_skeleton(skeleton_id, alias_width, classes, location_table, payload, cache)
                if skel:
                    result['skeletons'].append(skel)

        # Recurse into children and controls
        for list_key in ['children', 'controlsList']:
            items = fields.get(list_key, None)
            if isinstance(items, list):
                for item in items:
                    if isinstance(item, int) and item > 0:
                        walk_obj(item, depth + 1)

    walk_obj(root_id)

    return result


# ── Convert to glTF / JSON ──────────────────────────────────────────────────

def j3o_to_gltf(j3o_data, output_path):
    """Convert extracted J3O data to glTF or OBJ-like JSON."""
    geos = j3o_data['geometries']
    skels = j3o_data['skeletons']
    anims = j3o_data['animations']

    if not geos:
        print(f"  No geometry found in {j3o_data['file']}")
        return

    # For now, output as simple OBJ-like JSON (positions + indices)
    for geo_entry in geos:
        g = geo_entry['geometry']
        name = geo_entry['name'] or os.path.basename(j3o_data['file']).replace('.j3o', '')

        out = {
            'name': name,
            'vertexCount': g['vertexCount'],
            'positions': g['positions'],
            'normals': g['normals'],
            'indices': g['indices'],
            'hasBones': len(g['boneWeights']) > 0,
            'boneWeights': g['boneWeights'],
            'boneIndices': g['boneIndices'],
        }

        # Add skeleton if exists
        if skels:
            out['skeleton'] = skels[0]
        if anims:
            out['animations'] = anims

        out_path = os.path.join(output_path, f'{name}.json')
        with open(out_path, 'w') as f:
            json.dump(out, f, separators=(',', ':'))

        print(f"  -> {out_path} ({g['vertexCount']} verts, {len(g['indices'])//3} faces)")


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

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

    j3o_files = [
        ('spider', 'spider/spider.mesh.j3o'),
        ('staff', 'staff/staff.mesh.j3o'),
        ('cross', 'kross/kross.mesh.j3o'),
    ]

    for name, path in j3o_files:
        full = os.path.join(base, path)
        if not os.path.exists(full):
            print(f'Skipping {name}: not found')
            continue
        print(f'Extracting {name}...')
        try:
            data = extract_j3o(full)
            j3o_to_gltf(data, out_dir)
            print(f"  Skeletons: {len(data['skeletons'])}, Animations: {len(data['animations'])}")
        except Exception as e:
            print(f'  ERROR: {e}')
            import traceback
            traceback.print_exc()


if __name__ == '__main__':
    main()
