package com.tiedup.remake.client.gltf;

import com.mojang.blaze3d.systems.RenderSystem;
import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;
import org.joml.Matrix4f;
import org.joml.Quaternionf;
import org.joml.Vector3f;
import org.joml.Vector4f;

/**
 * CPU-based Linear Blend Skinning (LBS) engine.
 * Computes joint matrices purely from glTF data (rest translations + animation rotations).
 * All data is in MC-converted space (consistent with IBMs and vertex positions).
 *
 * <p><b>Scratch pool</b>: the {@code computeJointMatrices*} methods return a
 * reference to an internal grow-on-demand buffer. The caller MUST consume the
 * returned array before the next call to any {@code compute*} method on this
 * class. Storing the reference across frames produces corrupted output on the
 * next call.</p>
 */
@OnlyIn(Dist.CLIENT)
public final class GltfSkinningEngine {

    private GltfSkinningEngine() {}

    // Scratch pools for joint-matrix computation. Single-threaded access from
    // the render thread only (asserted at call sites). Pre-populated Matrix4f
    // slots are reused via set()/identity()/mul() instead of new Matrix4f(...).
    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 from glTF animation/rest data (default animation).
     * Each joint matrix = worldTransform * inverseBindMatrix.
     * Uses MC-converted glTF data throughout for consistency.
     *
     * @param data parsed glTF data (MC-converted)
     * @return live reference to an internal scratch buffer. Caller MUST consume
     *         before the next call to any {@code compute*} method; do not store.
     */
    public static Matrix4f[] computeJointMatrices(GltfData data) {
        assert RenderSystem.isOnRenderThread()
            : "GltfSkinningEngine.computeJointMatrices must run on the render thread (scratch buffers are not thread-safe)";
        return computeJointMatricesFromClip(data, data.animation());
    }

