#!/usr/bin/env python3
"""
Build the compact, animation-only wall-run library used by retarget.py.

The source clips are Uisco's authored UE4 Mannequin wall-run pair:
https://uisco.itch.io/wall-running-animations

Usage:
  blender --background --python tools/build_wallrun_library.py -- \
      <wall_run_left.FBX> <wall_run_right.FBX> <output.glb>

The two FBX files use the same UE4 mannequin skeleton.  This script keeps one
copy of that skeleton, names the actions explicitly, and exports no mesh.
"""
import os
import sys

import bpy


argv = sys.argv
argv = argv[argv.index("--") + 1:] if "--" in argv else []
if len(argv) != 3:
    print(__doc__)
    sys.exit(1)

LEFT, RIGHT, OUTPUT = argv


def import_clip(path, name):
    before_objects = set(bpy.data.objects)
    before_actions = set(bpy.data.actions)
    bpy.ops.import_scene.fbx(filepath=path)
    objects = [obj for obj in bpy.data.objects if obj not in before_objects]
    arms = [obj for obj in objects if obj.type == "ARMATURE"]
    actions = [action for action in bpy.data.actions if action not in before_actions]
    if len(arms) != 1 or len(actions) != 1:
        raise RuntimeError(
            f"{path}: expected one armature/action, got "
            f"{len(arms)} armatures and {len(actions)} actions"
        )
    actions[0].name = name
    actions[0].use_fake_user = True
    return arms[0], actions[0], objects


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    left_arm, left_action, left_objects = import_clip(LEFT, "WallRunLeft")
    right_arm, right_action, right_objects = import_clip(RIGHT, "WallRunRight")

    left_bones = [bone.name for bone in left_arm.data.bones]
    right_bones = [bone.name for bone in right_arm.data.bones]
    if left_bones != right_bones:
        raise RuntimeError("left/right source skeletons do not match")

    for obj in left_objects + right_objects:
        if obj not in (left_arm, right_arm):
            bpy.data.objects.remove(obj, do_unlink=True)
    bpy.data.objects.remove(right_arm, do_unlink=True)

    if left_arm.animation_data is None:
        left_arm.animation_data_create()
    left_arm.animation_data.action = None
    for track in list(left_arm.animation_data.nla_tracks):
        left_arm.animation_data.nla_tracks.remove(track)
    for action in (left_action, right_action):
        track = left_arm.animation_data.nla_tracks.new()
        track.name = action.name
        strip = track.strips.new(action.name, 0, action)
        strip.name = action.name
        track.mute = True

    for action in list(bpy.data.actions):
        if action not in (left_action, right_action):
            bpy.data.actions.remove(action)

    os.makedirs(os.path.dirname(os.path.abspath(OUTPUT)), exist_ok=True)
    bpy.ops.export_scene.gltf(
        filepath=OUTPUT,
        export_format="GLB",
        export_yup=True,
        export_animations=True,
        export_animation_mode="NLA_TRACKS",
        export_skins=True,
        export_bake_animation=False,
        export_optimize_animation_size=True,
    )
    print(
        f"Exported WallRunLeft/WallRunRight on "
        f"{len(left_arm.data.bones)} bones to {OUTPUT}"
    )


main()