#version 460 core

// hdr color
layout (location = 0) out vec3 h_color;

struct Light {
    // w contains the radius
    vec4 pos_radius;
    vec4 color;
};

layout(std430, binding = 0) buffer LightBuffer
{
    Light lights[];
};
layout (location = 0) uniform uint light_count;
layout (location = 1) uniform vec3 ambient_color;

// GBuffer
layout (location = 2) uniform sampler2D g_position;
layout (location = 3) uniform sampler2D g_albedo;
layout (location = 4) uniform sampler2D g_normal;
layout (location = 5) uniform sampler2D g_orm;

struct Surface {
    vec3 position;
    vec3 albedo;
    vec3 normal;
    float occlusion;
    float roughness;
    float metallic;
};

in VS_OUT {
    vec2 texcoord;
} fs_in;

const float PI = 3.1415926;

float GGXDistribution(float alpha, vec3 normal, vec3 H) {
    float a2 = alpha * alpha;
    float NdotH = dot(normal, H);
    float NdotH2 = NdotH * NdotH;
    float num = a2 * max(NdotH, 0.0); // both are length = 1, so this is heaviside
    float den = (NdotH2 * (a2 - 1.0) + 1);
    return num / (PI * den * den);
}

float SmithJoinMaskingShadowing(float alpha, vec3 normal, vec3 light_dir, vec3 view_dir, vec3 H) {
    float NdotL = abs(dot(normal, light_dir));
    float NdotV = abs(dot(normal, view_dir));
    float NdotL2 = NdotL * NdotL;
    float NdotV2 = NdotV * NdotV;
    float alpha2 = alpha * alpha;

    float num1 = 2 * NdotL * max(dot(H, light_dir), 0);
    float den1 = NdotL + sqrt(alpha2 + (1-alpha2)*NdotL2);
    float num2 = 2 * NdotV * max(dot(H, view_dir), 0);
    float den2 = NdotV + sqrt(alpha2 + (1-alpha2)*NdotV2);

    return (num1*num2)/(den1*den2);
}

vec3 FresnelSlick(vec3 f0, vec3 view_dir, vec3 H) {
    float VdotH = dot(view_dir, H);
    return f0 + (1.0 - f0) * pow(1 - abs(VdotH), 5);
}

vec3 Lerp(vec3 a, vec3 b, float t) {
    return a + (b - a) * t;
}

Surface ReadSurface() {
    Surface surf;
    surf.position = texture(g_position, fs_in.texcoord).rgb;
    surf.albedo = texture(g_albedo, fs_in.texcoord).rgb;
    surf.normal = texture(g_normal, fs_in.texcoord).rgb;
    vec3 orm = texture(g_orm, fs_in.texcoord).rgb;
    surf.occlusion = orm.r;
    surf.roughness = orm.g;
    surf.metallic = orm.b;
    return surf;
}

void main() {
    Surface surf = ReadSurface();
    // Eye location is always (0,0,0) because we calculate everything in view-space
    vec3 view_dir = normalize(-1 * surf.position);
    vec3 black = vec3(0, 0, 0);

    vec3 c_diff = Lerp(surf.albedo, black, surf.metallic);
    vec3 f0 = Lerp(vec3(0.04), surf.albedo, surf.metallic); 
    
    // We probably want a tiling based approach, but this is good enough for now
    vec3 Lo = vec3(0.0);
    for (uint i = 0; i < light_count; ++i) {
        vec3 light_pos = lights[i].pos_radius.xyz;
        float light_radius = lights[i].pos_radius.w;
        float light_dist = length(light_pos - surf.position);
        if (light_dist > light_radius)
            continue;
        float attenuation = 1.0 / (light_dist * light_dist);
        vec3 radiance = lights[i].color.rgb * attenuation;

        vec3 light_dir = normalize(-1 * light_pos);
        vec3 H = normalize(light_dir + view_dir);

        float NdotL = dot(surf.normal, light_dir);

        vec3 F = FresnelSlick(f0, view_dir, H);

        vec3 f_diffuse = (1 - F) * (1 / PI) * c_diff;
        vec3 f_specular = F * GGXDistribution(surf.roughness, surf.normal, H)
            * SmithJoinMaskingShadowing(surf.roughness, surf.normal, light_dir, view_dir, H)
            / (4 * abs(dot(view_dir, surf.normal)) * abs(NdotL) + 0.0001);

        vec3 material = f_diffuse + f_specular;
        Lo += material * radiance * NdotL;
    }

    vec3 ambient = ambient_color * surf.albedo * surf.occlusion;
    h_color = ambient + Lo;
}