#version 430 core

#define PI 3.14159265359
#define INV_PI 0.31830988618
#define EPSILON 1e-4
#define MAX_FLOAT 3.402823466e38
#define RR_THRESHOLD 0.1

#define MATERIAL_DIFFUSE 0
#define MATERIAL_METAL 1
#define MATERIAL_DIELECTRIC 2
#define MATERIAL_EMISSIVE 3

#define LIGHT_DIRECTIONAL 0
#define LIGHT_POINT 1
#define LIGHT_SPOT 2

layout(local_size_x = 16, local_size_y = 16) in;

// G-Buffer inputs
layout(binding = 0, rgba32f) uniform readonly image2D g_position;
layout(binding = 1, rgba32f) uniform readonly image2D g_normal;

// Material params + material id (for primary hit fast-path)
layout(binding = 5, rgba32f) uniform readonly image2D g_material;
layout(binding = 6, r32ui)   uniform readonly uimage2D g_material_id;

// Texcoord from G-Buffer
layout(binding = 7, rgba32f) uniform readonly image2D g_texcoord;

// Tangent from G-Buffer
layout(binding = 8, rgba32f) uniform readonly image2D g_tangent;

// Output
layout(binding = 3, rgba32f) uniform image2D output_image;
layout(binding = 4, rgba32f) uniform image2D accumulation_image;

struct Material {
    vec3 albedo;
    vec3 emission;
    float metallic;
    float roughness;
    int type;
    float ior;
    float ao;        // ambient occlusion
    float padding1;
    uint texture_handles[6];
};

struct Light {
    vec3 position;
    int type;
    vec3 direction;
    float intensity;
    vec3 color;
    float range;
    vec2 spot_angles;
    vec2 padding;
};

struct Ray {
    vec3 origin;
    vec3 direction;
};

struct HitInfo {
    bool hit;
    float t;
    vec3 position;
    vec3 normal;
    vec2 texcoord;
    vec3 tangent;
    uint material_id;
    int material_type;  // material type from G-Buffer
};

struct ScatterResult {
    bool scattered;
    vec3 attenuation;
    Ray scattered_ray;
};

struct BVHNodeGpu {
    vec4 aabb_min_left_first; // xyz min, w = left_first (uint bits in float)
    vec4 aabb_max_count;      // xyz max, w = count (uint bits in float)
};

struct TriangleGpu {
    vec4 v0_material; // xyz v0, w material_id (uint bits in float)
    vec4 v1;
    vec4 v2;
    vec4 n0;
    vec4 n1;
    vec4 n2;
    vec4 uv0_uv1;     // xy uv0, zw uv1
    vec4 uv2;         // xy uv2
    vec4 t0;          // tangent at v0
    vec4 t1;          // tangent at v1
};

layout(std430, binding = 0) readonly buffer MaterialBuffer { Material materials[]; };
layout(std430, binding = 1) readonly buffer LightBuffer { Light lights[]; };
layout(std430, binding = 2) readonly buffer BVHNodeBuffer { BVHNodeGpu bvh_nodes[]; };
layout(std430, binding = 3) readonly buffer TriangleBuffer { TriangleGpu bvh_tris[]; };

uniform uint u_frame_count;
uniform uint u_samples_per_pixel;
uniform uint u_max_depth;
uniform uint u_light_count;
uniform mat4 u_inv_view_projection;
uniform bool u_enable_accumulation;
uniform bool u_use_bvh;
uniform uint u_bvh_node_count;
uniform bool u_enable_textures;

// Global texture arrays for bindless sampling (6 arrays for each texture type)
layout(binding = 10) uniform sampler2DArray u_texture_albedo_array;
layout(binding = 11) uniform sampler2DArray u_texture_normal_array;
layout(binding = 12) uniform sampler2DArray u_texture_metallic_array;
layout(binding = 13) uniform sampler2DArray u_texture_roughness_array;
layout(binding = 14) uniform sampler2DArray u_texture_ao_array;
layout(binding = 15) uniform sampler2DArray u_texture_emission_array;

