#version 430 core

// Include shared modules
#include "../include/common.glsl"
#include "../include/structs.glsl"
#include "../include/math.glsl"
#include "../include/rng.glsl"
#include "../include/sobol.glsl"
#include "../include/sampling.glsl"

// Workgroup size
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, rg32f)  uniform readonly image2D g_normal;  // Octahedral encoded
layout(binding = 5, rgba32f) uniform readonly image2D g_material;
layout(binding = 6, r32ui)   uniform readonly uimage2D g_material_id;
layout(binding = 2, rgba32f) uniform readonly image2D g_texcoord;
layout(binding = 7, rgba32f) uniform readonly image2D g_tangent;

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

// SSBO bindings
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 { TriangleCompactGpu bvh_tris[]; };
layout(std430, binding = 4) readonly buffer AttrBuffer { TriangleAttrGpu bvh_attrs[]; };

// Uniforms
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;

// Texture arrays
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;

// Include material, BVH, and lighting modules (needs uniform declarations above)
#include "../include/material.glsl"
#include "../include/bvh.glsl"
#include "../include/lighting.glsl"

// Sobol sampling state
struct SobolState {
    uint sample_index;  // Which sample (0, 1, 2, ...)
    uint dimension;     // Current dimension being used
    uint scramble;      // Seed for Owen scrambling
};

// Initialize Sobol state
SobolState init_sobol(uint pixel_index, uint frame, uint sample_idx) {
    SobolState state;
    // Sample index combines frame and sample number for temporal variation
    // Add +1 to avoid degenerate index 0 (Sobol at index 0 produces all zeros before scrambling)
    // Add pixel_index offset to ensure spatial variation within same frame
    state.sample_index = sample_idx + frame * 1024u + pixel_index + 1u;
    state.dimension = 0u;
    // Use pixel index for per-pixel Owen scrambling
    state.scramble = pcg_hash(pixel_index + frame * 668265263u);
    return state;
}

// Get next Sobol float in [0, 1)
float sobol_next(inout SobolState state) {
    float value;
    if (state.dimension < 16u) {
        value = sobol_get(state.sample_index, state.dimension, state.scramble);
    } else {
        uint rng_state = pcg_hash(state.scramble + state.dimension * 2654435761u);
        value = float(rng_state) / 4294967296.0;
    }
    state.dimension++;
    return value;
}

// Sobol-based GGX half vector sampling
vec3 sobol_ggx_half_vector(float roughness, vec3 N, inout SobolState state) {
    float a = roughness * roughness;
    float a2 = a * a;

    float u1 = clamp(sobol_next(state), 0.001, 0.999);
    float u2 = sobol_next(state);

    float cos_theta = sqrt((1.0 - u1) / ((a2 - 1.0) * u1 + 1.0));
    cos_theta = clamp(cos_theta, 0.0, 1.0);
    float sin_theta = sqrt(max(0.0, 1.0 - cos_theta * cos_theta));
    float phi = 2.0 * PI * u2;

    vec3 H_tangent = vec3(sin_theta * cos(phi), sin_theta * sin(phi), cos_theta);
    vec3 up = (abs(N.y) < 0.999) ? vec3(0.0, 1.0, 0.0) : vec3(1.0, 0.0, 0.0);
    vec3 T = normalize(cross(up, N));
    vec3 B = cross(N, T);
    mat3 onb = mat3(T, B, N);
    return onb * H_tangent;
}

