#version 430 core

// Constants
#define PI 3.14159265359
#define INV_PI 0.31830988618
#define EPSILON 1e-4
#define MAX_FLOAT 3.402823466e38
#define MAX_RAY_DEPTH 8
#define MAX_LIGHTS 16
#define RR_THRESHOLD 0.1

// Material types
#define MATERIAL_DIFFUSE 0
#define MATERIAL_METAL 1
#define MATERIAL_DIELECTRIC 2
#define MATERIAL_EMISSIVE 3

// Light types
#define LIGHT_DIRECTIONAL 0
#define LIGHT_POINT 1
#define LIGHT_SPOT 2

// Structures
struct Material {
    vec3 albedo;
    float metallic;
    vec3 emission;
    float roughness;
    int type;
    float ior;
    vec2 padding;
};

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

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


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;
layout(binding = 2, rgba8) uniform readonly image2D g_albedo;

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

// Scene data
layout(std430, binding = 0) readonly buffer MaterialBuffer {
    Material materials[];
};

layout(std430, binding = 1) readonly buffer LightBuffer {
    Light lights[];
};

// Uniforms
uniform uint u_frame_count;
uniform uint u_samples_per_pixel;
uniform uint u_max_depth;
uniform uint u_light_count;
uniform vec3 u_camera_position;
uniform mat4 u_inv_view_projection;
uniform bool u_enable_accumulation;
uniform bool u_use_bvh;

// ============================================================================
// Utility Functions
// ============================================================================

/**
 * @brief Saturate float value to [0, 1]
 */
float saturate(float x) {
    return clamp(x, 0.0, 1.0);
}

/**
 * @brief Saturate vec3 value to [0, 1]
 */
vec3 saturate(vec3 x) {
    return clamp(x, vec3(0.0), vec3(1.0));
}

/**
 * @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;
}

// ============================================================================
// Random Number Generation
// ============================================================================

/**
 * @brief PCG hash function for random number generation
 */
uint pcg_hash(uint seed) {
    uint state = seed * 747796405u + 2891336453u;
    uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
    return (word >> 22u) ^ word;
}

/**
 * @brief Generate random float in [0, 1)
 */
float random_float(inout uint seed) {
    seed = pcg_hash(seed);
    return float(seed) / 4294967296.0;
}

/**
 * @brief Generate random vec2
 */
vec2 random_vec2(inout uint seed) {
    return vec2(random_float(seed), random_float(seed));
}

/**
 * @brief Generate random vec3
 */
vec3 random_vec3(inout uint seed) {
    return vec3(random_float(seed), random_float(seed), random_float(seed));
}

/**
 * @brief Generate random vector in unit sphere
 */
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;
    }
}

/**
 * @brief Generate random unit vector
 */
vec3 random_unit_vector(inout uint seed) {
    return normalize(random_in_unit_sphere(seed));
}

// ============================================================================
// Sampling Functions
// ============================================================================

/**
 * @brief Cosine-weighted hemisphere sampling
 * @return Sampled direction in local space (z-up)
 */
vec3 cosine_weighted_hemisphere(inout uint seed) {
    vec2 r = random_vec2(seed);
    float phi = 2.0 * PI * r.x;
    float cos_theta = sqrt(r.y);
    float sin_theta = sqrt(1.0 - r.y);
    
    return vec3(cos(phi) * sin_theta, sin(phi) * sin_theta, cos_theta);
}

/**
 * @brief Build orthonormal basis from normal
 */
void build_onb(vec3 normal, out vec3 tangent, out vec3 bitangent) {
    vec3 up = abs(normal.z) < 0.999 ? vec3(0.0, 0.0, 1.0) : vec3(1.0, 0.0, 0.0);
    tangent = normalize(cross(up, normal));
    bitangent = cross(normal, tangent);
}

/**
 * @brief Transform direction from local to world space
 */
vec3 local_to_world(vec3 local_dir, vec3 normal) {
    vec3 tangent, bitangent;
    build_onb(normal, tangent, bitangent);
    return tangent * local_dir.x + bitangent * local_dir.y + normal * local_dir.z;
}

/**
 * @brief Sample GGX distribution for microfacet normal
 */
vec3 sample_ggx(vec3 normal, float roughness, inout uint seed) {
    vec2 r = random_vec2(seed);
    float a = roughness * roughness;
    float a2 = a * a;
    
    float phi = 2.0 * PI * r.x;
    float cos_theta = sqrt((1.0 - r.y) / (1.0 + (a2 - 1.0) * r.y));
    float sin_theta = sqrt(1.0 - cos_theta * cos_theta);
    
    vec3 local_h = vec3(cos(phi) * sin_theta, sin(phi) * sin_theta, cos_theta);
    return local_to_world(local_h, normal);
}