// Helper function to sample texture from array by index
vec4 sample_texture_array(int slot, int index, vec2 uv) {
    if (index <= 0) return vec4(1.0);
    
    if (slot == 0) return texture(u_texture_albedo_array, vec3(uv, float(index - 1)));
    if (slot == 1) return texture(u_texture_normal_array, vec3(uv, float(index - 1)));
    if (slot == 2) return texture(u_texture_metallic_array, vec3(uv, float(index - 1)));
    if (slot == 3) return texture(u_texture_roughness_array, vec3(uv, float(index - 1)));
    if (slot == 4) return texture(u_texture_ao_array, vec3(uv, float(index - 1)));
    if (slot == 5) return texture(u_texture_emission_array, vec3(uv, float(index - 1)));
    
    return vec4(1.0);
}

// ============================================================================
// Utility
// ============================================================================

/**
 * @brief Check if vector is near zero
 */
bool near_zero(vec3 v) {
    return (abs(v.x) < EPSILON) && (abs(v.y) < EPSILON) && (abs(v.z) < EPSILON);
}

/**
 * @brief Reflect vector around normal
 */
vec3 reflect_vector(vec3 v, vec3 n) {
    return v - 2.0 * dot(v, n) * n;
}

/**
 * @brief Refract vector through surface
 */
vec3 refract_vector(vec3 uv, vec3 n, float etai_over_etat) {
    float cos_theta = min(dot(-uv, n), 1.0);
    vec3 r_out_perp = etai_over_etat * (uv + cos_theta * n);
    vec3 r_out_parallel = -sqrt(abs(1.0 - dot(r_out_perp, r_out_perp))) * n;
    return r_out_perp + r_out_parallel;
}

uint as_uint(float f) { return floatBitsToUint(f); }
float as_float(uint u) { return uintBitsToFloat(u); }

// ============================================================================
// RNG (PCG)
// ============================================================================

uint pcg_hash(uint seed) {
    uint state = seed * 747796405u + 2891336453u;
    uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
    return (word >> 22u) ^ word;
}

float random_float(inout uint seed) {
    seed = pcg_hash(seed);
    return float(seed) / 4294967296.0;
}

vec3 random_vec3(inout uint seed) {
    return vec3(random_float(seed), random_float(seed), random_float(seed));
}

vec3 random_in_unit_sphere(inout uint seed) {
    while (true) {
        vec3 p = 2.0 * random_vec3(seed) - vec3(1.0);
        if (dot(p, p) < 1.0) return p;
    }
}

vec3 random_unit_vector(inout uint seed) {
    return normalize(random_in_unit_sphere(seed));
}

// ============================================================================
// Camera ray
// ============================================================================

/**
 * @brief Generate primary ray in world space
 */
Ray generate_camera_ray(ivec2 pixel_coords, ivec2 image_size, inout uint seed) {
    vec2 jitter = vec2(random_float(seed), random_float(seed));
    vec2 uv = (vec2(pixel_coords) + jitter) / vec2(image_size);
    vec2 ndc = uv * 2.0 - 1.0;

    vec4 p_near = u_inv_view_projection * vec4(ndc, 0.0, 1.0);
    vec4 p_far  = u_inv_view_projection * vec4(ndc, 1.0, 1.0);
    vec3 near_ws = p_near.xyz / p_near.w;
    vec3 far_ws  = p_far.xyz / p_far.w;

    Ray r;
    r.origin = near_ws;
    r.direction = normalize(far_ws - near_ws);
    return r;
}

// ============================================================================
// Intersection
// ============================================================================

/**
 * @brief Ray-AABB intersection
 */
