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

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

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

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

    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 (v2+)
        per bone:
            uint32       nameLen
            nameLen bytes  UTF-8 name (no 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

    # --- Skip armature matrix (5 lines) ---
    for _ in range(5):
        next_line()

    # --- 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]]

    # --- Vertices ---
    line = next_line()  # "===Vertices: 5140"
    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: 8602"
    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)  # 6 floats

    # --- Normals ---
    next_line()  # "===Normals:"

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

    # --- Triangles ---
    line = next_line()  # "===Triangles: 8602"
    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 ..."

    vertex_weights = []
    for _ in range(num_vertices):
        next_line()  # "Vertex N:"
        line = next_line()  # "Vertex groups: 2"
        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)

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

    keyframes = []
    for _ in range(num_keyframes):
        line = next_line()  # "Frame: 0"
        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'BSAF')
        out.write(struct.pack('<I', 2))

        # 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 (v2+)
        for name in bone_names:
            name_bytes = name.encode('utf-8')
            out.write(struct.pack('<I', len(name_bytes)))
            out.write(name_bytes)

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

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

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

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

        # Vertex Weights
        for vw in 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}, Vertices: {num_vertices}, Faces: {num_faces}, "
          f"Triangles: {num_triangles}, Keyframes: {num_keyframes}")


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])