#version 460
#extension GL_EXT_ray_tracing : enable
#extension GL_EXT_shader_image_load_formatted : enable
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : enable
#extension GL_EXT_descriptor_heap : enable
#extension GL_EXT_nonuniform_qualifier : enable

// Specialization constant set from descriptorHeap.bufferStartElement —
// same pattern as the Sponza example. The TLAS lives at descriptor_heap
// slot `bufferStart`, the per-frame output image at heap slot 0.
layout(constant_id = 0) const uint16_t bufferStart = 0us;
layout(descriptor_heap) uniform accelerationStructureEXT topLevelAS[];
layout(descriptor_heap) uniform writeonly image2D image[];

layout(location = 0) rayPayloadEXT vec3 hitValue;

void main() {
    uvec2 pixel      = gl_LaunchIDEXT.xy;
    uvec2 resolution = gl_LaunchSizeEXT.xy;
    vec2  uv         = (vec2(pixel) + 0.5) / vec2(resolution);
    vec2  ndc        = uv * 2.0 - 1.0;

    // Fixed camera framing the 3x3x3 grid (extent = (kGrid-1)*kSpacing
    // = 6): position ext*(1.1, 0.8, 1.6) looking at the grid centre —
    // the WebGPU example's initial free-camera pose.
    vec3 origin  = vec3(6.6, 4.8, 9.6);
    vec3 forward = normalize(-origin);
    vec3 right   = normalize(cross(forward, vec3(0.0, 1.0, 0.0)));
    vec3 up      = cross(right, forward);

    float aspect  = float(resolution.x) / float(resolution.y);
    float tanHalf = tan(radians(70.0) * 0.5);
    vec3 direction = normalize(
        right * (ndc.x * aspect * tanHalf) +
        up    * (-ndc.y * tanHalf) +
        forward);

    traceRayEXT(
        topLevelAS[bufferStart],
        gl_RayFlagsNoneEXT,
        0xff,
        0, 0, 0,
        origin,
        0.01,
        direction,
        100000.0,
        0);

    imageStore(image[0], ivec2(pixel), vec4(hitValue, 1.0));
}
