#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

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

// Material structure - 使用 vec4 确保对齐
struct Material {
    vec4 albedo_metallic;      // xyz = albedo, w = metallic
    vec4 emission_roughness;   // xyz = emission, w = roughness
    int type;
    float ior;
    float padding1;
    float padding2;
};

// Light structure - 使用 vec4 确保对齐
struct Light {
    vec4 position_type;        // xyz = position, w = type
    vec4 direction_intensity;  // xyz = direction, w = intensity
    vec4 color_range;          // xyz = color, w = range
    vec4 spot_angles;          // xy = spot angles, zw = padding
};

// 辅助函数：从 Material 获取属性
vec3 get_material_albedo(Material mat) {
    return mat.albedo_metallic.xyz;
}

float get_material_metallic(Material mat) {
    return mat.albedo_metallic.w;
}

vec3 get_material_emission(Material mat) {
    return mat.emission_roughness.xyz;
}

float get_material_roughness(Material mat) {
    return mat.emission_roughness.w;
}

// 辅助函数：从 Light 获取属性
vec3 get_light_position(Light light) {
    return light.position_type.xyz;
}

int get_light_type(Light light) {
    return int(light.position_type.w);
}

vec3 get_light_direction(Light light) {
    return light.direction_intensity.xyz;
}

float get_light_intensity(Light light) {
    return light.direction_intensity.w;
}

vec3 get_light_color(Light light) {
    return light.color_range.xyz;
}

float get_light_range(Light light) {
    return light.color_range.w;
}

struct Ray {
    vec3 origin;
    vec3 direction;
};

struct HitInfo {
    bool hit;
    float t;
    vec3 position;
    vec3 normal;
    vec2 texcoord;
    uint material_id;
};

// Utility functions
float saturate(float x) {
    return clamp(x, 0.0, 1.0);
}

vec3 saturate(vec3 x) {
    return clamp(x, vec3(0.0), vec3(1.0));
}

// Random number generation (PCG Hash)
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;
}

vec2 random_vec2(inout uint seed) {
    return vec2(random_float(seed), random_float(seed));
}

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

// Random direction in hemisphere
vec3 random_hemisphere_direction(vec3 normal, inout uint seed) {
    float z = random_float(seed);
    float r = sqrt(max(0.0, 1.0 - z * z));
    float phi = 2.0 * PI * random_float(seed);
    
    vec3 dir = vec3(r * cos(phi), r * sin(phi), z);
    
    // Create coordinate system
    vec3 up = abs(normal.z) < 0.999 ? vec3(0.0, 0.0, 1.0) : vec3(1.0, 0.0, 0.0);
    vec3 tangent = normalize(cross(up, normal));
    vec3 bitangent = cross(normal, tangent);
    
    return normalize(tangent * dir.x + bitangent * dir.y + normal * dir.z);
}

// Cosine-weighted hemisphere sampling
// vec3 cosine_weighted_hemisphere(vec3 normal, inout uint seed) {
//     vec2 r = random_vec2(seed);
//     float r1 = 2.0 * PI * r.x;
//     float r2 = r.y;
//     float r2s = sqrt(r2);
//
//     vec3 up = abs(normal.z) < 0.999 ? vec3(0.0, 0.0, 1.0) : vec3(1.0, 0.0, 0.0);
//     vec3 tangent = normalize(cross(up, normal));
//     vec3 bitangent = cross(normal, tangent);
//
//     vec3 dir = tangent * cos(r1) * r2s + bitangent * sin(r1) * r2s + normal * sqrt(1.0 - r2);
//     return normalize(dir);
// }

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

// GGX distribution
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);
}

// Smith's geometry function
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;
}

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;
layout(binding = 3, r32ui) uniform readonly uimage2D g_material_id;

// Output
layout(binding = 4, rgba32f) uniform image2D output_image;
layout(binding = 5, 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;

// Evaluate direct lighting
vec3 evaluate_direct_lighting(vec3 position, vec3 normal, vec3 view_dir, 
                              Material material, inout uint seed) {
    vec3 direct_light = vec3(0.0);
    
    // 获取材质属性
    vec3 mat_albedo = get_material_albedo(material);
    float mat_metallic = get_material_metallic(material);
    float mat_roughness = get_material_roughness(material);
    
    for (uint i = 0u; i < u_light_count; i++) {
        Light light = lights[i];
        
        // 获取光源属性
        vec3 light_pos = get_light_position(light);
        int light_type = get_light_type(light);
        vec3 light_dir_raw = get_light_direction(light);
        float light_intensity = get_light_intensity(light);
        vec3 light_color = get_light_color(light);
        float light_range = get_light_range(light);
        
        vec3 light_dir;
        float light_distance;
        float attenuation = 1.0;
        
        if (light_type == LIGHT_POINT) {
            vec3 to_light = light_pos - 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_dir_raw);
            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);
            
            // PBR lighting
            vec3 F0 = mix(vec3(0.04), mat_albedo, mat_metallic);
            vec3 F = fresnel_schlick(max(dot(H, view_dir), 0.0), 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;
            direct_light += (kD * mat_albedo * INV_PI + specular) * radiance * NdotL;
        }
    }
    
    return direct_light;
}

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;
    }
    
    // ============ DEBUG MODE 6: 显示材质Metallic ============
    /*
    {
        uint mat_id = imageLoad(g_material_id, pixel_coords).r;
        if (imageLoad(g_position, pixel_coords).w > 0.5 && mat_id < uint(materials.length())) {
            Material mat = materials[mat_id];
            float metallic_vis = get_material_metallic(mat);
            imageStore(output_image, pixel_coords, vec4(metallic_vis, metallic_vis, metallic_vis, 1.0));
        } else {
            imageStore(output_image, pixel_coords, vec4(0.0, 0.0, 0.0, 1.0));
        }
        return;
    }
    */
    
    // 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);
    uint material_id = imageLoad(g_material_id, pixel_coords).r;
    
    // Background
    if (position_data.w < 0.5) {
        vec3 background = vec3(0.1, 0.1, 0.15);
        imageStore(output_image, pixel_coords, vec4(background, 1.0));
        imageStore(accumulation_image, pixel_coords, vec4(background, 1.0));
        return;
    }
    
    vec3 position = position_data.xyz;
    vec3 normal = normalize(normal_data.xyz);
    vec3 albedo = albedo_data.rgb;
    
    // Initialize random seed
    uint seed = uint(pixel_coords.x) + uint(pixel_coords.y) * uint(image_size.x) + u_frame_count * 719393u;
    
    vec3 color = vec3(0.0);
    
    // Get material from buffer
    Material material;
    uint mat_count = uint(materials.length());
    
    if (material_id < mat_count) {
        material = materials[material_id];
    } else {
        // Fallback - 创建默认材质
        material.albedo_metallic = vec4(albedo, 0.0);
        material.emission_roughness = vec4(0.0, 0.0, 0.0, 0.5);
        material.type = MATERIAL_DIFFUSE;
        material.ior = 1.5;
    }
    
    // Add emission
    color += get_material_emission(material);
    
    // Direct lighting
    vec3 view_dir = normalize(u_camera_position - position);
    
    if (u_light_count > 0u) {
        color += evaluate_direct_lighting(position, normal, view_dir, material, seed);
    }
    
    // Ambient lighting
    color += get_material_albedo(material) * 0.05;
    
    // 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));
}