bool intersect_aabb(Ray ray, vec3 aabb_min, vec3 aabb_max, float t_max) {
    vec3 inv_d = 1.0 / ray.direction;
    vec3 t0 = (aabb_min - ray.origin) * inv_d;
    vec3 t1 = (aabb_max - ray.origin) * inv_d;

    vec3 tmin3 = min(t0, t1);
    vec3 tmax3 = max(t0, t1);

    float tmin = max(max(tmin3.x, tmin3.y), tmin3.z);
    float tmax2 = min(min(tmax3.x, tmax3.y), tmax3.z);

    return (tmax2 >= max(tmin, 0.0)) && (tmin <= t_max);
}

/**
 * @brief Moller-Trumbore triangle intersection
 */
bool intersect_triangle(Ray ray, TriangleGpu tri, inout HitInfo hit) {
    vec3 v0 = tri.v0_material.xyz;
    vec3 v1 = tri.v1.xyz;
    vec3 v2 = tri.v2.xyz;

    vec3 e1 = v1 - v0;
    vec3 e2 = v2 - v0;
    vec3 pvec = cross(ray.direction, e2);
    float det = dot(e1, pvec);

    if (abs(det) < EPSILON) return false;
    float inv_det = 1.0 / det;

    vec3 tvec = ray.origin - v0;
    float u = dot(tvec, pvec) * inv_det;
    if (u < 0.0 || u > 1.0) return false;

    vec3 qvec = cross(tvec, e1);
    float v = dot(ray.direction, qvec) * inv_det;
    if (v < 0.0 || u + v > 1.0) return false;

    float t = dot(e2, qvec) * inv_det;
    if (t < EPSILON || t >= hit.t) return false;

    float w = 1.0 - u - v;
    vec3 n0 = tri.n0.xyz;
    vec3 n1 = tri.n1.xyz;
    vec3 n2 = tri.n2.xyz;

    vec2 uv0 = tri.uv0_uv1.xy;
    vec2 uv1 = tri.uv0_uv1.zw;
    vec2 uv2 = tri.uv2.xy;

    // Interpolate tangents
    vec3 t0 = tri.t0.xyz;
    vec3 t1 = tri.t1.xyz;
    // Compute t2 from normal and t0 (t2 = cross(n, t0))
    vec3 t2 = normalize(cross(n0, t0)); // approximate third tangent

    hit.hit = true;
    hit.t = t;
    hit.position = ray.origin + t * ray.direction;
    hit.normal = normalize(n0 * w + n1 * u + n2 * v);
    hit.texcoord = uv0 * w + uv1 * u + uv2 * v;
    
    // Interpolate tangent using barycentric coordinates
    hit.tangent = normalize(t0 * w + t1 * u + t2 * v);
    
    hit.material_id = as_uint(tri.v0_material.w);
    return true;
}

/**
 * @brief BVH traversal (closest hit)
 */
HitInfo trace_ray_bvh(Ray ray) {
    HitInfo hit;
    hit.hit = false;
    hit.t = MAX_FLOAT;

    if (!u_use_bvh || u_bvh_node_count == 0u) {
        return hit;
    }

    uint stack[64];
    int sp = 0;
    stack[sp++] = 0u;

    while (sp > 0) {
        uint node_idx = stack[--sp];
        if (node_idx >= u_bvh_node_count) continue;

        BVHNodeGpu node = bvh_nodes[node_idx];
        vec3 bmin = node.aabb_min_left_first.xyz;
        vec3 bmax = node.aabb_max_count.xyz;
        uint left_first = as_uint(node.aabb_min_left_first.w);
        uint count = as_uint(node.aabb_max_count.w);

        if (!intersect_aabb(ray, bmin, bmax, hit.t)) continue;

        if (count > 0u) {
            for (uint i = 0u; i < count; ++i) {
                TriangleGpu tri = bvh_tris[left_first + i];
                intersect_triangle(ray, tri, hit);
            }
        } else {
            uint left = left_first;
            uint right = left_first + 1u;
            if (sp < 63) stack[sp++] = right;
            if (sp < 63) stack[sp++] = left;
        }
    }

    return hit;
}

/**
 * @brief Any-hit BVH for shadow ray
 */
