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

package com.tiedup.remake.client.gltf;

import com.mojang.blaze3d.systems.RenderSystem;
import dev.kosmx.playerAnim.core.util.Pair;
import dev.kosmx.playerAnim.impl.IAnimatedPlayer;
import dev.kosmx.playerAnim.impl.animation.AnimationApplier;
import net.minecraft.client.model.HumanoidModel;
import net.minecraft.client.model.geom.ModelPart;
import net.minecraft.world.entity.LivingEntity;
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.Matrix4f;
import org.joml.Quaternionf;
import org.joml.Vector3f;

/**
 * Reads the LIVE skeleton state from HumanoidModel (after PlayerAnimator + bendy-lib
 * have applied all rotations for the current frame) and produces joint matrices
 * compatible with {@link GltfSkinningEngine#skinVertex}.
 * <p>
 * KEY INSIGHT: The ModelPart xRot/yRot/zRot values set by PlayerAnimator represent
 * DELTA rotations (difference from rest pose) expressed in the MC model-def frame.
 * GltfPoseConverter computed them as parent-frame deltas, decomposed to Euler ZYX.
 * <p>
 * To reconstruct the correct LOCAL rotation for the glTF hierarchy:
 * <pre>
 *   delta    = rotationZYX(zRot, yRot, xRot)   // MC-frame delta from ModelPart
 *   localRot = delta * restQ_mc                  // delta applied on top of local rest
 * </pre>
 * No de-parenting is needed because both delta and restQ_mc are already in the
 * parent's local frame. The MC-to-glTF conjugation (negate qx,qy) is a homomorphism,
 * so frame relationships are preserved through the conversion.
 * <p>
 * For bones WITHOUT a MC ModelPart (root, torso), use the MC-converted rest rotation
 * directly from GltfData.
 */
@OnlyIn(Dist.CLIENT)
public final class GltfLiveBoneReader {

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

    private GltfLiveBoneReader() {}

    // Scratch pools for joint-matrix computation. Render-thread-only
    // (asserted below). Pre-populated Matrix4f slots are reused via
    // set() / identity() / mul(). See GltfSkinningEngine for the twin pool.
    private static Matrix4f[] scratchJointMatrices = new Matrix4f[0];
    private static Matrix4f[] scratchWorldTransforms = new Matrix4f[0];
    private static final Matrix4f scratchLocal = new Matrix4f();

    private static Matrix4f[] ensureScratch(Matrix4f[] current, int needed) {
        if (current.length >= needed) return current;
        Matrix4f[] next = new Matrix4f[needed];
        int i = 0;
        for (; i < current.length; i++) next[i] = current[i];
        for (; i < needed; i++) next[i] = new Matrix4f();
        return next;
    }

