#extension GL_EXT_buffer_reference : require
#extension GL_EXT_nonuniform_qualifier : require

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

struct ray {
    vec3 origin;
    vec3 direction;
};

const float PI = 3.141592653589793;
const uint LAMBERTIAN_MATERIAL  = 0;
const uint METAL_MATERIAL       = 1;
const uint DIELECTRIC_MATERIAL  = 2;
const uint EMISSIVE_MATERIAL    = 3;

struct material {
    vec4 albedo;
    vec4 emissive;
    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 vertex {
    vec4 position;
};

struct triangle {
    vertex v1;
    vertex v2;
    vertex v3;
};

struct camera {
    vec4 position;

    vec4 forward;
    vec4 up;
    vec4 right;

    vec4 horizontal;
    vec4 vertical;
    vec4 first_pixel;

    float lens_radius;
    float[3] padding;
};

struct bvh_node {
    vec3 min;
    int next_id;
    vec3 max;
    int primitive_id;
};

layout(buffer_reference) readonly buffer scene_metadata {
    camera cam;

    uint max_bounce;
    uint width;
    uint height;
    uint sample_index;

    // uint spheres_count;
    uint triangles_count;
    bool enable_dof;
    bool debug_bvh;
    uint downscale_factor;
};

layout(buffer_reference) readonly buffer primitives_array {
    // sphere[] spheres;
    triangle triangles[];
};

layout(buffer_reference) readonly buffer nodes_array {
    bvh_node[] nodes;
};

layout(set = 0, binding = 2, rgba32f) uniform image2D images[];

layout(push_constant) uniform buffers {
    scene_metadata scene;
    primitives_array geometry;
    nodes_array bvh;
    uint output_image_index;
    uint accumulation_image_index;
    uint padding[2];
} bufs;

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

uint pcg_hash(inout uint seed)
{
    uint state = seed;
    seed *= 747796405u + 2891336453u;
    uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
    return (word >> 22u) ^ word;
}

float rand(inout uint seed)
{
    return float(pcg_hash(seed)) / 4294967296.0;
}

vec3 unit_vec(inout uint seed) {
    float z  = rand(seed) * 2.0 - 1.0;
    float a = 2.0 * PI * rand(seed);
    float r = sqrt(1.0 - z * z);

    return vec3(cos(a) * r, sin(a) * r, z);
}

vec3 diffuse_reflection_cosine_weighted(inout uint seed) {
    float r = rand(seed);
    float radius  = sqrt(r);
    float theta = 2.0 * PI * rand(seed);

    return vec3(cos(theta) * radius, sin(theta) * radius, sqrt(1 - r));
}

vec2 disk_vec(inout uint seed) {
    float z  = rand(seed) * 2.0 - 1.0;
    float a = 2.0 * PI * rand(seed);
    float r = sqrt(1.0 - z * z);

    return vec2(cos(a) * r, sin(a) * r);
}

bool hit_aabb(vec3 minimum, vec3 maximum, ray r, float min_t, float max_t) {
    const vec3 invdir = 1.0 / r.direction;
    const vec3 f = (maximum.xyz - r.origin.xyz) * invdir;
    const vec3 n = (minimum.xyz - r.origin.xyz) * invdir;

    const vec3 tmax = max(f, n);
    const vec3 tmin = min(f, n);

    const float t1 = min(tmax.x, min(tmax.y, tmax.z));
    const float t0 = max(max(tmin.x, max(tmin.y, tmin.z)), 0.0f);

    return t1 >= t0;
}

uint hit_aabbs(ray r, float min_t, float max_t) {
    uint hit_count = 0;
    int id = 0;

    while(id != -1) {
        if (bufs.bvh.nodes[id].primitive_id != -1) {
            triangle tri = bufs.geometry.triangles[bufs.bvh.nodes[id].primitive_id];
            vec3 minimum = min(tri.v1.position, min(tri.v2.position, tri.v3.position)).xyz;
            vec3 maximum = max(tri.v1.position, max(tri.v2.position, tri.v3.position)).xyz;

            if (hit_aabb(minimum, maximum, r, min_t, max_t)) {
                hit_count++;
            }
            id = bufs.bvh.nodes[id].next_id;
        } else if (hit_aabb(bufs.bvh.nodes[id].min, bufs.bvh.nodes[id].max, r, min_t, max_t)) {
                hit_count++;
                id++;
        } else {
            id = bufs.bvh.nodes[id].next_id;
        }
    }

    return hit_count;
}

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_triangle(triangle tri, ray r, float min_t, float max_t, out hit_info info) {
    vec3 v2v1 = tri.v2.position.xyz - tri.v1.position.xyz;
    vec3 v3v1 = tri.v3.position.xyz - tri.v1.position.xyz;

    vec3 originv1 = r.origin - tri.v1.position.xyz;
    vec3 normal = cross(v2v1, v3v1);
    vec3 q = cross(originv1, r.direction);
    float d = 1.0 / dot(r.direction, normal);
    float u = d * dot(-q, v3v1);
    float v = d * dot(q, v2v1);
    float t = d * dot(-normal, originv1);

    if (u < 0.0 || u > 1.0 || v < 0.0 || u + v > 1.0 || t < min_t || t > max_t) {
        return false;
    }

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

    // vec3 outward_normal = normal;
    // info.front_face = dot(outward_normal, r.direction) < 0.0;
    info.normal = normal;
    // info.mat = material(vec4(u, v, 0.0, 0.0), vec4(0.0), 1.0, 1.5, LAMBERTIAN_MATERIAL, 0);
    info.mat = material(vec4(0.55), vec4(0.0), 1.0, 1.5, LAMBERTIAN_MATERIAL, 0);

    return true;
}

bool hit_node(ray r, float min_t, float max_t, out hit_info info) {
    hit_info temp_info;
    float closest_t = max_t;
    bool hit = false;
    int id = 0;

    while(id != -1) {
        // Leafs
        if (bufs.bvh.nodes[id].primitive_id != -1) {
            // if (hit_sphere(spheres[bufs.bvh.nodes[id].primitive_id], r, min_t, closest_t, temp_info)) {
            //     info = temp_info;
            //     closest_t = temp_info.t;
            //     hit = true;
            // }
            if (hit_triangle(bufs.geometry.triangles[bufs.bvh.nodes[id].primitive_id], r, min_t, closest_t, temp_info)) {
                info = temp_info;
                closest_t = temp_info.t;
                hit = true;
            }
            id = bufs.bvh.nodes[id].next_id;
        } else { // Internal nodes
            if (hit_aabb(bufs.bvh.nodes[id].min, bufs.bvh.nodes[id].max, r, min_t, closest_t)) {
                id++;
            } else {
                id = bufs.bvh.nodes[id].next_id;
            }
        }
    }

    return hit;
}

bool lambertian_scattering(material mat, ray incident, inout uint seed, hit_info info, out vec4 attenuation, out ray scattered) {
    attenuation = mat.albedo;
    scattered = ray(info.point, normalize(info.normal + unit_vec(seed)));

    return true;
}

bool metal_scattering(material mat, ray incident, inout uint seed, hit_info info, out vec4 attenuation, out ray scattered) {
    vec3 reflected = reflect(normalize(incident.direction), info.normal);

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

    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, inout uint 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 seed) {
    hit_info info;
    uint bounces = 0;

    vec4 final_color = vec4(0.0);
    vec4 throughput = vec4(1.0);
    ray current_ray = r;
    while (hit_node(current_ray, 0.001, 1e15, info) && bounces < bufs.scene.max_bounce) {
        ray scattered;
        vec4 attenuation = vec4(0.0);
        // 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)) ||
        //     info.mat.type == EMISSIVE_MATERIAL) {
            lambertian_scattering(info.mat, current_ray, seed, info, attenuation, scattered);
            // final_color += info.mat.emissive * throughput;
            throughput *= attenuation;

            float NdotR = dot(info.normal, scattered.direction);
            float PDF = NdotR / PI;
            vec4 BRDF = attenuation / PI;
            throughput = (throughput * NdotR / PDF) * BRDF;
            current_ray = scattered;

            // Russian Roulette
            // As the throughput gets smaller, the ray is more likely to get terminated early.
            // Survivors have their value boosted to make up for fewer samples being in the average.
            // {
            //     float p = max(throughput.r, max(throughput.g, throughput.b));
            //     if (rand(seed) > p) {
            //         break;
            //     }

            //     // Add the energy we 'lose' by randomly terminating paths
            //     throughput *= 1.0 / p;
            // }
        // }
        bounces++;
    }

    return throughput;
    // return final_color;
}

void main() {
    uint x = gl_GlobalInvocationID.x * bufs.scene.downscale_factor;
    uint y = gl_GlobalInvocationID.y * bufs.scene.downscale_factor;

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

    uint seed = uint(x * uint(1973) + y * uint(9277) + bufs.scene.sample_index * uint(26699)) | uint(1);
    vec2 rand_disk = disk_vec(seed);
    vec2 jittered_uvs = uv + vec2(rand(seed) / (bufs.scene.width - 1), rand(seed) / (bufs.scene.height - 1));
    vec2 lens_disk = vec2(0.0);
    if (bufs.scene.enable_dof) {
        lens_disk = rand_disk * bufs.scene.cam.lens_radius;
    }

    vec3 offset = lens_disk.x * bufs.scene.cam.right.xyz + lens_disk.y * bufs.scene.cam.up.xyz;
    vec4 proj_plane_pos = bufs.scene.cam.first_pixel + jittered_uvs.x * bufs.scene.cam.horizontal + jittered_uvs.y * bufs.scene.cam.vertical;

    ray r = ray(bufs.scene.cam.position.xyz + offset, proj_plane_pos.xyz - bufs.scene.cam.position.xyz - offset);
    vec4 out_color = vec4(0.0);

    if (bufs.scene.debug_bvh) {
        out_color = hit_aabbs(r, 0.001, 1e15) * vec4(0.001, 0.0, 0.0, 1.0);
    } else {
        out_color = ray_color(r, seed);
    }

    float spp_scale = (1.0 / bufs.scene.sample_index);
    vec4 acc_color = imageLoad(images[nonuniformEXT(bufs.accumulation_image_index)], ivec2(x, y));
    out_color = vec4(mix(acc_color, out_color, spp_scale));

    for (uint i = 0; i < bufs.scene.downscale_factor; i++) {
        for (uint j = 0; j < bufs.scene.downscale_factor; j++) {
            imageStore(images[nonuniformEXT(bufs.output_image_index)], ivec2(x + i, y + j), out_color);
        }
    }
}