#version 460

#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_buffer_reference : require

#include "vertex.glsl"

struct SkinningDataType {
    ivec4 jointIds;
    vec4 weights;
};

// Keep buffer-reference pointer alignment explicit so the C++ push-constant
// struct can safely use 64-bit VkDeviceAddress fields.
layout (buffer_reference, std430, buffer_reference_align = 8) readonly buffer SkinningData {
    SkinningDataType data[];
};

layout (buffer_reference, std430, buffer_reference_align = 8) readonly buffer JointMatrices {
    mat4 matrices[];
};

layout (push_constant) uniform constants
{
    JointMatrices jointMatrices;
    uint jointMatricesStartIndex;
    uint numVertices;
    VertexBuffer inputBuffer;
    SkinningData skinningData;
    WritableVertexBuffer outputBuffer;
} pcs;

layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in;

mat4 getJointMatrix(int jointId) {
    if (jointId < 0) {
        return mat4(1.0);
    }

    return pcs.jointMatrices.matrices[pcs.jointMatricesStartIndex + uint(jointId)];
}

vec3 safeNormalize(vec3 v, vec3 fallback) {
    float len2 = dot(v, v);
    if (len2 <= 1e-20) {
        return fallback;
    }
    return v * inversesqrt(len2);
}

mat4 buildSkinMatrix(SkinningDataType sd) {
    // If imported weights are slightly imperfect, normalize them here as a
    // second line of defense. If a vertex has no weights at all, pass it
    // through unchanged instead of collapsing it to vec3(0).
    float weightSum = sd.weights.x + sd.weights.y + sd.weights.z + sd.weights.w;
    if (weightSum <= 1e-8) {
        return mat4(1.0);
    }

    vec4 w = sd.weights / weightSum;

    return
        w.x * getJointMatrix(sd.jointIds.x) +
        w.y * getJointMatrix(sd.jointIds.y) +
        w.z * getJointMatrix(sd.jointIds.z) +
        w.w * getJointMatrix(sd.jointIds.w);
}

void main()
{
    uint index = gl_GlobalInvocationID.x;
    if (index >= pcs.numVertices) {
        return;
    }

    Vertex v = pcs.inputBuffer.vertices[index];
    SkinningDataType sd = pcs.skinningData.data[index];
    mat4 skinMatrix = buildSkinMatrix(sd);

    v.position = vec3(skinMatrix * vec4(v.position, 1.0));

    // For rigid joint matrices this is correct. If you allow non-uniform scale
    // in bones, replace this with a proper inverse-transpose normal matrix path.
    mat3 skinMat3 = mat3(skinMatrix);
    v.normal = safeNormalize(skinMat3 * v.normal, v.normal);
    v.tangent.xyz = safeNormalize(skinMat3 * v.tangent.xyz, v.tangent.xyz);
    // v.tangent.w is handedness; do not transform it.

    pcs.outputBuffer.vertices[index] = v;
}