    /**
     * Compute joint matrices by reading live skeleton state from the HumanoidModel.
     * <p>
     * For upper bones: reconstructs the MC-frame delta from ModelPart euler angles,
     * then composes with the MC-converted rest rotation to get the local rotation.
     * For lower bones: reads bend values from the entity's AnimationApplier and
     * composes the bend delta with the local rest rotation.
     * For non-animated bones: uses rest rotation from GltfData directly.
     * <p>
     * The resulting joint matrices should match {@link GltfSkinningEngine#computeJointMatrices}
     * when the player is in the rest pose (no animation active).
     *
     * @param model  the HumanoidModel after PlayerAnimator has applied rotations
     * @param data   parsed glTF data (MC-converted)
     * @param entity the living entity being rendered
     * @return live reference to an internal scratch buffer (or null on failure).
     *         Caller MUST consume before the next call to any {@code compute*}
     *         method on this class; do not store.
     */
    public static Matrix4f[] computeJointMatricesFromModel(
        HumanoidModel<?> model,
        GltfData data,
        LivingEntity entity
    ) {
        if (model == null || data == null || entity == null) return null;
        assert RenderSystem.isOnRenderThread()
            : "GltfLiveBoneReader.computeJointMatricesFromModel must run on the render thread (scratch buffers are not thread-safe)";

        int jointCount = data.jointCount();
        scratchJointMatrices = ensureScratch(scratchJointMatrices, jointCount);
        scratchWorldTransforms = ensureScratch(
            scratchWorldTransforms,
            jointCount
        );
        Matrix4f[] jointMatrices = scratchJointMatrices;
        Matrix4f[] worldTransforms = scratchWorldTransforms;

        int[] parents = data.parentJointIndices();
        String[] jointNames = data.jointNames();
        Quaternionf[] restRotations = data.restRotations();
        Vector3f[] restTranslations = data.restTranslations();

        // Get the AnimationApplier for bend values (may be null)
        AnimationApplier emote = getAnimationApplier(entity);

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

            if (GltfBoneMapper.isLowerBone(boneName)) {
                // --- Lower bone: reconstruct from bend values ---
                localRot = computeLowerBoneLocalRotation(
                    boneName,
                    j,
                    restRotations,
                    emote
                );
            } else if (hasUniqueModelPart(boneName)) {
                // --- Upper bone with a unique ModelPart ---
                ModelPart part = GltfBoneMapper.getModelPart(model, boneName);
                if (part != null) {
                    localRot = computeUpperBoneLocalRotation(
                        part,
                        j,
                        restRotations
                    );
                } else {
                    // Fallback: use rest rotation
                    localRot = new Quaternionf(restRotations[j]);
                }
            } else {
                // --- Non-animated bone (root, torso, etc.): use rest rotation ---
                localRot = new Quaternionf(restRotations[j]);
            }

            // Build local transform: translate(restTranslation) * rotate(localRot)
            scratchLocal.identity();
            scratchLocal.translate(restTranslations[j]);
            scratchLocal.rotate(localRot);

            // Compose with parent to get world transform.
            // Same semantics as pre-refactor: treat as root when parent hasn't
            // been processed yet (parents[j] >= j was a null in the old array).
            Matrix4f world = worldTransforms[j];
            if (parents[j] >= 0 && parents[j] < j) {
                world.set(worldTransforms[parents[j]]).mul(scratchLocal);
            } else {
                world.set(scratchLocal);
            }

            // Final joint matrix = worldTransform * inverseBindMatrix
            jointMatrices[j].set(world).mul(data.inverseBindMatrices()[j]);
        }