// Sobol-based diffuse scattering with cosine-weighted importance sampling
ScatterResult scatter_diffuse_sobol(Ray ray_in, HitInfo hit, Material mat, inout SobolState state) {
    ScatterResult r;
    r.scattered = true;
    r.attenuation = mat.albedo;

    // Cosine-weighted importance sampling for Lambertian BRDF
    float r_sqrt = sqrt(sobol_next(state));
    float phi = 2.0 * PI * sobol_next(state);
    float x = r_sqrt * cos(phi);
    float y = r_sqrt * sin(phi);
    float z = sqrt(max(0.0, 1.0 - x * x - y * y));
    
    vec3 up = (abs(hit.normal.y) < 0.999) ? vec3(0.0, 1.0, 0.0) : vec3(1.0, 0.0, 0.0);
    vec3 T = normalize(cross(up, hit.normal));
    vec3 B = cross(hit.normal, T);
    mat3 onb = mat3(T, B, hit.normal);
    vec3 dir = onb * vec3(x, y, z);

    if (near_zero(dir)) dir = hit.normal;

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

// Sobol-based metal scattering (GGX)
ScatterResult scatter_metal_sobol(Ray ray_in, HitInfo hit, Material mat, inout SobolState state) {
    ScatterResult r;

    vec3 V = normalize(-ray_in.direction);
    vec3 N = hit.normal;
    float roughness = clamp(mat.roughness, 0.04, 1.0);

    vec3 H = sobol_ggx_half_vector(roughness, N, state);
    if (dot(H, N) < 0.0) H = -H;

    vec3 L = reflect(-V, H);

    float NdotL = dot(N, L);
    if (NdotL <= 0.0) {
        r.scattered = false;
        r.attenuation = vec3(0.0);
        return r;
    }

    float HdotV = max(dot(H, V), 0.001);
    vec3 F = fresnel_schlick(HdotV, mat.albedo);

    r.attenuation = clamp(F, vec3(0.0), vec3(1.0));
    r.scattered = true;
    r.scattered_ray.origin = hit.position + N * EPSILON;
    r.scattered_ray.direction = L;
    return r;
}

// Sobol-based dielectric scattering
ScatterResult scatter_dielectric_sobol(Ray ray_in, HitInfo hit, Material mat, inout SobolState state) {
    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));
    
    bool entering = cos_theta > 0.0;
    float eta = entering ? (1.0 / mat.ior) : mat.ior;
    vec3 normal = entering ? hit.normal : -hit.normal;
    
    float sin_theta_t = eta * sin_theta;
    bool total_internal_reflection = sin_theta_t >= 1.0;
    
    float f = fresnel_dielectric(cos_theta, mat.ior);
    
    vec3 dir;
    if (total_internal_reflection || sobol_next(state) < f) {
        dir = reflect_vector(unit_dir, normal);
    } else {
        dir = refract_vector(unit_dir, normal, eta);
    }

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

// Sobol-based scatter dispatcher
ScatterResult scatter_ray_sobol(Ray ray_in, HitInfo hit, Material mat, inout SobolState state) {
    if (mat.type == MATERIAL_DIFFUSE) return scatter_diffuse_sobol(ray_in, hit, mat, state);
    if (mat.type == MATERIAL_METAL) return scatter_metal_sobol(ray_in, hit, mat, state);
    if (mat.type == MATERIAL_DIELECTRIC) return scatter_dielectric_sobol(ray_in, hit, mat, state);

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

// Generate camera ray (center pixel, no jitter)
Ray generate_camera_ray(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;
}

// Path tracing with G-Buffer acceleration for primary ray (Sobol sampling)
vec3 trace_path_sobol(ivec2 pixel_coords, ivec2 image_size, inout SobolState sobol) {
    Ray ray = generate_camera_ray(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);
        
        if (hit0.material_type >= 0) {
            mat0.type = hit0.material_type;
        }
        
        apply_material_textures(mat0, hit0.normal, hit0.texcoord, hit0.tangent);

        radiance += throughput * mat0.emission;

        ScatterResult sc0 = scatter_ray_sobol(ray, hit0, mat0, sobol);
        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_material_textures(mat, hit.normal, hit.texcoord, hit.tangent);

        radiance += throughput * mat.emission;

        ScatterResult sc = scatter_ray_sobol(ray, hit, mat, sobol);
        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 || sobol_next(sobol) > 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 pixel_index = uint(pixel_coords.x) + uint(pixel_coords.y) * uint(image_size.x);
    uint seed = init_seed(pixel_coords, image_size, u_frame_count);

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

    for (uint s = 0u; s < spp; ++s) {
        SobolState sobol = init_sobol(pixel_index, u_frame_count, s);
        color += trace_path_sobol(pixel_coords, image_size, sobol);
    }
    color /= float(spp);

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

    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);
    }
    
    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));
}