#!/usr/bin/env python3
"""
Convert a text-based multi-mesh bone animation file to the BSMF binary format.

Usage:
    python convert_anim_to_binary_new.py <input.txt> <output.bin>

Binary format (BSMF v2) -- all values little-endian:

    HEADER
        4 bytes   magic "BSMF"
        uint32    version (2)

    ARMATURE MATRIX
        16 x float   4x4 matrix (row-major)

    BONES
        uint32    numBones
        per bone:
            3 x float    boneStartWorld  (from HEAD_LOCAL)
            float        boneLength
            9 x float    3x3 rotation matrix (row-major)
            int32        parentIndex (-1 if none)
            uint32       numChildren
            numChildren x int32  childIndices

    BONE NAMES
        per bone:
            uint32       nameLen
            nameLen bytes  UTF-8 name (no terminator)

    MESHES
        uint32    numMeshes
        per mesh:
            uint32    nameLength
            nameLength x char    meshName (UTF-8, no null terminator)

            VERTICES
                uint32    numVertices
                numVertices x 3 x float   positions

            UV COORDINATES
                uint32    numFaces
                numFaces x 6 x float     3 UV pairs per face (u0,v0,u1,v1,u2,v2)

            NORMALS
                numVertices x 3 x float   normals

            TRIANGLES
                uint32    numTriangles
                numTriangles x 3 x int32  vertex indices

            VERTEX WEIGHTS
                per vertex (numVertices):
                    uint32    numGroups
                    numGroups x (int32 boneIndex, float weight)

    ANIMATION KEYFRAMES
        uint32    numKeyframes
        per keyframe:
            int32     frameNumber
            per bone (numBones, in index order 0..N-1):
                3 x float    location
                16 x float   4x4 matrix (row-major)
"""

import struct
import re
import sys


def parse_floats(line):
    return [float(x) for x in re.findall(r'[-]?\d+\.\d+', line)]


def parse_first_int(line):
    m = re.search(r'\d+', line)
    if m:
        return int(m.group())
    raise ValueError(f"No integer found in: {line}")


def parse_children(line):
    return re.findall(r"'([^']+)'", line)


