TiedUp-/src/main/java/com/tiedup/remake/client/gltf/GltfPoseConverter.java

package com.tiedup.remake.client.gltf;

import dev.kosmx.playerAnim.core.data.AnimationFormat;
import dev.kosmx.playerAnim.core.data.KeyframeAnimation;
import dev.kosmx.playerAnim.core.util.Ease;
import java.util.HashSet;
import java.util.Set;
import org.jetbrains.annotations.Nullable;
import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.joml.Quaternionf;
import org.joml.Vector3f;

/**
 * Converts glTF rest pose + animation quaternions into a PlayerAnimator KeyframeAnimation.
 * <p>
 * Data is expected to be already in MC coordinate space (converted by GlbParser).
 * For upper bones: computes delta quaternion, decomposes to Euler ZYX (pitch/yaw/roll).
 * For lower bones: extracts bend angle from delta quaternion.
 * <p>
 * The GLB model's arm pivots are expected to match MC's exactly (world Y=1.376),
 * so no angle scaling is needed. If the pivots don't match, fix the Blender model.
 * <p>
 * Produces a static looping pose (beginTick=0, endTick=1, looped).
 */
@OnlyIn(Dist.CLIENT)
public final class GltfPoseConverter {

    private static final Logger LOGGER = LogManager.getLogger("GltfPipeline");

    private GltfPoseConverter() {}

    /**
     * Convert a GltfData's rest pose (or first animation frame) to a KeyframeAnimation.
     * Uses the default (first) animation clip.
     * GltfData must already be in MC coordinate space.
     *
     * @param data parsed glTF data (in MC space)
     * @return a static looping KeyframeAnimation suitable for PlayerAnimator
     */
    public static KeyframeAnimation convert(GltfData data) {
        return convertClip(data, data.rawGltfAnimation(), "gltf_pose");
    }

    /**
     * Convert a specific named animation from GltfData to a KeyframeAnimation.
     * Falls back to the default animation if the name is not found.
     *
     * @param data          parsed glTF data (in MC space)
     * @param animationName the name of the animation to convert (e.g. "Struggle", "Idle")
     * @return a static looping KeyframeAnimation suitable for PlayerAnimator
     */
    public static KeyframeAnimation convert(GltfData data, String animationName) {
        GltfData.AnimationClip rawClip = data.getRawAnimation(animationName);
        if (rawClip == null) {
            LOGGER.warn("[GltfPipeline] Animation '{}' not found, falling back to default", animationName);
            return convert(data);
        }
        return convertClip(data, rawClip, "gltf_" + animationName);
    }

    /**
     * Convert a GLB animation with selective part enabling and free-bone support.
     *
     * <p>Owned parts are always enabled in the output animation. Free parts (in
     * {@code enabledParts} but not in {@code ownedParts}) are only enabled if the
     * GLB contains actual keyframe data for them. Parts not in {@code enabledParts}
     * at all are always disabled (pass through to lower layers).</p>
     *
     * @param data          parsed glTF data (in MC space)
     * @param animationName animation name in GLB, or null for default
     * @param ownedParts    parts the item explicitly owns (always enabled)
     * @param enabledParts  parts the item may animate (owned + free); free parts
     *                      are only enabled if the GLB has keyframes for them
     * @return KeyframeAnimation with selective parts active
     */
    public static KeyframeAnimation convertSelective(GltfData data, @Nullable String animationName,
            Set<String> ownedParts, Set<String> enabledParts) {
        GltfData.AnimationClip rawClip;
        String animName;
        if (animationName != null) {
            rawClip = data.getRawAnimation(animationName);
            animName = "gltf_" + animationName;
        } else {
            rawClip = null;
            animName = "gltf_pose";
        }
        if (rawClip == null) {
            rawClip = data.rawGltfAnimation();
        }
        return convertClipSelective(data, rawClip, animName, ownedParts, enabledParts);
    }

