#!/usr/bin/env python3
"""
Retarget the authored left/right wall-run clips into an existing character GLB.

Unlike retarget.py, this is deliberately animation-only: it does not rebuild
the hierarchy, touch skin weights, subdivide cloth, scale meshes, or rewrite
the rig sidecar.  It is safe to use on a character that has already passed the
full one-time character build.

Usage:
  blender --background --python tools/inject_wallrun.py -- \
      <character.glb> <animation_library.glb> <output.glb> [ClipName ...]

With no clip names it keeps the original wall-run behavior. Supplying names
also makes this the small post-export repair used for libraries that Blender's
first, very large NLA export occasionally omits (notably RunBackward on some
Auto-Rig Pro characters).
"""
import math
import os
import sys

import bpy
from mathutils import Matrix, Quaternion, Vector

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from rig_map import RigRoles, build_map
from retarget_pose import authored_world_rotation, build_segment_pairs


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

CHARACTER, LIBRARY, OUTPUT = argv[:3]
UP = Vector((0.0, 0.0, 1.0))
RAW_CLIPS = tuple(argv[3:]) or ("WallRunLeft", "WallRunRight")
# ``Source=Destination`` lets the compact repair give a source action its
# canonical runtime name (used by strafe_back-cycle -> RunBackward).
CLIP_SPECS = tuple(
    tuple(raw.split("=", 1)) if "=" in raw else (raw, raw)
    for raw in RAW_CLIPS
)
CLIPS = tuple(destination for _source, destination in CLIP_SPECS)


def find_armature(objects):
    return next((obj for obj in objects if obj.type == "ARMATURE"), None)


def assign_action(obj, action):
    if obj.animation_data is None:
        obj.animation_data_create()
    obj.animation_data.action = action
    if action is None:
        return
    try:
        if len(action.slots):
            obj.animation_data.action_slot = action.slots[0]
    except Exception:
        pass


def world_rest(arm):
    return {
        bone.name: arm.matrix_world @ bone.matrix_local
        for bone in arm.data.bones
    }


def bone_order(arm):
    out = []

    def walk(bone):
        out.append(bone.name)
        for child in bone.children:
            walk(child)

    for bone in arm.data.bones:
        if bone.parent is None:
            walk(bone)
    return out


def rig_forward(arm, roles):
    acc = Vector((0.0, 0.0, 0.0))
    for side in ("L", "R"):
        for role in ("toe", "foot"):
            name = roles.limb.get((role, side))
            if not name:
                continue
            bone = arm.data.bones[name]
            direction = (
                (arm.matrix_world @ bone.tail_local)
                - (arm.matrix_world @ bone.matrix_local.translation)
            )
            direction.z = 0.0
            if direction.length > 1e-5:
                acc += direction.normalized()
            break
    return acc.normalized() if acc.length > 1e-5 else None


def facing_correction(src_arm, src_roles, target_arm, target_roles):
    source = rig_forward(src_arm, src_roles)
    target = rig_forward(target_arm, target_roles)
    if source is None or target is None:
        return Quaternion()
    angle = Vector((source.x, source.y)).angle_signed(
        Vector((target.x, target.y)), 0.0
    )
    return Quaternion(UP, angle) if abs(angle) >= 1e-4 else Quaternion()


def solve_pose(arm, order, rest_world, desired_rot, hips, hips_head):
    pose_world = {}
    basis = {}
    bones = arm.data.bones
    for name in order:
        bone = bones[name]
        rest = rest_world[name]
        if bone.parent is not None:
            parent = bone.parent.name
            matrix = pose_world[parent] @ rest_world[parent].inverted() @ rest
        else:
            matrix = rest
        rotation = Quaternion()
        if name in desired_rot:
            rotation = matrix.to_quaternion().inverted() @ desired_rot[name]
        location = Vector((0.0, 0.0, 0.0))
        if name == hips and hips_head is not None:
            location = matrix.inverted() @ hips_head
        basis[name] = (location, rotation)
        pose_world[name] = (
            matrix
            @ Matrix.Translation(location)
            @ rotation.to_matrix().to_4x4()
        )
    return basis


