struct ray {
    vec3 origin;
    vec3 direction;
};

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

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

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

    uint samples_per_pixel;
    float viewport_width;
    float viewport_height;
    float proj_plane_distance;

    float[200] random_numbers;

    vec4[1000] random_in_sphere;

    sphere[5] spheres;

    uint max_bounce;
    float [3] padding;

    vec4[] 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 < 5; sphere_index++) {
        if (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;
    uint width = gl_NumWorkGroups.x;
    uint height = gl_NumWorkGroups.y;

    vec2 uv = vec2(float(x) / (width - 1), 1.0 - (float(y) / (height - 1)));
    vec3 first = camera_pos.xyz - vec3(viewport_width / 2.0, viewport_height / 2.0, proj_plane_distance);

    image[y * width + x] = vec4(0.0);

    for (uint sample_index = 0; sample_index < samples_per_pixel; sample_index++) {
        vec2 jittered_uvs = uv + vec2(random_numbers[sample_index * 2] / (width - 1), random_numbers[sample_index * 2 + 1] / (height - 1));
        vec3 proj_plane_pos = first + vec3(jittered_uvs.x * viewport_width, jittered_uvs.y * viewport_height, 0.0);

        image[y * width + x] += ray_color(ray(camera_pos.xyz, proj_plane_pos - camera_pos.xyz), sample_index);
    }

    float spp_scale = (1.0 / samples_per_pixel);
    image[y * width + x] = sqrt(image[y * width + x] * spp_scale);
}