bool trace_any_bvh(Ray ray, float t_max) {
    if (!u_use_bvh || u_bvh_node_count == 0u) return false;

    uint stack[64];
    int sp = 0;
    stack[sp++] = 0u;

    HitInfo hit;
    hit.hit = false;
    hit.t = t_max;

    while (sp > 0) {
        uint node_idx = stack[--sp];
        if (node_idx >= u_bvh_node_count) continue;

        BVHNodeGpu node = bvh_nodes[node_idx];
        vec3 bmin = node.aabb_min_left_first.xyz;
        vec3 bmax = node.aabb_max_count.xyz;
        uint left_first = as_uint(node.aabb_min_left_first.w);
        uint count = as_uint(node.aabb_max_count.w);

        if (!intersect_aabb(ray, bmin, bmax, hit.t)) continue;

        if (count > 0u) {
            for (uint i = 0u; i < count; ++i) {
                TriangleGpu tri = bvh_tris[left_first + i];
                if (intersect_triangle(ray, tri, hit)) return true;
            }
        } else {
            uint left = left_first;
            uint right = left_first + 1u;
            if (sp < 63) stack[sp++] = right;
            if (sp < 63) stack[sp++] = left;
        }
    }

    return false;
}

// ============================================================================
// Primary-ray fast path via G-Buffer
// ============================================================================

/**
 * @brief Read primary hit from G-Buffer if current pixel has geometry
 * @note Uses g_position.w as "valid" marker (your gbuffer writes 1.0 on hits, clear is 0).
 */
HitInfo trace_primary_gbuffer(Ray ray, ivec2 pixel_coords) {
    HitInfo hit;
    hit.hit = false;
    hit.t = MAX_FLOAT;
    hit.position = vec3(0.0);
    hit.normal = vec3(0.0, 1.0, 0.0);
    hit.texcoord = vec2(0.0);
    hit.tangent = vec3(0.0);
    hit.material_id = 0u;
    hit.material_type = 0;

    vec4 pos = imageLoad(g_position, pixel_coords);
    if (pos.w <= 0.5) {
        return hit;
    }

    vec3 p = pos.xyz;
    vec3 n = normalize(imageLoad(g_normal, pixel_coords).xyz);

    // integer material id
    uint mid = imageLoad(g_material_id, pixel_coords).r;
    
    // material type stored in g_material.w
    vec4 mat = imageLoad(g_material, pixel_coords);
    int mtype = int(mat.w);
    
    // Read texcoord from G-Buffer
    vec4 texcoord_tangent = imageLoad(g_texcoord, pixel_coords);
    vec2 texcoord = texcoord_tangent.xy;
    
    // Read tangent from G-Buffer
    vec4 tangent_data = imageLoad(g_tangent, pixel_coords);
    vec3 tangent = tangent_data.xyz;

    hit.hit = true;
    hit.position = p;
    hit.normal = n;
    hit.texcoord = texcoord;
    hit.tangent = tangent;
    hit.material_id = mid;
    hit.material_type = mtype;

    // For RR/any debug usage; path tracing uses this as starting point only.
    hit.t = length(p - ray.origin);

    return hit;
}

// ============================================================================
// Material + scattering
// ============================================================================

// Apply normal map in world space
vec3 apply_normal_map(vec3 normal, vec2 texcoord, vec3 tangent, uint normal_handle) {
    if (normal_handle == 0 || !u_enable_textures) return normal;
    
    vec3 T = normalize(tangent - normal * dot(tangent, normal));
    vec3 B = cross(normal, T);
    mat3 TBN = mat3(T, B, normal);
    
    vec3 map_n = sample_texture_array(1, int(normal_handle), texcoord).xyz * 2.0 - 1.0;
    return normalize(TBN * map_n);
}