        return jointMatrices;
    }

    /**
     * Compute local rotation for an upper bone that has a unique ModelPart.
     * <p>
     * ModelPart xRot/yRot/zRot are DELTA rotations (set by PlayerAnimator) expressed
     * as ZYX Euler angles in the MC model-def frame. These deltas were originally
     * computed by GltfPoseConverter as parent-frame quantities.
     * <p>
     * The local rotation for the glTF hierarchy is simply:
     * <pre>
     *   delta    = rotationZYX(zRot, yRot, xRot)
     *   localRot = delta * restQ_mc
     * </pre>
     * No de-parenting is needed: both delta and restQ_mc are already in the parent's
     * frame. The MC-to-glTF negate-xy conjugation is a group homomorphism, preserving
     * the frame relationship.
     */
    private static Quaternionf computeUpperBoneLocalRotation(
        ModelPart part,
        int jointIndex,
        Quaternionf[] restRotations
    ) {
        // Reconstruct the MC-frame delta from ModelPart euler angles.
        Quaternionf delta = new Quaternionf().rotationZYX(
            part.zRot,
            part.yRot,
            part.xRot
        );
        // Local rotation = delta applied on top of the local rest rotation.
        return new Quaternionf(delta).mul(restRotations[jointIndex]);
    }

    /**
     * Compute local rotation for a lower bone (elbow/knee) from bend values.
     * <p>
     * Bend values are read from the entity's AnimationApplier. The bend delta is
     * reconstructed as a quaternion rotation around the bend axis, then composed
     * with the local rest rotation:
     * <pre>
     *   bendQuat = axisAngle(cos(bendAxis)*s, 0, sin(bendAxis)*s, cos(halfAngle))
     *   localRot = bendQuat * restQ_mc
     * </pre>
     * No de-parenting needed — same reasoning as upper bones.
     */
    private static Quaternionf computeLowerBoneLocalRotation(
        String boneName,
        int jointIndex,
        Quaternionf[] restRotations,
        AnimationApplier emote
    ) {
        if (emote != null) {
            // Get the MC part name for the upper bone of this lower bone
            String upperBone = GltfBoneMapper.getUpperBoneFor(boneName);
            String animPartName = (upperBone != null)
                ? GltfBoneMapper.getAnimPartName(upperBone)
                : null;

            if (animPartName != null) {
                Pair<Float, Float> bend = emote.getBend(animPartName);
                if (bend != null) {
                    float bendAxis = bend.getLeft();
                    float bendValue = bend.getRight();

                    // Reconstruct bend as quaternion (this is the delta)
                    float ax = (float) Math.cos(bendAxis);
                    float az = (float) Math.sin(bendAxis);
                    float halfAngle = bendValue * 0.5f;
                    float s = (float) Math.sin(halfAngle);
                    Quaternionf bendQuat = new Quaternionf(
                        ax * s,
                        0,
                        az * s,
                        (float) Math.cos(halfAngle)
                    );

                    // Local rotation = bend delta applied on top of local rest rotation
                    return new Quaternionf(bendQuat).mul(
                        restRotations[jointIndex]
                    );
                }
            }
        }

        // No bend data or no AnimationApplier — use rest rotation (identity delta)
        return new Quaternionf(restRotations[jointIndex]);
    }

    /**
     * Check if a bone name corresponds to a bone that has its OWN unique ModelPart
     * (not just a mapping — it must be the PRIMARY bone for that ModelPart).
     * <p>
     * "torso" maps to model.body but "body" is the primary bone for it.
     * Lower bones share a ModelPart with their upper bone.
     * Unknown bones (e.g., "PlayerArmature") have no ModelPart at all.
     */
    private static boolean hasUniqueModelPart(String boneName) {
        // Bones that should read their rotation from the live HumanoidModel.
        //
        // NOTE: "body" is deliberately EXCLUDED. MC's HumanoidModel is FLAT —
        // body, arms, legs, head are all siblings with ABSOLUTE rotations.
        // But the GLB skeleton is HIERARCHICAL (body → torso → arms).
        // If we read body's live rotation (e.g., attack swing yRot), it propagates
        // to arms/head through the hierarchy, but MC's flat model does NOT do this.
        // Result: cuffs mesh rotates with body during attack while arms stay put.
        //
        // Body rotation effects that matter (sneak lean, sitting) are handled by
        // LivingEntityRenderer's PoseStack transform, which applies to the entire
        // mesh uniformly. No need to read body rotation into joint matrices.
        return switch (boneName) {
            case "head" -> true;
            case "leftUpperArm" -> true;
            case "rightUpperArm" -> true;
            case "leftUpperLeg" -> true;
            case "rightUpperLeg" -> true;
            default -> false; // body, torso, lower bones, unknown
        };
    }

    /**
     * Get the AnimationApplier from an entity, if available.
     * Works for both players (via mixin) and NPCs implementing IAnimatedPlayer.
     */
    private static AnimationApplier getAnimationApplier(LivingEntity entity) {
        if (entity instanceof IAnimatedPlayer animated) {
            try {
                return animated.playerAnimator_getAnimation();
            } catch (Exception e) {
                LOGGER.debug(
                    "[GltfPipeline] Could not get AnimationApplier for {}: {}",
                    entity.getClass().getSimpleName(),
                    e.getMessage()
                );
            }
        }
        return null;
    }
}