    /**
     * Internal: convert a specific raw animation clip with selective part enabling
     * and free-bone support.
     *
     * <p>Tracks which PlayerAnimator parts received actual keyframe data from the GLB.
     * A bone has keyframes if {@code rawClip.rotations()[jointIndex] != null}.
     * This information is used by {@link #enableSelectiveParts} to decide whether
     * free parts should be enabled.</p>
     *
     * @param ownedParts   parts the item explicitly owns (always enabled)
     * @param enabledParts parts the item may animate (owned + free)
     */
    private static KeyframeAnimation convertClipSelective(GltfData data, GltfData.AnimationClip rawClip,
            String animName, Set<String> ownedParts, Set<String> enabledParts) {
        KeyframeAnimation.AnimationBuilder builder =
            new KeyframeAnimation.AnimationBuilder(AnimationFormat.JSON_EMOTECRAFT);

        builder.beginTick = 0;
        builder.endTick = 1;
        builder.stopTick = 1;
        builder.isLooped = true;
        builder.returnTick = 0;
        builder.name = animName;

        String[] jointNames = data.jointNames();
        Quaternionf[] rawRestRotations = data.rawGltfRestRotations();

        // Track which PlayerAnimator part names received actual animation data
        Set<String> partsWithKeyframes = new HashSet<>();

        for (int j = 0; j < data.jointCount(); j++) {
            String boneName = jointNames[j];
            if (!GltfBoneMapper.isKnownBone(boneName)) continue;

            // Check if this joint has explicit animation data (not just rest pose fallback).
            // A bone counts as explicitly animated if it has rotation OR translation keyframes.
            boolean hasExplicitAnim = rawClip != null && (
                (j < rawClip.rotations().length && rawClip.rotations()[j] != null)
                || (rawClip.translations() != null
                    && j < rawClip.translations().length
                    && rawClip.translations()[j] != null)
            );

            Quaternionf animQ = getRawAnimQuaternion(rawClip, rawRestRotations, j);
            Quaternionf restQ = rawRestRotations[j];

            // delta_local = inverse(rest_q) * anim_q  (in bone-local frame)
            Quaternionf deltaLocal = new Quaternionf(restQ).invert().mul(animQ);

            // Convert to PARENT frame: delta_parent = rest * delta_local * inv(rest)
            Quaternionf deltaParent = new Quaternionf(restQ).mul(deltaLocal)
                .mul(new Quaternionf(restQ).invert());

            // Convert from glTF parent frame to MC model-def frame.
            // 180deg rotation around Z (X and Y differ): negate qx and qy.
            Quaternionf deltaQ = new Quaternionf(deltaParent);
            deltaQ.x = -deltaQ.x;
            deltaQ.y = -deltaQ.y;

            if (GltfBoneMapper.isLowerBone(boneName)) {
                convertLowerBone(builder, boneName, deltaQ);
            } else {
                convertUpperBone(builder, boneName, deltaQ);
            }

            // Record which PlayerAnimator part received data
            if (hasExplicitAnim) {
                String animPart = GltfBoneMapper.getAnimPartName(boneName);
                if (animPart != null) {
                    partsWithKeyframes.add(animPart);
                }
                // For lower bones, the keyframe data goes to the upper bone's part
                if (GltfBoneMapper.isLowerBone(boneName)) {
                    String upperBone = GltfBoneMapper.getUpperBoneFor(boneName);
                    if (upperBone != null) {
                        String upperPart = GltfBoneMapper.getAnimPartName(upperBone);
                        if (upperPart != null) {
                            partsWithKeyframes.add(upperPart);
                        }
                    }
                }
            }
        }

        // Selective: enable owned parts always, free parts only if they have keyframes
        enableSelectiveParts(builder, ownedParts, enabledParts, partsWithKeyframes);

        KeyframeAnimation anim = builder.build();
        LOGGER.debug("[GltfPipeline] Converted selective animation '{}' (owned: {}, enabled: {}, withKeyframes: {})",
            animName, ownedParts, enabledParts, partsWithKeyframes);
        return anim;
    }