// ============================================================================
// BRDF Functions
// ============================================================================

/**
 * @brief Schlick's approximation for Fresnel reflectance
 */
vec3 fresnel_schlick(float cos_theta, vec3 f0) {
    return f0 + (1.0 - f0) * pow(1.0 - cos_theta, 5.0);
}

/**
 * @brief Fresnel reflectance for dielectrics
 */
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);
}

/**
 * @brief GGX normal distribution function
 */
float distribution_ggx(vec3 N, vec3 H, float roughness) {
    float a = roughness * roughness;
    float a2 = a * a;
    float NdotH = max(dot(N, H), 0.0);
    float NdotH2 = NdotH * NdotH;
    
    float nom = a2;
    float denom = (NdotH2 * (a2 - 1.0) + 1.0);
    denom = PI * denom * denom;
    
    return nom / max(denom, EPSILON);
}

/**
 * @brief Smith's geometry function for GGX
 */
float geometry_smith(vec3 N, vec3 V, vec3 L, float roughness) {
    float NdotV = max(dot(N, V), 0.0);
    float NdotL = max(dot(N, L), 0.0);
    float r = roughness + 1.0;
    float k = (r * r) / 8.0;
    
    float ggx1 = NdotV / (NdotV * (1.0 - k) + k);
    float ggx2 = NdotL / (NdotL * (1.0 - k) + k);
    
    return ggx1 * ggx2;
}

// ============================================================================
// Material Scattering
// ============================================================================

/**
 * @brief Scatter ray for diffuse material
 */
ScatterResult scatter_diffuse(Ray ray_in, HitInfo hit, Material mat, inout uint seed) {
    ScatterResult result;
    result.scattered = true;
    result.attenuation = mat.albedo;
    
    // Cosine-weighted hemisphere sampling
    vec3 local_dir = cosine_weighted_hemisphere(seed);
    vec3 scatter_direction = local_to_world(local_dir, hit.normal);
    
    // Prevent degenerate scatter direction
    if (near_zero(scatter_direction)) {
        scatter_direction = hit.normal;
    }
    
    result.scattered_ray.origin = hit.position + hit.normal * EPSILON;
    result.scattered_ray.direction = normalize(scatter_direction);
    result.pdf = max(dot(hit.normal, result.scattered_ray.direction), 0.0) * INV_PI;
    
    return result;
}

/**
 * @brief Scatter ray for metal material
 */
ScatterResult scatter_metal(Ray ray_in, HitInfo hit, Material mat, inout uint seed) {
    ScatterResult result;
    
    vec3 reflected = reflect_vector(normalize(ray_in.direction), hit.normal);
    
    // Add roughness perturbation
    vec3 fuzz = mat.roughness * random_in_unit_sphere(seed);
    vec3 scatter_direction = reflected + fuzz;
    
    result.scattered = dot(scatter_direction, hit.normal) > 0.0;
    
    if (result.scattered) {
        result.scattered_ray.origin = hit.position + hit.normal * EPSILON;
        result.scattered_ray.direction = normalize(scatter_direction);
        
        vec3 V = -normalize(ray_in.direction);
        vec3 H = normalize(V + result.scattered_ray.direction);
        vec3 F0 = mat.albedo;
        
        result.attenuation = fresnel_schlick(max(dot(H, V), 0.0), F0);
        result.pdf = 1.0; // Delta distribution approximation
    } else {
        result.attenuation = vec3(0.0);
        result.pdf = 0.0;
    }
    
    return result;
}

/**
 * @brief Scatter ray for dielectric material (glass)
 */
ScatterResult scatter_dielectric(Ray ray_in, HitInfo hit, Material mat, inout uint seed) {
    ScatterResult result;
    result.scattered = true;
    result.attenuation = vec3(1.0);
    
    float refraction_ratio = dot(ray_in.direction, hit.normal) < 0.0 ? 
        (1.0 / mat.ior) : mat.ior;
    
    vec3 unit_direction = normalize(ray_in.direction);
    float cos_theta = min(dot(-unit_direction, hit.normal), 1.0);
    float sin_theta = sqrt(1.0 - cos_theta * cos_theta);
    
    bool cannot_refract = refraction_ratio * sin_theta > 1.0;
    float reflectance = fresnel_dielectric(cos_theta, refraction_ratio);
    
    vec3 direction;
    if (cannot_refract || reflectance > random_float(seed)) {
        direction = reflect_vector(unit_direction, hit.normal);
    } else {
        direction = refract_vector(unit_direction, hit.normal, refraction_ratio);
    }
    
    result.scattered_ray.origin = hit.position + direction * EPSILON;
    result.scattered_ray.direction = normalize(direction);
    result.pdf = 1.0; // Delta distribution
    
    return result;
}