// Apply material textures to get final PBR values
void apply_material_textures(inout Material mat, inout vec3 normal, vec2 texcoord, vec3 tangent) {
    if (!u_enable_textures) return;
    
    // Albedo texture (replace)
    if (mat.texture_handles[0] != 0) {
        mat.albedo = sample_texture_array(0, int(mat.texture_handles[0]), texcoord).rgb;
    }
    
    // Normal map
    if (mat.texture_handles[1] != 0) {
        normal = apply_normal_map(normal, texcoord, tangent, mat.texture_handles[1]);
    }
    
    // Metallic texture (replace)
    if (mat.texture_handles[2] != 0) {
        mat.metallic = sample_texture_array(2, int(mat.texture_handles[2]), texcoord).r;
    }
    
    // Roughness texture (replace)
    if (mat.texture_handles[3] != 0) {
        mat.roughness = sample_texture_array(3, int(mat.texture_handles[3]), texcoord).r;
    }
    
    // AO texture (store in material, apply during lighting)
    if (mat.texture_handles[4] != 0) {
        mat.ao = sample_texture_array(4, int(mat.texture_handles[4]), texcoord).r;
    }
    
    // Emission texture (replace)
    if (mat.texture_handles[5] != 0) {
        mat.emission = sample_texture_array(5, int(mat.texture_handles[5]), texcoord).rgb;
    }
}

vec3 fresnel_schlick(float cos_theta, vec3 f0) {
    return f0 + (1.0 - f0) * pow(1.0 - cos_theta, 5.0);
}

float fresnel_dielectric(float cos_theta, float ior) {
    float r0 = (1.0 - ior) / (1.0 + ior);
    r0 = r0 * r0;
    return r0 + (1.0 - r0) * pow(1.0 - cos_theta, 5.0);
}

ScatterResult scatter_diffuse(Ray ray_in, HitInfo hit, Material mat, inout uint seed) {
    ScatterResult r;
    r.scattered = true;
    r.attenuation = mat.albedo;

    vec3 dir = hit.normal + random_unit_vector(seed);
    if (near_zero(dir)) dir = hit.normal;

    r.scattered_ray.origin = hit.position + hit.normal * EPSILON;
    r.scattered_ray.direction = normalize(dir);
    return r;
}

ScatterResult scatter_metal(Ray ray_in, HitInfo hit, Material mat, inout uint seed) {
    ScatterResult r;

    vec3 reflected = reflect_vector(normalize(ray_in.direction), hit.normal);
    vec3 fuzz = mat.roughness * random_in_unit_sphere(seed);
    vec3 dir = reflected + fuzz;

    r.scattered = dot(dir, hit.normal) > 0.0;
    r.attenuation = mat.albedo;
    r.scattered_ray.origin = hit.position + hit.normal * EPSILON;
    r.scattered_ray.direction = normalize(dir);
    return r;
}

ScatterResult scatter_dielectric(Ray ray_in, HitInfo hit, Material mat, inout uint seed) {
    ScatterResult r;
    r.scattered = true;
    r.attenuation = vec3(1.0);

    vec3 unit_dir = normalize(ray_in.direction);
    float cos_theta = dot(-unit_dir, hit.normal);
    float sin_theta = sqrt(max(0.0, 1.0 - cos_theta * cos_theta));
    
    // Determine if ray is entering or exiting the material
    // If dot(dir, normal) < 0, ray is entering (from air into material)
    bool entering = cos_theta > 0.0;
    
    // eta: ratio of indices (etai/etat)
    // Entering: eta = 1.0/ior (air to material)
    // Exiting: eta = ior/1.0 (material to air)
    float eta = entering ? (1.0 / mat.ior) : mat.ior;
    
    // Use correct normal for refraction calculation
    // When exiting, we need to use -normal
    vec3 normal = entering ? hit.normal : -hit.normal;
    
    // Check for total internal reflection
    float sin_theta_t = eta * sin_theta;
    bool total_internal_reflection = sin_theta_t >= 1.0;
    
    // Fresnel reflectance (Schlick approximation)
    float f0 = pow((1.0 - mat.ior) / (1.0 + mat.ior), 2.0);
    float f = f0 + (1.0 - f0) * pow(1.0 - abs(cos_theta), 5.0);
    
    vec3 dir;
    if (total_internal_reflection || random_float(seed) < f) {
        // Reflect
        dir = reflect_vector(unit_dir, normal);
    } else {
        // Refract
        dir = refract_vector(unit_dir, normal, eta);
    }

    r.scattered_ray.origin = hit.position + dir * EPSILON;
    r.scattered_ray.direction = normalize(dir);
    return r;
}