    /**
     * Add keyframes for specific owned parts from a GLB animation clip to an existing builder.
     *
     * <p>Only writes keyframes for bones that map to a part in {@code ownedParts}.
     * Other bones are skipped entirely. This allows multiple items to contribute
     * to the same animation builder without overwriting each other's keyframes.</p>
     *
     * @param builder     the shared animation builder to add keyframes to
     * @param data        parsed glTF data
     * @param rawClip     the raw animation clip, or null for rest pose
     * @param ownedParts  parts this item exclusively owns (only these get keyframes)
     * @return set of part names that received actual keyframe data from the GLB
     */
    public static Set<String> addBonesToBuilder(
            KeyframeAnimation.AnimationBuilder builder,
            GltfData data, @Nullable GltfData.AnimationClip rawClip,
            Set<String> ownedParts) {

        String[] jointNames = data.jointNames();
        Quaternionf[] rawRestRotations = data.rawGltfRestRotations();
        Set<String> partsWithKeyframes = new HashSet<>();

        for (int j = 0; j < data.jointCount(); j++) {
            String boneName = jointNames[j];
            if (!GltfBoneMapper.isKnownBone(boneName)) continue;

            // Only process bones that belong to this item's owned parts
            String animPart = GltfBoneMapper.getAnimPartName(boneName);
            if (animPart == null || !ownedParts.contains(animPart)) continue;

            // For lower bones, check if the UPPER bone's part is owned
            // (lower bone keyframes go to the upper bone's StateCollection)
            if (GltfBoneMapper.isLowerBone(boneName)) {
                String upperBone = GltfBoneMapper.getUpperBoneFor(boneName);
                if (upperBone != null) {
                    String upperPart = GltfBoneMapper.getAnimPartName(upperBone);
                    if (upperPart == null || !ownedParts.contains(upperPart)) continue;
                }
            }

            boolean hasExplicitAnim = rawClip != null && (
                (j < rawClip.rotations().length && rawClip.rotations()[j] != null)
                || (rawClip.translations() != null
                    && j < rawClip.translations().length
                    && rawClip.translations()[j] != null)
            );

            Quaternionf animQ = getRawAnimQuaternion(rawClip, rawRestRotations, j);
            Quaternionf restQ = rawRestRotations[j];

            Quaternionf deltaLocal = new Quaternionf(restQ).invert().mul(animQ);
            Quaternionf deltaParent = new Quaternionf(restQ).mul(deltaLocal)
                .mul(new Quaternionf(restQ).invert());

            Quaternionf deltaQ = new Quaternionf(deltaParent);
            deltaQ.x = -deltaQ.x;
            deltaQ.y = -deltaQ.y;

            if (GltfBoneMapper.isLowerBone(boneName)) {
                convertLowerBone(builder, boneName, deltaQ);
            } else {
                convertUpperBone(builder, boneName, deltaQ);
            }

            if (hasExplicitAnim) {
                partsWithKeyframes.add(animPart);
                if (GltfBoneMapper.isLowerBone(boneName)) {
                    String upperBone = GltfBoneMapper.getUpperBoneFor(boneName);
                    if (upperBone != null) {
                        String upperPart = GltfBoneMapper.getAnimPartName(upperBone);
                        if (upperPart != null) partsWithKeyframes.add(upperPart);
                    }
                }
            }
        }

        return partsWithKeyframes;
    }

    /**
     * Convert an animation clip using skeleton data from a separate source.
     *
     * <p>This is useful when the animation clip is stored separately from the
     * skeleton (e.g., furniture seat animations where the Player_* armature's
     * clips are parsed into a separate map from the skeleton GltfData).</p>
     *
     * <p>The resulting animation has all parts fully enabled. Callers should
     * create a mutable copy and selectively disable parts as needed.</p>
     *
     * @param skeleton the GltfData providing rest pose, joint names, and joint count
     * @param clip     the raw animation clip (in glTF space) to convert
     * @param animName debug name for the resulting animation
     * @return a static looping KeyframeAnimation with all parts enabled
     */
    public static KeyframeAnimation convertWithSkeleton(
            GltfData skeleton, GltfData.AnimationClip clip, String animName) {
        return convertClip(skeleton, clip, animName);
    }

