layout(local_size_x = 8, local_size_y = 8, local_size_z = 1) in;

struct ray {
    vec3 origin;
    vec3 direction;
};

const uint LAMBERTIAN_MATERIAL  = 0;
const uint METAL_MATERIAL       = 1;
const uint DIELECTRIC_MATERIAL  = 2;

const vec3 light_position = vec3(10, -25, 0);

struct material {
    vec4 albedo;
    float fuzz;
    float ior;
    uint type;
    uint padding;
};

struct hit_info {
    float t;
    vec3 point;
    vec3 normal;
    material mat;
    bool front_face;
};

struct sphere {
    vec4 position;
    material mat;
    float radius;
    float[3] padding;
};

struct camera {
    vec4 position;

    vec4 forward;
    vec4 up;
    vec4 right;

    vec4 horizontal;
    vec4 vertical;
    vec4 first_pixel;

    float lens_radius;
    float[3] padding;
};

layout(set = 0, binding = 0) buffer Inputs {
    vec4 sky_color;
    vec4 ground_color;
    camera cam;

    vec2[1000] random_offset;
    vec2[1000] random_disk;
    vec4[1000] random_in_sphere;

    sphere[512] spheres;

    uint max_bounce;
    uint samples_per_pixel;
    uint width;
    uint height;

    uint sample_index;
    uint[3] padding;
};

layout(set = 0, binding = 1, rgba8) uniform writeonly image2D output_image;
layout(set = 0, binding = 2, rgba8) uniform readonly image2D accumulation_image;

vec3 at(ray r, float t) {
    return r.origin + r.direction * t;
}

vec3 to_hemisphere_vec(vec3 unit_vec, vec3 normal) {
    if (dot(unit_vec, normal) > 0.0) {
        return unit_vec;
    } else {
        return -unit_vec;
    }
}

float rand(vec2 co) {
    return fract(sin(dot(co.xy ,vec2(12.9898,78.233))) * 43758.5453);
}

bool hit_sphere(sphere s, ray r, float min_t, float max_t, out hit_info info) {
    vec3 origin_center = r.origin - s.position.xyz;

    float a = dot(r.direction, r.direction);
    float half_b = dot(r.direction, origin_center);
    float c = dot(origin_center, origin_center) - s.radius * s.radius;
    float discriminant = half_b * half_b - a * c;

    if (discriminant < 0) {
        return false;
    }

    float sqrt_discriminant = sqrt(discriminant);
    float root = (-half_b - sqrt_discriminant) / a;
    if (root < min_t || root > max_t) {
        root = (-half_b + sqrt_discriminant) / a;

        if (root < min_t || root > max_t) {
            return false;
        }
    }

    info.t = root;
    info.point = at(r, info.t);

    vec3 outward_normal = (info.point - s.position.xyz) / s.radius;
    info.front_face = dot(outward_normal, r.direction) < 0.0;
    info.normal = info.front_face ? outward_normal : -outward_normal;
    info.mat = s.mat;

    return true;
}

bool hit_spheres(ray r, float min_t, float max_t, out hit_info info) {
    float closest_t = max_t;
    hit_info temp_info;
    bool hit = false;
    for (uint sphere_index = 0; sphere_index < 512; sphere_index++) {
        if (spheres[sphere_index].radius > 0.0 && hit_sphere(spheres[sphere_index], r, min_t, closest_t, temp_info)) {
            info = temp_info;
            closest_t = temp_info.t;
            hit = true;
        }
    }

    return hit;
}

bool lambertian_scattering(material mat, ray incident, vec2 seed, hit_info info, out vec4 attenuation, out ray scattered) {
    vec3 unit_hemisphere_vec = to_hemisphere_vec(random_in_sphere[uint(rand(seed) * 999.0)].xyz, info.normal);

    attenuation = mat.albedo;
    scattered = ray(info.point, info.normal + unit_hemisphere_vec);

    return true;
}

