#include "F:/csharp/GhostEngine/src/Runtime//Ghost.Graphics/Shaders/Includes/Properties.hlsl"
#include "F:/csharp/GhostEngine/src/Runtime//Ghost.Graphics/Shaders/Includes/Common.hlsl"

struct PixelInput
{
    float4 position : SV_POSITION;
    float4 color : COLOR;
    float4 uv : TEXCOORD0;
    nointerpolation uint meshletID : MESHLET_ID;
};

[numthreads(128, 1, 1)] // 128 threads to cover max 64 vertices and 124 triangles
[outputtopology("triangle")]
void MSMain(
    uint3 groupThreadID : SV_GroupThreadID,
    uint3 groupID : SV_GroupID,
    out vertices PixelInput outVerts[64],
    out indices uint3 outTris[124])
{
    InstanceData instanceData = LoadData<InstanceData>(g_PushConstantData.instanceBuffer, g_PushConstantData.instanceIndex);
    MeshData meshData = LoadData<MeshData>(instanceData.meshBuffer, 0);

    ByteAddressBuffer meshletBuffer = GET_BUFFER(meshData.meshletBuffer);
    Meshlet m = meshletBuffer.Load<Meshlet>(groupID.x * sizeof(Meshlet));

    uint vertexCount = m.packedCounts & 0xFF;
    uint triangleCount = (m.packedCounts >> 8) & 0xFF;

    SetMeshOutputCounts(vertexCount, triangleCount);

    ByteAddressBuffer meshletVerticesBuffer = GET_BUFFER(meshData.meshletVerticesBuffer);
    ByteAddressBuffer meshletTrianglesBuffer = GET_BUFFER(meshData.meshletTrianglesBuffer);

    // Write vertex output
    if (groupThreadID.x < vertexCount)
    {
        uint vertexIndex = meshletVerticesBuffer.Load((m.vertexOffset + groupThreadID.x) * 4);
        ByteAddressBuffer vertices = GET_BUFFER(meshData.vertexBuffer);
        Vertex v = vertices.Load<Vertex>(vertexIndex * sizeof(Vertex));

        FrameData globalFrameData = LoadData<FrameData>(g_PushConstantData.frameBuffer, 0);
        ViewData viewData = LoadData<ViewData>(g_PushConstantData.viewBuffer, 0);

        float4 worldPos = mul(instanceData.localToWorld, float4(v.position.xyz, 1.0f));
        float4 viewPos = mul(viewData.viewMatrix, worldPos);

        outVerts[groupThreadID.x].position = mul(viewData.projectionMatrix, viewPos);

        outVerts[groupThreadID.x].color = v.color;
        outVerts[groupThreadID.x].uv = v.uv;
        outVerts[groupThreadID.x].meshletID = groupID.x;
    }

    // Write triangle output (1 thread processes 1 triangle)
    if (groupThreadID.x < triangleCount)
    {
        uint triangleIndex = groupThreadID.x;

        // Load the packed 32-bit integer containing the 3 local indices
        uint packedIndices = meshletTrianglesBuffer.Load((m.triangleOffset + triangleIndex) * 4);

        uint i0 = packedIndices & 0xFF;
        uint i1 = (packedIndices >> 8) & 0xFF;
        uint i2 = (packedIndices >> 16) & 0xFF;

        outTris[triangleIndex] = uint3(i0, i1, i2);
    }
}

float4 PSMain(PixelInput input) : SV_TARGET
{
    // PerMaterialData perMaterialData = LoadData<PerMaterialData>(g_PushConstantData.materialIndex, 0);
    //
    // float4 color1 = SAMPLE_TEXTURE2D(perMaterialData.texture1, perMaterialData.tex_sampler, input.uv.xy);
    // float4 color2 = SAMPLE_TEXTURE2D(perMaterialData.texture2, perMaterialData.tex_sampler, input.uv.xy);
    // float4 color3 = SAMPLE_TEXTURE2D(perMaterialData.texture3, perMaterialData.tex_sampler, input.uv.xy);
    // float4 color4 = SAMPLE_TEXTURE2D(perMaterialData.texture4, perMaterialData.tex_sampler, input.uv.xy);
    //
    // float4 blendedColor = (color1 + color2 + color3 + color4) * 0.25f;
    // return perMaterialData.color * blendedColor + input.color;

    // TODO: Randome color on meshlet.
    // return float4(1, 0, 0, 1);
    uint hash = PCGHash(input.meshletID);

    float r = float((hash & 0xFF0000u) >> 16) / 255.0;
    float g = float((hash & 0x00FF00u) >> 8)  / 255.0;
    float b = float((hash & 0x0000FFu))       / 255.0;

    return float4(r, g, b, 1.0);
}