/**
 * @brief Scatter ray based on material type
 */
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);
    } else if (mat.type == MATERIAL_METAL) {
        return scatter_metal(ray_in, hit, mat, seed);
    } else if (mat.type == MATERIAL_DIELECTRIC) {
        return scatter_dielectric(ray_in, hit, mat, seed);
    } else {
        // Emissive material doesn't scatter
        ScatterResult result;
        result.scattered = false;
        result.attenuation = vec3(0.0);
        result.pdf = 0.0;
        return result;
    }
}

// ============================================================================
// Scene Intersection (G-Buffer based)
// ============================================================================

/**
 * @brief Trace ray against G-Buffer (single bounce only)
 * @note This is a simplified version - full path tracing needs scene geometry
 */
HitInfo trace_ray_gbuffer(Ray ray, ivec2 pixel_coords) {
    HitInfo hit;
    
    // Read from G-Buffer at current pixel
    vec4 position_data = imageLoad(g_position, pixel_coords);
    vec4 normal_data = imageLoad(g_normal, pixel_coords);
    vec4 albedo_data = imageLoad(g_albedo, pixel_coords);
    
    if (position_data.w > 0.5) {
        hit.hit = true;
        hit.position = position_data.xyz;
        hit.normal = normalize(normal_data.xyz);
        hit.material_id = uint(albedo_data.a * 255.0 + 0.5);
        hit.t = length(hit.position - ray.origin);
    } else {
        hit.hit = false;
        hit.t = MAX_FLOAT;
    }
    
    return hit;
}

// ============================================================================
// Path Tracing Core
// ============================================================================


/**
 * @brief Sample direct lighting from light sources
 */
vec3 sample_direct_lighting(vec3 position, vec3 normal, Material mat, inout uint seed) {
    if (u_light_count == 0u) return vec3(0.0);
    
    vec3 direct_light = vec3(0.0);
    
    // Sample one random light (could be improved with MIS)
    uint light_idx = uint(random_float(seed) * float(u_light_count)) % u_light_count;
    Light light = lights[light_idx];
    
    vec3 light_dir;
    float light_distance;
    float pdf_light = 1.0 / float(u_light_count);
    
    if (light.type == LIGHT_POINT) {
        vec3 to_light = light.position - position;
        light_distance = length(to_light);
        light_dir = to_light / light_distance;
        
        if (light_distance > light.range) return vec3(0.0);
        
        float NdotL = max(dot(normal, light_dir), 0.0);
        if (NdotL > 0.0) {
            float attenuation = 1.0 / max(light_distance * light_distance, 0.01);
            vec3 radiance = light.color * light.intensity * attenuation;
            
            // Simple BRDF evaluation (diffuse)
            direct_light = mat.albedo * INV_PI * radiance * NdotL / pdf_light;
        }
    } else if (light.type == LIGHT_DIRECTIONAL) {
        light_dir = normalize(-light.direction);
        float NdotL = max(dot(normal, light_dir), 0.0);
        
        if (NdotL > 0.0) {
            vec3 radiance = light.color * light.intensity;
            direct_light = mat.albedo * INV_PI * radiance * NdotL / pdf_light;
        }
    }
    
    return direct_light;
}

/**
 * @brief Trace path and accumulate radiance
 */
vec3 trace_path(Ray initial_ray, ivec2 pixel_coords, inout uint seed) {
    vec3 radiance = vec3(0.0);
    vec3 throughput = vec3(1.0);
    Ray current_ray = initial_ray;
    
    uint mat_count = uint(materials.length());
    
    for (uint depth = 0u; depth < u_max_depth; depth++) {
        // Trace ray (only first bounce uses G-Buffer)
        HitInfo hit;
        if (depth == 0u) {
            hit = trace_ray_gbuffer(current_ray, pixel_coords);
        } else {
            // For subsequent bounces, we can't trace without full scene geometry
            // This is a limitation of G-Buffer based approach
            // In a full path tracer, you'd trace against the actual scene here
            hit.hit = false;
        }
        
        if (!hit.hit) {
            // Hit sky/background
            vec3 sky_color = vec3(0.1, 0.1, 0.15);
            radiance += throughput * sky_color;
            break;
        }
        
        // Get material
        Material mat;
        if (hit.material_id < mat_count) {
            mat = materials[hit.material_id];
        } else {
            // Fallback material
            mat.albedo = vec3(0.5);
            mat.metallic = 0.0;
            mat.roughness = 0.5;
            mat.emission = vec3(0.0);
            mat.type = MATERIAL_DIFFUSE;
            mat.ior = 1.5;
        }
        
        // Add emission
        radiance += throughput * mat.emission;
        
        // Sample direct lighting (only for diffuse surfaces)
        if (mat.type == MATERIAL_DIFFUSE && depth == 0u) {
            radiance += throughput * sample_direct_lighting(hit.position, hit.normal, mat, seed);
        }
        
        // Scatter ray
        ScatterResult scatter = scatter_ray(current_ray, hit, mat, seed);
        
        if (!scatter.scattered || scatter.pdf < EPSILON) {
            break;
        }
        
		// Update throughput
        throughput *= scatter.attenuation;
        
        // Russian roulette path termination
        if (depth > 3u) {
            float rr_probability = max(throughput.r, max(throughput.g, throughput.b));
            if (rr_probability < RR_THRESHOLD || random_float(seed) > rr_probability) {
                break;
            }
            throughput /= rr_probability;
        }
        
        // Continue with scattered ray
        current_ray = scatter.scattered_ray;
        
        // Safety check for throughput
        if (all(lessThan(throughput, vec3(EPSILON)))) {
            break;
        }
    }
    
    return radiance;
}