bool metal_scattering(material mat, ray incident, vec2 seed, hit_info info, out vec4 attenuation, out ray scattered) {
    vec3 reflected = reflect(normalize(incident.direction), info.normal);
    vec3 unit_sphere_vec = random_in_sphere[uint(rand(seed) * 999.0)].xyz;

    attenuation = mat.albedo;
    scattered = ray(info.point, reflected + mat.fuzz * unit_sphere_vec); 

    return (dot(scattered.direction, info.normal) > 0.0);
}

float reflectance(float cosine, float eta) {
    // Use Schlick's approximation for reflectance.
    float r0 = (1 - eta) / (1 + eta);
    r0 = r0 * r0;
    return r0 + (1 - r0) * pow((1 - cosine), 5);
}

bool dielectric_scattering(material mat, ray incident, vec2 seed, hit_info info, out vec4 attenuation, out ray scattered) {
    attenuation = vec4(1.0);

    vec3 unit_dir = normalize(incident.direction);
    float cos_theta = min(dot(-unit_dir, info.normal), 1.0);
    float sin_theta = sqrt(1.0 - cos_theta*cos_theta);

    float eta = info.front_face ? 1.0 / mat.ior : mat.ior;
    vec3 direction;

    if (eta * sin_theta > 1.0 || reflectance(cos_theta, eta) > rand(seed)) {
        direction = reflect(unit_dir, info.normal);
    } else {
        direction = refract(unit_dir, info.normal, eta);
    }

    scattered = ray(info.point, direction); 

    return true;
}

vec4 ray_color(ray r, uint sample_index) {
    hit_info info;
    uint bounces = 0;

    vec4 final_color = vec4(1.0);
    ray current_ray = r;
    while (hit_spheres(current_ray, 0.0001, 1e15, info) && bounces < max_bounce) {
        vec2 seed = vec2(info.point.x * sample_index, info.point.y * bounces) * gl_GlobalInvocationID.xy;
        ray scattered;
        vec4 attenuation;
        if ((info.mat.type == LAMBERTIAN_MATERIAL && lambertian_scattering(info.mat, current_ray, seed, info, attenuation, scattered)) ||
        (info.mat.type == METAL_MATERIAL && metal_scattering(info.mat, current_ray, seed, info, attenuation, scattered)) ||
        (info.mat.type == DIELECTRIC_MATERIAL && dielectric_scattering(info.mat, current_ray, seed, info, attenuation, scattered))) {
            final_color *= attenuation;
            current_ray = scattered;
        } else {
            return vec4(0.0);
        }

        bounces++;
    }

    if (bounces == max_bounce) {
        return vec4(0.0);
    }

    return final_color * mix(ground_color, sky_color, (normalize(r.direction).y + 1.0) * 0.5);
}

void main() {
    uint x = gl_GlobalInvocationID.x;
    uint y = gl_GlobalInvocationID.y;

    vec2 uv = vec2(float(x) / (width - 1), 1.0 - (float(y) / (height - 1)));

    vec2 jittered_uvs = uv + vec2(random_disk[sample_index].x / (width - 1), random_disk[sample_index].y / (height - 1));
    vec2 lens_disk = random_disk[sample_index] * cam.lens_radius;
    vec3 offset = lens_disk.x * cam.right.xyz + lens_disk.y * cam.up.xyz;
    vec4 proj_plane_pos = cam.first_pixel + jittered_uvs.x * cam.horizontal + jittered_uvs.y * cam.vertical;

    ray r = ray(cam.position.xyz + offset, proj_plane_pos.xyz - cam.position.xyz - offset);

    vec4 out_color = ray_color(r, sample_index);

    float spp_scale = (1.0 / samples_per_pixel);
    vec4 acc_color = imageLoad(accumulation_image, ivec2(x, y));
    imageStore(output_image, ivec2(x, y), acc_color + out_color * spp_scale);
}