    /**
     * Internal: convert a specific raw animation clip to a KeyframeAnimation.
     */
    private static KeyframeAnimation convertClip(GltfData data, GltfData.AnimationClip rawClip, String animName) {
        KeyframeAnimation.AnimationBuilder builder =
            new KeyframeAnimation.AnimationBuilder(AnimationFormat.JSON_EMOTECRAFT);

        builder.beginTick = 0;
        builder.endTick = 1;
        builder.stopTick = 1;
        builder.isLooped = true;
        builder.returnTick = 0;
        builder.name = animName;

        String[] jointNames = data.jointNames();
        Quaternionf[] rawRestRotations = data.rawGltfRestRotations();

        for (int j = 0; j < data.jointCount(); j++) {
            String boneName = jointNames[j];

            if (!GltfBoneMapper.isKnownBone(boneName)) continue;

            Quaternionf animQ = getRawAnimQuaternion(rawClip, rawRestRotations, j);
            Quaternionf restQ = rawRestRotations[j];

            // delta_local = inverse(rest_q) * anim_q  (in bone-local frame)
            Quaternionf deltaLocal = new Quaternionf(restQ).invert().mul(animQ);

            // Convert to PARENT frame: delta_parent = rest * delta_local * inv(rest)
            // Simplifies algebraically to: animQ * inv(restQ)
            Quaternionf deltaParent = new Quaternionf(restQ).mul(deltaLocal)
                .mul(new Quaternionf(restQ).invert());

            // Convert from glTF parent frame to MC model-def frame.
            // 180° rotation around Z (X and Y differ): negate qx and qy.
            Quaternionf deltaQ = new Quaternionf(deltaParent);
            deltaQ.x = -deltaQ.x;
            deltaQ.y = -deltaQ.y;

            LOGGER.debug(String.format(
                "[GltfPipeline] Bone '%s': restQ=(%.3f,%.3f,%.3f,%.3f) animQ=(%.3f,%.3f,%.3f,%.3f) deltaQ=(%.3f,%.3f,%.3f,%.3f)",
                boneName,
                restQ.x, restQ.y, restQ.z, restQ.w,
                animQ.x, animQ.y, animQ.z, animQ.w,
                deltaQ.x, deltaQ.y, deltaQ.z, deltaQ.w));

            if (GltfBoneMapper.isLowerBone(boneName)) {
                convertLowerBone(builder, boneName, deltaQ);
            } else {
                convertUpperBone(builder, boneName, deltaQ);
            }
        }

        builder.fullyEnableParts();

        KeyframeAnimation anim = builder.build();
        LOGGER.debug("[GltfPipeline] Converted glTF animation '{}' to KeyframeAnimation", animName);
        return anim;
    }

    /**
     * Get the raw animation quaternion for a joint from a specific clip.
     * Falls back to rest rotation if the clip is null or has no data for this joint.
     */
    private static Quaternionf getRawAnimQuaternion(
        GltfData.AnimationClip rawClip, Quaternionf[] rawRestRotations, int jointIndex
    ) {
        if (rawClip != null && jointIndex < rawClip.rotations().length
                && rawClip.rotations()[jointIndex] != null) {
            return rawClip.rotations()[jointIndex][0]; // first frame
        }
        return rawRestRotations[jointIndex]; // fallback to rest
    }

    private static void convertUpperBone(
        KeyframeAnimation.AnimationBuilder builder,
        String boneName, Quaternionf deltaQ
    ) {
        // Decompose delta quaternion to Euler ZYX
        // JOML's getEulerAnglesZYX stores: euler.x = X rotation, euler.y = Y rotation, euler.z = Z rotation
        // (the "ZYX" refers to rotation ORDER, not storage order)
        Vector3f euler = new Vector3f();
        deltaQ.getEulerAnglesZYX(euler);
        float pitch = euler.x;  // X rotation (pitch)
        float yaw = euler.y;    // Y rotation (yaw)
        float roll = euler.z;   // Z rotation (roll)

        LOGGER.debug(String.format(
            "[GltfPipeline] Upper bone '%s': pitch=%.1f° yaw=%.1f° roll=%.1f°",
            boneName,
            Math.toDegrees(pitch),
            Math.toDegrees(yaw),
            Math.toDegrees(roll)));

        // Get the StateCollection for this body part
        String animPart = GltfBoneMapper.getAnimPartName(boneName);
        if (animPart == null) return;

        KeyframeAnimation.StateCollection part = getPartByName(builder, animPart);
        if (part == null) return;

        part.pitch.addKeyFrame(0, pitch, Ease.CONSTANT);
        part.yaw.addKeyFrame(0, yaw, Ease.CONSTANT);
        part.roll.addKeyFrame(0, roll, Ease.CONSTANT);
    }