    /**
     * Compute joint matrices with frame interpolation for animated entities.
     * Uses SLERP for rotations and LERP for translations between adjacent keyframes.
     *
     * <p>The {@code time} parameter is in frame-space: 0.0 corresponds to the first
     * keyframe and {@code frameCount - 1} to the last. Values between integer frames
     * are interpolated. Out-of-range values are clamped.</p>
     *
     * @param data the parsed glTF data (MC-converted)
     * @param clip the animation clip to sample (null = rest pose for all joints)
     * @param time time in frame-space (0.0 = first frame, N-1 = last frame)
     * @return live reference to an internal scratch buffer. Caller MUST consume
     *         before the next call to any {@code compute*} method; do not store.
     */
    public static Matrix4f[] computeJointMatricesAnimated(
        GltfData data,
        GltfData.AnimationClip clip,
        float time
    ) {
        assert RenderSystem.isOnRenderThread()
            : "GltfSkinningEngine.computeJointMatricesAnimated 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();

        for (int j = 0; j < jointCount; j++) {
            // Build local transform: translate(interpT) * rotate(interpQ)
            scratchLocal.identity();
            scratchLocal.translate(getInterpolatedTranslation(data, clip, j, time));
            scratchLocal.rotate(getInterpolatedRotation(data, clip, j, time));

            // Compose with parent. Same semantics as the previous allocating
            // code path: only use the parent when its index is already processed
            // (parents[j] < j). Out-of-order/root → treat as identity parent.
            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;
    }

    /**
     * Internal: compute joint matrices from a specific animation clip.
     */
    private static Matrix4f[] computeJointMatricesFromClip(
        GltfData data,
        GltfData.AnimationClip clip
    ) {
        assert RenderSystem.isOnRenderThread()
            : "GltfSkinningEngine.computeJointMatricesFromClip 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();

        for (int j = 0; j < jointCount; j++) {
            // Build local transform: translate(animT or restT) * rotate(animQ or restQ)
            scratchLocal.identity();
            scratchLocal.translate(getAnimTranslation(data, clip, j));
            scratchLocal.rotate(getAnimRotation(data, clip, j));

            // Compose with parent — see note in computeJointMatricesAnimated.
            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;
    }

    /**
     * Get the animation rotation for a joint (MC-converted).
     * Falls back to rest rotation if no animation or the channel is empty.
     */
    private static Quaternionf getAnimRotation(
        GltfData data,
        GltfData.AnimationClip clip,
        int jointIndex
    ) {
        if (
            clip != null &&
            clip.rotations() != null &&
            jointIndex < clip.rotations().length &&
            clip.rotations()[jointIndex] != null &&
            clip.rotations()[jointIndex].length > 0
        ) {
            return clip.rotations()[jointIndex][0]; // first frame
        }
        return data.restRotations()[jointIndex];
    }

    /**
     * Get the animation translation for a joint (MC-converted).
     * Falls back to rest translation if no animation translation exists or the
     * channel is empty.
     */
    private static Vector3f getAnimTranslation(
        GltfData data,
        GltfData.AnimationClip clip,
        int jointIndex
    ) {
        if (
            clip != null &&
            clip.translations() != null &&
            jointIndex < clip.translations().length &&
            clip.translations()[jointIndex] != null &&
            clip.translations()[jointIndex].length > 0
        ) {
            return clip.translations()[jointIndex][0]; // first frame
        }
        return data.restTranslations()[jointIndex];
    }

    // ---- Interpolated accessors (for computeJointMatricesAnimated) ----

    /**
     * Get an interpolated rotation for a joint at a fractional frame time.
     * Uses SLERP between the two bounding keyframes.
     *
     * <p>Falls back to rest rotation when the clip is null or has no rotation
     * data for the given joint. A single-frame channel returns that frame directly.</p>
     *
     * @param data       parsed glTF data
     * @param clip       animation clip (may be null)
     * @param jointIndex joint to query
     * @param time       frame-space time (clamped internally)
     * @return new Quaternionf with the interpolated rotation (never mutates source data)
     */
    private static Quaternionf getInterpolatedRotation(
        GltfData data,
        GltfData.AnimationClip clip,
        int jointIndex,
        float time
    ) {
        if (
            clip == null ||
            jointIndex >= clip.rotations().length ||
            clip.rotations()[jointIndex] == null
        ) {
            // No animation data for this joint -- use rest pose (copy to avoid mutation)
            Quaternionf rest = data.restRotations()[jointIndex];
            return new Quaternionf(rest);
        }

        Quaternionf[] frames = clip.rotations()[jointIndex];
        if (frames.length == 1) {
            return new Quaternionf(frames[0]);
        }

        // Clamp time to valid range [0, frameCount-1]
        float clamped = Math.max(0.0f, Math.min(time, frames.length - 1));
        int f0 = (int) Math.floor(clamped);
        int f1 = Math.min(f0 + 1, frames.length - 1);
        float alpha = clamped - f0;

        if (alpha < 1e-6f || f0 == f1) {
            return new Quaternionf(frames[f0]);
        }

        // SLERP: create a copy of frame0 and slerp toward frame1
        return new Quaternionf(frames[f0]).slerp(frames[f1], alpha);
    }

    /**
     * Get an interpolated translation for a joint at a fractional frame time.
     * Uses LERP between the two bounding keyframes.
     *
     * <p>Falls back to rest translation when the clip is null, the clip has no
     * translation data at all, or has no translation data for the given joint.
     * A single-frame channel returns that frame directly.</p>
     *
     * @param data       parsed glTF data
     * @param clip       animation clip (may be null)
     * @param jointIndex joint to query
     * @param time       frame-space time (clamped internally)
     * @return new Vector3f with the interpolated translation (never mutates source data)
     */
    private static Vector3f getInterpolatedTranslation(
        GltfData data,
        GltfData.AnimationClip clip,
        int jointIndex,
        float time
    ) {
        if (
            clip == null ||
            clip.translations() == null ||
            jointIndex >= clip.translations().length ||
            clip.translations()[jointIndex] == null
        ) {
            // No animation data for this joint -- use rest pose (copy to avoid mutation)
            Vector3f rest = data.restTranslations()[jointIndex];
            return new Vector3f(rest);
        }

        Vector3f[] frames = clip.translations()[jointIndex];
        if (frames.length == 1) {
            return new Vector3f(frames[0]);
        }

        // Clamp time to valid range [0, frameCount-1]
        float clamped = Math.max(0.0f, Math.min(time, frames.length - 1));
        int f0 = (int) Math.floor(clamped);
        int f1 = Math.min(f0 + 1, frames.length - 1);
        float alpha = clamped - f0;

        if (alpha < 1e-6f || f0 == f1) {
            return new Vector3f(frames[f0]);
        }

        // LERP: create a copy of frame0 and lerp toward frame1
        return new Vector3f(frames[f0]).lerp(frames[f1], alpha);
    }

    /**
     * Skin a single vertex using Linear Blend Skinning.
     *
     * <p>Callers should pre-allocate {@code tmpPos} and {@code tmpNorm} and reuse
     * them across all vertices in a mesh to avoid per-vertex allocations (12k+
     * allocations per frame for a typical mesh).</p>
     *
     * @param data         parsed glTF data
     * @param vertexIdx    index into the vertex arrays
     * @param jointMatrices joint matrices from computeJointMatrices
     * @param outPos       output skinned position (3 floats)
     * @param outNormal    output skinned normal (3 floats)
     * @param tmpPos       pre-allocated scratch Vector4f for position transforms
     * @param tmpNorm      pre-allocated scratch Vector4f for normal transforms
     */
    public static void skinVertex(
        GltfData data,
        int vertexIdx,
        Matrix4f[] jointMatrices,
        float[] outPos,
        float[] outNormal,
        Vector4f tmpPos,
        Vector4f tmpNorm
    ) {
        float[] positions = data.positions();
        float[] normals = data.normals();
        int[] joints = data.joints();
        float[] weights = data.weights();

        // Rest position
        float vx = positions[vertexIdx * 3];
        float vy = positions[vertexIdx * 3 + 1];
        float vz = positions[vertexIdx * 3 + 2];

        // Rest normal
        float nx = normals[vertexIdx * 3];
        float ny = normals[vertexIdx * 3 + 1];
        float nz = normals[vertexIdx * 3 + 2];

        // LBS: v_skinned = Σ(w[i] * jointMatrix[j[i]] * v_rest)
        float sx = 0,
            sy = 0,
            sz = 0;
        float snx = 0,
            sny = 0,
            snz = 0;

        int jointCount = data.jointCount();
        for (int i = 0; i < 4; i++) {
            int ji = joints[vertexIdx * 4 + i];
            float w = weights[vertexIdx * 4 + i];
            // Guard against jointCount, NOT jointMatrices.length. The scratch
            // pool (P2-04) means jointMatrices may be longer than jointCount,
            // with trailing slots holding stale matrices from the previous
            // item's skeleton. Tightens the bound back to the pre-scratch
            // semantics. Closes B-batch review RISK-E01.
            if (w <= 0.0f || ji >= jointCount) continue;

            Matrix4f jm = jointMatrices[ji];

            // Transform position
            tmpPos.set(vx, vy, vz, 1.0f);
            jm.transform(tmpPos);
            sx += w * tmpPos.x;
            sy += w * tmpPos.y;
            sz += w * tmpPos.z;

            // Transform normal (ignore translation)
            tmpNorm.set(nx, ny, nz, 0.0f);
            jm.transform(tmpNorm);
            snx += w * tmpNorm.x;
            sny += w * tmpNorm.y;
            snz += w * tmpNorm.z;
        }

        outPos[0] = sx;
        outPos[1] = sy;
        outPos[2] = sz;

        // Normalize the normal
        float len = (float) Math.sqrt(snx * snx + sny * sny + snz * snz);
        if (len > 0.0001f) {
            outNormal[0] = snx / len;
            outNormal[1] = sny / len;
            outNormal[2] = snz / len;
        } else {
            outNormal[0] = 0;
            outNormal[1] = 1;
            outNormal[2] = 0;
        }
    }
}