ScatterResult scatter_ray(Ray ray_in, HitInfo hit, Material mat, inout uint seed) {
    if (mat.type == MATERIAL_DIFFUSE) return scatter_diffuse(ray_in, hit, mat, seed);
    if (mat.type == MATERIAL_METAL) return scatter_metal(ray_in, hit, mat, seed);
    if (mat.type == MATERIAL_DIELECTRIC) return scatter_dielectric(ray_in, hit, mat, seed);

    ScatterResult r;
    r.scattered = false;
    r.attenuation = vec3(0.0);
    return r;
}

// ============================================================================
// Direct lighting (with shadow ray)
// ============================================================================

vec3 eval_direct_lighting(inout HitInfo hit, Material mat, inout uint seed) {
    if (u_light_count == 0u) return vec3(0.0);

    uint light_idx = uint(random_float(seed) * float(u_light_count)) % u_light_count;
    Light light = lights[light_idx];

    vec3 L;
    float dist = MAX_FLOAT;
    vec3 radiance = vec3(0.0);

    if (light.type == LIGHT_POINT) {
        vec3 to_light = light.position - hit.position;
        dist = length(to_light);
        if (dist > light.range) return vec3(0.0);
        L = to_light / dist;

        float atten = 1.0 / max(dist * dist, 0.01);
        radiance = light.color * light.intensity * atten;
    } else if (light.type == LIGHT_DIRECTIONAL) {
        L = normalize(-light.direction);
        radiance = light.color * light.intensity;
    } else {
        return vec3(0.0);
    }

    float n_dot_l = max(dot(hit.normal, L), 0.0);
    if (n_dot_l <= 0.0) return vec3(0.0);

    Ray shadow_ray;
    shadow_ray.origin = hit.position + hit.normal * EPSILON;
    shadow_ray.direction = L;

    float t_max = (light.type == LIGHT_POINT) ? (dist - EPSILON) : MAX_FLOAT;
    if (trace_any_bvh(shadow_ray, t_max)) return vec3(0.0);

    float pdf_light = 1.0 / float(u_light_count);
    vec3 brdf = mat.albedo * INV_PI;
    // Apply AO to the final lighting
    return brdf * radiance * n_dot_l * mat.ao / max(pdf_light, EPSILON);
}

// ============================================================================
// Path tracing
// ============================================================================

Material fetch_material(uint material_id) {
    uint cnt = uint(materials.length());
    if (material_id < cnt) return materials[material_id];

    Material m;
    m.albedo = vec3(0.5);
    m.metallic = 0.0;
    m.emission = vec3(0.0);
    m.roughness = 0.5;
    m.type = MATERIAL_DIFFUSE;
    m.ior = 1.5;
    m.ao = 1.0;  // default: no AO
    return m;
}

vec3 environment_color(vec3 dir) {
    return vec3(0.1, 0.1, 0.15);
}

Ray generate_camera_ray_center(ivec2 pixel_coords, ivec2 image_size) {
    vec2 uv = (vec2(pixel_coords) + vec2(0.5)) / vec2(image_size);
    vec2 ndc = uv * 2.0 - 1.0;

    vec4 p_near = u_inv_view_projection * vec4(ndc, 0.0, 1.0);
    vec4 p_far  = u_inv_view_projection * vec4(ndc, 1.0, 1.0);
    vec3 near_ws = p_near.xyz / p_near.w;
    vec3 far_ws  = p_far.xyz / p_far.w;

    Ray r;
    r.origin = near_ws;
    r.direction = normalize(far_ws - near_ws);
    return r;
}

