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

// New: material params + material id
layout(binding = 5, rgba32f) uniform readonly image2D g_material;
layout(binding = 6, r32ui)   uniform readonly uimage2D g_material_id;

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

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

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

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 vec3 u_camera_position;
uniform mat4 u_inv_view_projection;
uniform bool u_enable_accumulation;
uniform bool u_use_bvh;
uniform uint u_bvh_node_count;

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

    // Interpolate normal/uv
    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;

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

    // Small fixed stack
    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 {
            // Interior: push children
            uint left = left_first;
            uint right = left_first + 1u;

            // Depth-first; no sorting (simple)
            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;
}

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

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 = min(dot(-unit_dir, hit.normal), 1.0);
    float sin_theta = sqrt(max(0.0, 1.0 - cos_theta * cos_theta));

    float refraction_ratio = dot(unit_dir, hit.normal) < 0.0 ? (1.0 / mat.ior) : mat.ior;
    bool cannot_refract = refraction_ratio * sin_theta > 1.0;
    float reflect_prob = fresnel_dielectric(cos_theta, refraction_ratio);

    vec3 dir;
    if (cannot_refract || random_float(seed) < reflect_prob) {
        dir = reflect_vector(unit_dir, hit.normal);
    } else {
        dir = refract_vector(unit_dir, hit.normal, refraction_ratio);
    }

    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(HitInfo hit, Material mat, inout uint seed) {
    if (u_light_count == 0u) return vec3(0.0);

    // sample one light
    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);

    // shadow ray
    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; // diffuse direct only (simple)
    return brdf * radiance * n_dot_l / 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;
    return m;
}

vec3 environment_color(vec3 dir) {
    // simple dark sky
    return vec3(0.1, 0.1, 0.15);
}

vec3 trace_path(Ray ray, inout uint seed) {
    vec3 radiance = vec3(0.0);
    vec3 throughput = vec3(1.0);

    for (uint depth = 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);

        // emission
        radiance += throughput * mat.emission;

        // direct light (only for diffuse to keep simple)
        if (mat.type == MATERIAL_DIFFUSE) {
            radiance += throughput * eval_direct_lighting(hit, mat, seed);
        }

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

        throughput *= sc.attenuation;

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

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

    // Multi-sample
    uint spp = max(u_samples_per_pixel, 1u);
    for (uint s = 0u; s < spp; ++s) {
        Ray cam_ray = generate_camera_ray(pixel_coords, image_size, seed);
        color += trace_path(cam_ray, seed);
    }
    color /= float(spp);

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

    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);
        color = mix(accumulated, color, w);
    }

    imageStore(accumulation_image, pixel_coords, vec4(color, 1.0));
    imageStore(output_image, pixel_coords, vec4(color, 1.0));
}