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

layout(set = 0, binding = 0) uniform sampler samplers[];
layout(set = 0, binding = 1) uniform texture2D textures[];
layout(set = 0, binding = 2, rgba32f) uniform image2D images[];

layout(buffer_reference) readonly buffer draw_params {
    uint output_image_index;
    uint width;
    uint height;
    uint pad;
};

layout(push_constant) uniform buffers {
    layout(offset = 16) draw_params draw;
} regs;

vec2 sphere_intersection ( in vec3 orig, in vec3 dir, float radius ) {
    float b = dot(orig, dir);
    float c = dot(orig, orig) - radius * radius;
    float h = b * b - c;
    if (h < 0.f) return vec2(-1.f); // no intersection
    h = sqrt(h);
    return vec2(-b - h, -b + h);
}

vec3 rayleigh_scattering = vec3(5.802f, 13.558f, 33.1f);
vec3 rayleigh_absorption = vec3(0.f);

vec3 mie_scattering = vec3(3.996f);
vec3 mie_absorption = vec3(4.40f);

void main() {
    ivec2 coords = ivec2(gl_GlobalInvocationID.xy);

    const uint raymarching_steps = 10U;

    const float ground_radius = 6360.f;
    const float top_radius = 6420.f;
    const float atmosphere_radius = top_radius - ground_radius;

    vec2 uvs = vec2(coords) / vec2(regs.draw.width, regs.draw.height);

    float cos_mu = uvs.y * 2.f - 1.f;
    float mu = acos(cos_mu);
    float height = atmosphere_radius *  uvs.x;

    vec2 pos = vec2(0.f, ground_radius + height);
    vec2 dir = normalize(vec2(cos_mu, sin(mu)));

    // find ground hit
    float t = 0.f;
    vec2 ground_intersection = sphere_intersection(vec3(pos, 0.f), vec3(dir, 0.f), ground_radius);
    t = max(max(ground_intersection.x, 0.f), ground_intersection.y);
    t = ground_intersection.x;
    // vec4 val = vec4(vec3(0.f), 1.f);

    vec2 top_intersection = sphere_intersection(vec3(pos, 0.f), vec3(dir, 0.f), top_radius);
    // t = max(max(top_intersection.x, 0.f), top_intersection.y);

    // find atmosphere hit
    if (ground_intersection.y < 0.f) {
        vec2 top_intersection = sphere_intersection(vec3(pos, 0.f), vec3(dir, 0.f), top_radius);

        t = top_intersection.y;
        // val.z = 1.f;
    } else {
        // val.y = 1.f;
    }

    // val.x = t;
    // imageStore(images[nonuniformEXT(regs.draw.output_image_index)], coords, val);
    // return;
    imageStore(images[nonuniformEXT(regs.draw.output_image_index)], coords, vec4(t, 0.f, 0.f, 1.f));
    return;

    vec3 transmittance = vec3(0.f);

    // raymarching
    float step = t / raymarching_steps;
    for (uint step_index = 0U; step_index < raymarching_steps; step_index++) {
        float r_d = exp(-height / 8.f);
        vec3 r_e = rayleigh_scattering * r_d;
        float m_d = exp(-height / 1.2f);
        vec3 m_e = mie_scattering * m_d + mie_absorption;

        transmittance += (r_e + m_e) * step;
    }

    vec4 color = vec4(exp(-transmittance), 1.f);
    color.xyz = vec3(0.f);
    color.x = t;

    imageStore(images[nonuniformEXT(regs.draw.output_image_index)], coords, color);
}