/**
 * @brief Enhanced direct lighting with G-Buffer
 */
vec3 render_direct_lighting(ivec2 pixel_coords, inout uint seed) {
    // Read G-Buffer
    vec4 position_data = imageLoad(g_position, pixel_coords);
    vec4 normal_data = imageLoad(g_normal, pixel_coords);
    vec4 albedo_data = imageLoad(g_albedo, pixel_coords);
    
    if (position_data.w < 0.5) {
        return vec3(0.1, 0.1, 0.15); // Sky
    }
    
    vec3 position = position_data.xyz;
    vec3 normal = normalize(normal_data.xyz);
    uint material_id = uint(albedo_data.a * 255.0 + 0.5);
    
    // Get material
    Material mat;
    uint mat_count = uint(materials.length());
    
    if (material_id < mat_count) {
        mat = materials[material_id];
    } else {
        // Fallback: use G-Buffer albedo
        mat.albedo = albedo_data.rgb;
        mat.metallic = 0.0;
        mat.roughness = 0.5;
        mat.emission = vec3(0.0);
        mat.type = MATERIAL_DIFFUSE;
        mat.ior = 1.5;
    }
    
    vec3 color = vec3(0.0);
    
    // Emission
    color += mat.emission;
    
    // View direction
    vec3 view_dir = normalize(u_camera_position - position);
    
    // Direct lighting from all lights
    for (uint i = 0u; i < u_light_count; i++) {
        Light light = lights[i];
        vec3 light_dir;
        float light_distance;
        float attenuation = 1.0;
        
        if (light.type == LIGHT_POINT) {
            vec3 to_light = light.position - position;
            light_distance = length(to_light);
            light_dir = to_light / light_distance;
            attenuation = 1.0 / max(light_distance * light_distance, 0.01);
            
            if (light_distance > light.range) continue;
        } else if (light.type == LIGHT_DIRECTIONAL) {
            light_dir = normalize(-light.direction);
            light_distance = MAX_FLOAT;
        } else {
            continue;
        }
        
        float NdotL = max(dot(normal, light_dir), 0.0);
        
        if (NdotL > 0.0) {
            vec3 H = normalize(view_dir + light_dir);
            float NdotV = max(dot(normal, view_dir), 0.0);
            float NdotH = max(dot(normal, H), 0.0);
            float HdotV = max(dot(H, view_dir), 0.0);
            
            // Cook-Torrance BRDF
            vec3 F0 = mix(vec3(0.04), mat.albedo, mat.metallic);
            vec3 F = fresnel_schlick(HdotV, F0);
            float D = distribution_ggx(normal, H, max(mat.roughness, 0.04));
            float G = geometry_smith(normal, view_dir, light_dir, mat.roughness);
            
            vec3 numerator = D * G * F;
            float denominator = 4.0 * NdotV * NdotL + EPSILON;
            vec3 specular = numerator / denominator;
            
            vec3 kS = F;
            vec3 kD = (vec3(1.0) - kS) * (1.0 - mat.metallic);
            
            vec3 radiance = light.color * light.intensity * attenuation;
            color += (kD * mat.albedo * INV_PI + specular) * radiance * NdotL;
        }
    }
    
    // Ambient occlusion approximation
    color += mat.albedo * 0.03;
    
    return color;
}

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;
    }
    
    // Initialize random seed
    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 = render_direct_lighting(pixel_coords, seed);
    
    // Clamp
    color = clamp(color, vec3(0.0), vec3(10.0));
    
    // Accumulation
    if (u_enable_accumulation && u_frame_count > 0u) {
        vec3 accumulated = imageLoad(accumulation_image, pixel_coords).rgb;
        float weight = 1.0 / float(u_frame_count + 1u);
        color = mix(accumulated, color, weight);
    }
    
    imageStore(accumulation_image, pixel_coords, vec4(color, 1.0));
    imageStore(output_image, pixel_coords, vec4(color, 1.0));
}