/**
 * @brief Trace path with primary-ray G-Buffer acceleration
 */
vec3 trace_path_primary_gbuffer(ivec2 pixel_coords, ivec2 image_size, inout uint seed) {
    Ray ray = generate_camera_ray_center(pixel_coords, image_size);

    vec3 radiance = vec3(0.0);
    vec3 throughput = vec3(1.0);

    // Depth 0: try G-Buffer hit first
    HitInfo hit0 = trace_primary_gbuffer(ray, pixel_coords);
    if (hit0.hit) {
        Material mat0 = fetch_material(hit0.material_id);
        
        // Override material type from G-Buffer if available
        if (hit0.material_type >= 0) {
            mat0.type = hit0.material_type;
        }
        
        // Apply PBR textures (use tangent from G-Buffer if available)
        apply_material_textures(mat0, hit0.normal, hit0.texcoord, hit0.tangent);

        radiance += throughput * mat0.emission;

        ScatterResult sc0 = scatter_ray(ray, hit0, mat0, seed);
        if (!sc0.scattered) return radiance;

        throughput *= sc0.attenuation;
        ray = sc0.scattered_ray;
    }

    // Subsequent bounces: BVH
    for (uint depth = (hit0.hit ? 1u : 0u); depth < u_max_depth; ++depth) {
        HitInfo hit = trace_ray_bvh(ray);
        if (!hit.hit) {
            radiance += throughput * environment_color(ray.direction);
            break;
        }

        Material mat = fetch_material(hit.material_id);
        
        // Apply PBR textures (use tangent from intersection)
        apply_material_textures(mat, hit.normal, hit.texcoord, hit.tangent);

        radiance += throughput * mat.emission;

        ScatterResult sc = scatter_ray(ray, hit, mat, seed);
        if (!sc.scattered) break;

        throughput *= sc.attenuation;

        if (depth > 3u) {
            float p = max(throughput.r, max(throughput.g, throughput.b));
            p = clamp(p, 0.0, 0.95);
            if (p < RR_THRESHOLD || random_float(seed) > p) break;
            throughput /= p;
        }

        ray = sc.scattered_ray;

        if (all(lessThan(throughput, vec3(EPSILON)))) break;
    }

    return radiance;
}

// ACES Filmic Tone Mapping
vec3 aces_tonemap(vec3 x) {
    float a = 2.51;
    float b = 0.03;
    float c = 2.43;
    float d = 0.59;
    float e = 0.14;
    return clamp((x * (a * x + b)) / (x * (c * x + d) + e), 0.0, 1.0);
}

void main() {
    ivec2 pixel_coords = ivec2(gl_GlobalInvocationID.xy);
    ivec2 image_size = imageSize(output_image);
    if (pixel_coords.x >= image_size.x || pixel_coords.y >= image_size.y) return;

    uint base_seed = uint(pixel_coords.x) + uint(pixel_coords.y) * uint(image_size.x);
    uint seed = base_seed + u_frame_count * 719393u;

    vec3 color = vec3(0.0);
    uint spp = max(u_samples_per_pixel, 1u);

    for (uint s = 0u; s < spp; ++s) {
        color += trace_path_primary_gbuffer(pixel_coords, image_size, seed);
    }
    color /= float(spp);

    color = clamp(color, vec3(0.0), vec3(100.0));

    // Store HDR color to accumulation buffer BEFORE tone mapping
    vec3 accumulation_color = color;
    
    if (u_enable_accumulation && u_frame_count > 0u) {
        vec3 accumulated = imageLoad(accumulation_image, pixel_coords).rgb;
        float w = 1.0 / float(u_frame_count + 1u);
        accumulation_color = mix(accumulated, color, w);
    }
    
    // Apply ACES tone mapping to output (not accumulation)
    vec3 output_color = aces_tonemap(accumulation_color);

    imageStore(accumulation_image, pixel_coords, vec4(accumulation_color, 1.0));
    imageStore(output_image, pixel_coords, vec4(output_color, 1.0));
}