#version 430 core

#include "../include/common.glsl"
#include "../include/structs.glsl"
#include "../include/math.glsl"
#include "../include/rng.glsl"
#include "../include/sobol.glsl"
#include "../include/sampling.glsl"
#include "../include/tonemap.glsl"

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

layout(binding = 0, rgba32f) uniform readonly image2D g_position;
layout(binding = 1, rg32f) uniform readonly image2D g_normal;
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;

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

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[];
};

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;

uniform uint u_sr_enabled;
uniform uint u_sr_scaling; // pixel ratio, e.g. 4 → 4× fewer pixels
uniform uint u_sr_block; // sqrt(scaling), block side length in pixels
uniform uint u_sr_jitter; // frame index within one jitter cycle (0 .. scaling-1)
uniform uint u_sr_full_width;
uniform uint u_sr_full_height;

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 "../include/material.glsl"
#include "../include/bvh.glsl"
#include "../include/lighting.glsl"

struct SobolState {
    uint sample_index;
    uint dimension;
    uint scramble;
};

SobolState init_sobol(uint pixel_index, uint frame, uint sample_idx) {
    SobolState state;
    state.sample_index = sample_idx + frame * 1024u + pixel_index + 1u;
    state.dimension = 0u;
    state.scramble = pcg_hash(pixel_index + frame * 668265263u);
    return state;
}

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

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 ct = sqrt((1.0 - u1) / ((a2 - 1.0) * u1 + 1.0));
    ct = clamp(ct, 0.0, 1.0);
    float st = sqrt(max(0.0, 1.0 - ct * ct));
    float ph = 2.0 * PI * u2;
    vec3 Ht = vec3(st * cos(ph), st * sin(ph), ct);
    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);
    return mat3(T, B, N) * Ht;
}

ScatterResult scatter_diffuse_sobol(Ray ray_in, HitInfo hit, Material mat, inout SobolState state) {
    ScatterResult r;
    r.scattered = true;
    r.attenuation = mat.albedo;
    float rs = sqrt(sobol_next(state));
    float ph = 2.0 * PI * sobol_next(state);
    float x = rs * cos(ph), y = rs * sin(ph);
    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);
    vec3 dir = mat3(T, B, hit.normal) * 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;
}

ScatterResult scatter_metal_sobol(Ray ray_in, HitInfo hit, Material mat, inout SobolState state) {
    ScatterResult r;
    vec3 V = normalize(-ray_in.direction);
    float roughness = clamp(mat.roughness, 0.04, 1.0);
    vec3 H = sobol_ggx_half_vector(roughness, hit.normal, state);
    if (dot(H, hit.normal) < 0.0) H = -H;
    vec3 L = reflect(-V, H);
    if (dot(hit.normal, L) <= 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 + hit.normal * EPSILON;
    r.scattered_ray.direction = L;
    return r;
}

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 ud = normalize(ray_in.direction);
    float ct = dot(-ud, hit.normal);
    float st = sqrt(max(0.0, 1.0 - ct * ct));
    bool ent = ct > 0.0;
    float eta = ent ? (1.0 / mat.ior) : mat.ior;
    vec3 N = ent ? hit.normal : -hit.normal;
    float st_t = eta * st;
    float f = fresnel_dielectric(ct, mat.ior);
    vec3 dir;
    if (st_t >= 1.0 || sobol_next(state) < f)
        dir = reflect_vector(ud, N);
    else
        dir = refract_vector(ud, N, eta);
    r.scattered_ray.origin = hit.position + dir * EPSILON;
    r.scattered_ray.direction = dir;
    return r;
}

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

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 pn = u_inv_view_projection * vec4(ndc, 0.0, 1.0);
    vec4 pf = u_inv_view_projection * vec4(ndc, 1.0, 1.0);
    Ray r;
    r.origin = pn.xyz / pn.w;
    r.direction = normalize(pf.xyz / pf.w - r.origin);
    return r;
}

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

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

    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(max(throughput.r, 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;
}

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

    ivec2 pixel_coords;
    ivec2 image_size;

    if (u_sr_enabled != 0u) {
        uint jx = u_sr_jitter % u_sr_block;
        uint jy = u_sr_jitter / u_sr_block;
        pixel_coords = rt_coord * int(u_sr_block) + ivec2(int(jx), int(jy));
        image_size = ivec2(int(u_sr_full_width), int(u_sr_full_height));
    } else {
        pixel_coords = rt_coord;
        image_size = output_size;
    }

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

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

    // ── Super resolution: per‑cycle running average inside the RT shader ──
    if (u_sr_enabled != 0u) {
        vec4 old = imageLoad(accumulation_image, pixel_coords);
        uint cyc = u_frame_count / u_sr_scaling;
        float w = 1.0 / float(cyc + 1u);
        vec3 acc = mix(old.rgb, color, w);
        imageStore(accumulation_image, pixel_coords, vec4(acc, 1.0));
        return;
    }

    // ── Standard non‑SR accumulation ─────────────────────────────────────
    vec3 acc_color = color;
    if (u_enable_accumulation && u_frame_count > 0u) {
        vec3 old = imageLoad(accumulation_image, rt_coord).rgb;
        float w = 1.0 / float(u_frame_count + 1u);
        acc_color = mix(old, color, w);
    }
    vec3 out_color = aces_tonemap(acc_color);
    imageStore(accumulation_image, rt_coord, vec4(acc_color, 1.0));
    imageStore(output_image, rt_coord, vec4(out_color, 1.0));
}