def retarget_action(target_arm, target_roles, src_arm, src_roles, mapping,
                    src_action, clip_name, yaw, scale):
    src_rest = world_rest(src_arm)
    target_rest = world_rest(target_arm)
    target_order = bone_order(target_arm)
    src_rest_rot = {
        name: matrix.to_quaternion() for name, matrix in src_rest.items()
    }
    target_rest_rot = {
        name: matrix.to_quaternion() for name, matrix in target_rest.items()
    }
    yaw_inverse = yaw.inverted()
    segment_pairs = build_segment_pairs(mapping, src_roles, target_roles)
    src_hips = src_roles.hips
    target_hips = target_roles.hips
    src_hips_rest = src_rest[src_hips].translation.copy()
    target_hips_rest = target_rest[target_hips].translation.copy()

    assign_action(src_arm, src_action)
    frame_start, frame_end = (
        int(round(value)) for value in src_action.frame_range
    )
    baked = bpy.data.actions.new(clip_name)
    assign_action(target_arm, baked)
    for pose_bone in target_arm.pose.bones:
        pose_bone.rotation_mode = "QUATERNION"

    scene = bpy.context.scene
    for frame in range(frame_start, frame_end + 1):
        scene.frame_set(frame)
        depsgraph = bpy.context.evaluated_depsgraph_get()
        src_eval = src_arm.evaluated_get(depsgraph)
        source_world = src_eval.matrix_world
        desired = {}
        for target_name, source_name in mapping.items():
            if (
                source_name not in src_eval.pose.bones
                or target_name not in target_rest_rot
            ):
                continue
            desired[target_name] = authored_world_rotation(
                target_name,
                source_name,
                src_eval,
                source_world,
                src_rest_rot,
                target_rest,
                target_rest_rot,
                yaw,
                yaw_inverse,
                segment_pairs,
            )

        moved = (
            source_world @ src_eval.pose.bones[src_hips].matrix
        ).translation
        displacement = yaw @ ((moved - src_hips_rest) * scale)
        # Gameplay owns horizontal travel.  Preserve vertical compression and
        # lift from the performance, but keep the cycle in-place.
        displacement.x = 0.0
        displacement.y = 0.0
        hips_head = target_hips_rest + displacement
        basis = solve_pose(
            target_arm,
            target_order,
            target_rest,
            desired,
            target_hips,
            hips_head,
        )
        for name in mapping:
            if name not in basis:
                continue
            pose_bone = target_arm.pose.bones[name]
            pose_bone.rotation_quaternion = basis[name][1]
            pose_bone.keyframe_insert("rotation_quaternion", frame=frame)
        pose_bone = target_arm.pose.bones[target_hips]
        pose_bone.location = basis[target_hips][0]
        pose_bone.keyframe_insert("location", frame=frame)

    assign_action(target_arm, None)
    return baked


def add_clip_track(arm, action, name):
    action.name = name
    action.use_fake_user = True
    track = arm.animation_data.nla_tracks.new()
    track.name = name
    strip = track.strips.new(name, 0, action)
    strip.name = name
    track.mute = True


def main():
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=CHARACTER)
    target_arm = find_armature(bpy.data.objects)
    if target_arm is None:
        raise RuntimeError(f"{CHARACTER}: no armature")
    target_roles = RigRoles(target_arm)
    missing = target_roles.missing_core()
    if missing:
        raise RuntimeError(f"{CHARACTER}: missing target roles {missing}")

    # Remove the old relabelled sprint, plus stale versions from a prior run.
    for track in list(target_arm.animation_data.nla_tracks):
        if track.name in ("WallRun",) + CLIPS:
            target_arm.animation_data.nla_tracks.remove(track)
    for action in list(bpy.data.actions):
        if action.name in ("WallRun",) + CLIPS:
            bpy.data.actions.remove(action)
    assign_action(target_arm, None)

    before_objects = set(bpy.data.objects)
    before_actions = set(bpy.data.actions)
    bpy.ops.import_scene.gltf(filepath=LIBRARY)
    new_objects = [obj for obj in bpy.data.objects if obj not in before_objects]
    new_actions = [
        action for action in bpy.data.actions if action not in before_actions
    ]
    src_arm = find_armature(new_objects)
    if src_arm is None:
        raise RuntimeError(f"{LIBRARY}: no source armature")
    src_roles = RigRoles(src_arm)
    mapping = build_map(src_roles, target_roles, {})
    if len(mapping) < 20:
        raise RuntimeError(f"{LIBRARY}: only {len(mapping)} mapped bones")

    yaw = facing_correction(src_arm, src_roles, target_arm, target_roles)
    src_h = (
        src_arm.matrix_world
        @ src_arm.data.bones[src_roles.hips].matrix_local
    ).translation.z
    target_h = (
        target_arm.matrix_world
        @ target_arm.data.bones[target_roles.hips].matrix_local
    ).translation.z
    scale = target_h / src_h if src_h > 1e-5 else 1.0

    by_name = {action.name.split(".")[0]: action for action in new_actions}
    for source_name, clip in CLIP_SPECS:
        source = by_name.get(source_name)
        if source is None:
            raise RuntimeError(f"{LIBRARY}: missing {source_name}")
        baked = retarget_action(
            target_arm,
            target_roles,
            src_arm,
            src_roles,
            mapping,
            source,
            clip,
            yaw,
            scale,
        )
        add_clip_track(target_arm, baked, clip)

    for obj in new_objects:
        bpy.data.objects.remove(obj, do_unlink=True)
    for action in new_actions:
        if action.users == 0:
            bpy.data.actions.remove(action)
    assign_action(target_arm, None)

    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_apply=False,
        export_skins=True,
        export_animations=True,
        export_animation_mode="NLA_TRACKS",
        export_bake_animation=False,
        export_optimize_animation_size=True,
        export_optimize_animation_keep_anim_armature=False,
    )
    print(
        f"Injected {', '.join(CLIPS)} into {OUTPUT}; "
        f"{len(bpy.data.actions)} total actions"
    )


main()