    private static void convertLowerBone(
        KeyframeAnimation.AnimationBuilder builder,
        String boneName, Quaternionf deltaQ
    ) {
        // Extract bend angle and axis from the delta quaternion
        float angle = 2.0f * (float) Math.acos(
            Math.min(1.0, Math.abs(deltaQ.w))
        );

        // Determine bend direction from axis
        float bendDirection = 0.0f;
        if (deltaQ.x * deltaQ.x + deltaQ.z * deltaQ.z > 0.001f) {
            bendDirection = (float) Math.atan2(deltaQ.z, deltaQ.x);
        }

        // Sign: if w is negative, the angle wraps
        if (deltaQ.w < 0) {
            angle = -angle;
        }

        LOGGER.debug(String.format(
            "[GltfPipeline] Lower bone '%s': bendAngle=%.1f° bendDir=%.1f°",
            boneName,
            Math.toDegrees(angle),
            Math.toDegrees(bendDirection)));

        // Apply bend to the upper bone's StateCollection
        String upperBone = GltfBoneMapper.getUpperBoneFor(boneName);
        if (upperBone == null) return;

        String animPart = GltfBoneMapper.getAnimPartName(upperBone);
        if (animPart == null) return;

        KeyframeAnimation.StateCollection part = getPartByName(builder, animPart);
        if (part == null || !part.isBendable) return;

        part.bend.addKeyFrame(0, angle, Ease.CONSTANT);
        part.bendDirection.addKeyFrame(0, bendDirection, Ease.CONSTANT);
    }

    private static KeyframeAnimation.StateCollection getPartByName(
        KeyframeAnimation.AnimationBuilder builder, String name
    ) {
        return switch (name) {
            case "head" -> builder.head;
            case "body" -> builder.body;
            case "rightArm" -> builder.rightArm;
            case "leftArm" -> builder.leftArm;
            case "rightLeg" -> builder.rightLeg;
            case "leftLeg" -> builder.leftLeg;
            default -> null;
        };
    }

    /**
     * Enable parts selectively based on ownership and keyframe presence.
     *
     * <ul>
     *   <li>Owned parts: always enabled (the item controls these bones)</li>
     *   <li>Free parts WITH keyframes: enabled (the GLB has animation data for them)</li>
     *   <li>Free parts WITHOUT keyframes: disabled (no data to animate, pass through to context)</li>
     *   <li>Other items' parts: disabled (pass through to their own layer)</li>
     * </ul>
     *
     * @param builder            the animation builder with keyframes already added
     * @param ownedParts         parts the item explicitly owns (always enabled)
     * @param enabledParts       parts the item may animate (owned + free)
     * @param partsWithKeyframes parts that received actual animation data from the GLB
     */
    private static void enableSelectiveParts(
            KeyframeAnimation.AnimationBuilder builder,
            Set<String> ownedParts, Set<String> enabledParts,
            Set<String> partsWithKeyframes) {
        String[] allParts = {"head", "body", "rightArm", "leftArm", "rightLeg", "leftLeg"};
        for (String partName : allParts) {
            KeyframeAnimation.StateCollection part = getPartByName(builder, partName);
            if (part != null) {
                if (ownedParts.contains(partName)) {
                    // Always enable owned parts — the item controls these bones
                    part.fullyEnablePart(false);
                } else if (enabledParts.contains(partName) && partsWithKeyframes.contains(partName)) {
                    // Free part WITH keyframes: enable so the GLB animation drives it
                    part.fullyEnablePart(false);
                } else {
                    // Other item's part, or free part without keyframes: disable.
                    // Disabled parts pass through to the lower-priority context layer.
                    part.setEnabled(false);
                }
            }
        }
    }
}