def convert(input_path, output_path):
    with open(input_path, 'r', encoding='utf-8', errors='replace') as f:
        lines = f.readlines()

    idx = 0

    def next_line():
        nonlocal idx
        line = lines[idx].rstrip()
        idx += 1
        return line

    # --- Armature matrix header + 4 rows ---
    next_line()  # "=== Armature Matrix ==="
    armature_matrix = []
    for _ in range(4):
        armature_matrix.extend(parse_floats(next_line())[:4])

    # --- Bone count ---
    line = next_line()  # "=== Armature Bones: 65"
    num_bones = parse_first_int(line)

    bone_names = []
    bones = []
    bone_parent_names = []
    bone_children_names = []

    for _ in range(num_bones):
        bone = {}

        # "Bone: mixamorig:Hips"
        line = next_line()
        bone_name = line[6:]
        bone_names.append(bone_name)

        # "  HEAD_LOCAL: <Vector (x, y, z)>"
        line = next_line()
        bone['head'] = parse_floats(line)[:3]

        # "  TAIL_LOCAL: ..." -- skip
        next_line()

        # "  Length: 0.123"
        line = next_line()
        bone['length'] = parse_floats(line)[0]

        # 3x3 matrix (3 rows)
        mat = []
        for _ in range(3):
            mat.extend(parse_floats(next_line()))
        bone['matrix_3x3'] = mat

        # "  Parent: None" or "  Parent: boneName"
        line = next_line()
        if line == "  Parent: None":
            bone_parent_names.append(None)
        else:
            bone_parent_names.append(line[10:])

        # "  Children: ['a', 'b'] or []"
        line = next_line()
        bone_children_names.append(parse_children(line))

        bones.append(bone)

    # Build name -> index map
    name_to_idx = {name: i for i, name in enumerate(bone_names)}

    # Resolve parent / child indices
    for i in range(num_bones):
        if bone_parent_names[i] is None:
            bones[i]['parent'] = -1
        else:
            bones[i]['parent'] = name_to_idx[bone_parent_names[i]]
        bones[i]['children'] = [name_to_idx[c] for c in bone_children_names[i]]

    # --- Multi-mesh header ---
    line = next_line()  # "=== TOTAL MESHES TO EXPORT: 7 ==="
    num_meshes = parse_first_int(line)

    meshes = []

    for _ in range(num_meshes):
        # "=== Mesh Object: Name ==="
        line = next_line()
        m = re.match(r"===\s*Mesh Object:\s*(.+?)\s*===$", line)
        if not m:
            raise ValueError(f"Invalid mesh header: {line}")
        mesh_name = m.group(1)

        # --- Vertices ---
        line = next_line()  # "===Vertices: N"
        num_vertices = parse_first_int(line)

        vertices = []
        for _ in range(num_vertices):
            vertices.append(parse_floats(next_line())[:3])

        # --- UV Coordinates ---
        next_line()  # "===UV Coordinates:"
        line = next_line()  # "Face count: M"
        num_faces = parse_first_int(line)

        uvs = []
        for _ in range(num_faces):
            next_line()  # "Face N"
            next_line()  # "UV Count: 3"
            face_uvs = []
            for _ in range(3):
                face_uvs.extend(parse_floats(next_line())[:2])
            uvs.append(face_uvs)

        # --- Normals ---
        next_line()  # "===Normals:"
        normals = []
        for _ in range(num_vertices):
            normals.append(parse_floats(next_line())[:3])

        # --- Triangles ---
        line = next_line()  # "===Triangles: M"
        num_triangles = parse_first_int(line)

        triangles = []
        for _ in range(num_triangles):
            line = next_line()
            ints = [int(x) for x in re.findall(r'[-]?\d+', line)]
            triangles.append(ints[:3])

        # --- Vertex Weights ---
        next_line()  # "=== Vertex Weights (Max 5 bones per vertex) ==="
        vertex_weights = []
        for _ in range(num_vertices):
            next_line()  # "Vertex N:"
            line = next_line()  # "Vertex groups: K"
            num_groups = parse_first_int(line)

            groups = []
            for _ in range(num_groups):
                line = next_line()
                m = re.search(r"'([^']+)'.*?([-]?\d+\.\d+)", line)
                bone_name = m.group(1)
                weight = float(m.group(2))
                groups.append((name_to_idx[bone_name], weight))

            vertex_weights.append(groups)

        meshes.append({
            'name': mesh_name,
            'num_vertices': num_vertices,
            'vertices': vertices,
            'num_faces': num_faces,
            'uvs': uvs,
            'normals': normals,
            'num_triangles': num_triangles,
            'triangles': triangles,
            'vertex_weights': vertex_weights,
        })

    # --- Animation Keyframes ---
    next_line()  # "=== Animation Keyframes ==="
    next_line()  # "=== Bone Transforms per Keyframe ==="
    line = next_line()  # "Keyframes: N"
    num_keyframes = parse_first_int(line)

    keyframes = []
    for _ in range(num_keyframes):
        line = next_line()  # "Frame: N"
        frame_number = parse_first_int(line)

        bone_data = {}
        for _ in range(num_bones):
            line = next_line()  # "  Bone: mixamorig:Hips"
            bone_name = line.strip()
            if bone_name.startswith("Bone: "):
                bone_name = bone_name[6:]
            bone_idx = name_to_idx[bone_name]

            # Location
            location = parse_floats(next_line())[:3]

            # Rotation (skip)
            next_line()

            # "    Matrix:" (skip header)
            next_line()

            # 4 rows of 4 floats
            matrix = []
            for _ in range(4):
                matrix.extend(parse_floats(next_line()))

            bone_data[bone_idx] = {
                'location': location,
                'matrix': matrix,
            }

        keyframes.append((frame_number, bone_data))

    # ================================================================
    # Write binary file
    # ================================================================
    with open(output_path, 'wb') as out:
        # Header
        out.write(b'BSMF')
        out.write(struct.pack('<I', 2))

        # Armature matrix (16 floats, row-major)
        out.write(struct.pack('<16f', *armature_matrix))

        # Bones
        out.write(struct.pack('<I', num_bones))
        for i in range(num_bones):
            b = bones[i]
            out.write(struct.pack('<3f', *b['head']))
            out.write(struct.pack('<f', b['length']))
            out.write(struct.pack('<9f', *b['matrix_3x3']))
            out.write(struct.pack('<i', b['parent']))
            out.write(struct.pack('<I', len(b['children'])))
            for c in b['children']:
                out.write(struct.pack('<i', c))

        # Bone names
        for name in bone_names:
            name_bytes = name.encode('utf-8')
            out.write(struct.pack('<I', len(name_bytes)))
            out.write(name_bytes)

        # Meshes
        out.write(struct.pack('<I', num_meshes))
        for md in meshes:
            name_bytes = md['name'].encode('utf-8')
            out.write(struct.pack('<I', len(name_bytes)))
            out.write(name_bytes)

            # Vertices
            out.write(struct.pack('<I', md['num_vertices']))
            for v in md['vertices']:
                out.write(struct.pack('<3f', *v))

            # UV Coordinates
            out.write(struct.pack('<I', md['num_faces']))
            for uv in md['uvs']:
                out.write(struct.pack('<6f', *uv))

            # Normals
            for n in md['normals']:
                out.write(struct.pack('<3f', *n))

            # Triangles
            out.write(struct.pack('<I', md['num_triangles']))
            for t in md['triangles']:
                out.write(struct.pack('<3i', *t))

            # Vertex weights
            for vw in md['vertex_weights']:
                out.write(struct.pack('<I', len(vw)))
                for bone_idx, weight in vw:
                    out.write(struct.pack('<if', bone_idx, weight))

        # Animation Keyframes
        out.write(struct.pack('<I', num_keyframes))
        for frame_num, bone_data in keyframes:
            out.write(struct.pack('<i', frame_num))
            for i in range(num_bones):
                bd = bone_data[i]
                out.write(struct.pack('<3f', *bd['location']))
                out.write(struct.pack('<16f', *bd['matrix']))

    input_size = sum(len(l) for l in lines)
    import os
    output_size = os.path.getsize(output_path)
    print(f"Converted: {input_path} ({input_size:,} bytes text) -> {output_path} ({output_size:,} bytes binary)")
    print(f"  Bones: {num_bones}, Meshes: {num_meshes}, Keyframes: {num_keyframes}")
    for md in meshes:
        print(f"    - {md['name']}: {md['num_vertices']} verts, "
              f"{md['num_faces']} faces, {md['num_triangles']} tris")


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print(f"Usage: {sys.argv[0]} <input.txt> <output.bin>")
        sys.exit(1)

    convert(sys.argv[1], sys.argv[2])