#version 450
#extension GL_GOOGLE_include_directive : require

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

#include "lights.glsl"
#include "disney_bsdf.glsl"

#define BIAS 0.0002
#define VARIANCE_MIN 0.00000001
#define SHADOW_CUTOFF 0.0001

layout(set = 0, binding = 0) uniform Locals {
    mat4 VP;
    mat4 P;
    uvec4 light_count;
    vec4 cam_pos;
};
layout(std430, set = 0, binding = 1) buffer readonly Materials { Material materials[]; };
layout(set = 0, binding = 2, rgba16f) uniform writeonly image2D LightingTexture;

layout(set = 1, binding = 0, rgba16f) uniform readonly image2D Albedo;
layout(set = 1, binding = 1, rgba16f) uniform readonly image2D Normal;
layout(set = 1, binding = 2, rgba16f) uniform readonly image2D WorldPos;
layout(set = 1, binding = 3, rgba16f) uniform readonly image2D MatParams;

// layout(set = 2, binding = 0) uniform PointLights { PointLight point_lights[128]; };
layout(set = 2, binding = 1) uniform AreaLights { AreaLight area_lights[128]; };
layout(set = 2, binding = 2) uniform SpotLights { SpotLight spot_lights[128]; };
layout(set = 2, binding = 3) uniform DirectionalLights { DirectionalLight dir_lights[128]; };

layout(set = 2, binding = 4) uniform sampler ShadowSampler;

// layout(set = 2, binding = 5) uniform textureCubeArray PointShadowMaps;
layout(set = 2, binding = 6) uniform texture2DArray AreaShadowMaps;
layout(set = 2, binding = 7) uniform texture2DArray SpotShadowMaps;
layout(set = 2, binding = 8) uniform texture2DArray DirShadowMaps;

// layout(set = 2, binding = 9) uniform textureCubeArray PointShadowMatrices { LightInfo PointTransforms[128]; };

layout(set = 2, binding = 10) uniform AreaShadowMatrices { LightInfo AreaTransforms[128]; };
layout(set = 2, binding = 11) uniform SpotShadowMatrices { LightInfo SpotTransforms[128]; };
layout(set = 2, binding = 12) uniform DirShadowMatrices { LightInfo DirTransforms[128]; };

// float fetch_point_shadow(uint light_id, float bias, vec3 worldToLight);
float fetch_area_shadow(uint light_id, float bias, vec4 ls_coords);
float fetch_spot_shadow(uint light_id, float bias, vec4 ls_coords);
float fetch_dir_shadow(uint light_id, float bias, vec4 ls_coords);

void main() {
    const ivec2 pixel = ivec2(gl_GlobalInvocationID.xy);
    const ivec2 image_size = imageSize(Albedo);
    if (pixel.x >= image_size.x || pixel.y >= image_size.y) {
        return;
    }

    const vec3 normal = imageLoad(Normal, pixel).xyz;
    const vec3 V = imageLoad(WorldPos, pixel).xyz - BIAS * normal;
    const vec4 albedo_mid = imageLoad(Albedo, pixel).xyzw;
    const vec4 mat_params = imageLoad(MatParams, pixel).xyzw;

    const vec3 albedo = albedo_mid.xyz;
    const uint MID = uint(albedo_mid.w);
    const uvec4 parameters = materials[MID].parameters;

    ShadingData params = extractParameters(albedo, vec3(1.0), materials[MID].specular.xyz, parameters);
    params.roughness = max(params.roughness, mat_params.x);
    params.metallic = max(params.metallic, mat_params.y);
    params.sheen = max(params.sheen, mat_params.z);

    const vec3 D = normalize(cam_pos.xyz - V);

    vec3 final = vec3(0.0);

    // const uint point_light_count = light_count.x;
    const uint spot_light_count = light_count.y;
    const uint area_light_count = light_count.z;
    const uint dir_light_count = light_count.w;

    // for (uint i = 0; i < point_light_count; i++)
    // {
    //     vec3 L = vec3(point_lights[i].position_energy) - V;
    //     const float distance2 = dot(L, L);
    //     L /= sqrt(distance2);
    //     const float NdotL = dot(normal, L);
    //     if (NdotL <= 0.0) {
    //         continue;
    //     }

    //     const float bias = max(0.05 * (1.0 - NdotL), 0.005);  
    //     const float shadow = fetch_point_shadow(i, bias, -L);

    //     final += shadow * vec3(point_lights[i].radiance) * NdotL * (1.0 / distance2);
    // }

    //    for (uint i = 0; i < spot_light_countspot_light_count; i++)
    for (uint i = 0; i < spot_light_count; i++)
    {
        const vec3 pos = spot_lights[i].position_cos_inner.xyz;
        const float cos_inner = spot_lights[i].position_cos_inner.w;
        const float cos_outer = spot_lights[i].radiance_cos_outer.w;
        const vec3 direction = spot_lights[i].direction_energy.xyz;

        vec3 L = pos - V.xyz;
        const float dist2 = dot(L, L);
        L = normalize(L);

        const float d = max(0.0, -dot(L, direction) - cos_outer) / (cos_inner - cos_outer);
        const float NdotL = dot(normal, L);
        const float LNdotL = min(1.0, d);
        if (NdotL <= 0.0 || LNdotL <= 0.0) {
            continue;
        }

        const vec4 ls_vertex = SpotTransforms[i].MP * vec4(V.xyz, 1.0);
        const float shadow = fetch_spot_shadow(i, BIAS, ls_vertex);
        if (shadow < SHADOW_CUTOFF) {
            continue;
        }
        const vec3 bsdf = evalLighting(params, normal, D, L);
        final += bsdf * shadow * spot_lights[i].radiance_cos_outer.xyz * NdotL * LNdotL * (1.0 / dist2);
    }

    for (uint i = 0; i < area_light_count; i++)
    {
        const vec3 pos = area_lights[i].position_energy.xyz;
        const vec3 LN = area_lights[i].normal_area.xyz;
        const float area = area_lights[i].normal_area.w;

        vec3 L = pos - V.xyz;
        const float dist2 = dot(L, L);
        L = normalize(L);

        const float NdotL = dot(normal, L);
        const float LNdotL = -dot(LN, L);
        if (NdotL <= 0.0 || LNdotL <= 0.0) {
            continue;
        }

        const vec4 ls_vertex = AreaTransforms[i].MP * vec4(V.xyz, 1.0);
        const float shadow = fetch_area_shadow(i, BIAS, ls_vertex);
        if (shadow < SHADOW_CUTOFF) {
            continue;
        }
        const vec3 bsdf = evalLighting(params, normal, D, L);
        final += bsdf * shadow * vec3(area_lights[i].radiance_x, area_lights[i].radiance_y, area_lights[i].radiance_z) * NdotL * LNdotL * (1.0 / dist2);
    }

    for (uint i = 0; i < dir_light_count; i++)
    {
        const vec3 L = -dir_lights[i].direction_energy.xyz;
        const float NdotL = dot(L, normal);
        if (NdotL < 0.0) {
            continue;
        }

        const vec4 ls_vertex = DirTransforms[i].MP * vec4(V.xyz, 1.0);
        const float shadow = fetch_dir_shadow(i, BIAS, ls_vertex);
        if (shadow < SHADOW_CUTOFF) {
            continue;
        }
        const vec3 bsdf = evalLighting(params, normal, D, L);
        final += bsdf * shadow * dir_lights[i].radiance.xyz * NdotL;
    }

    imageStore(LightingTexture, pixel, vec4(final, 1.0));
}

// float fetch_point_shadow(uint light_id, float bias, vec3 worldToLight) {
//     vec4 p = vec4(-worldToLight, light_id);
//     float depth = length(worldToLight);
//     float shadow_depth = texture(samplerCubeArray(PointShadowMaps, ShadowSampler), p).r * 1e3;
//     return (shadow_depth - bias) < depth ? 1.0 : 0.0;
// }

float linstep(float low, float high, float v){
    return clamp((v - low) / (high - low), 0.0, 1.0);
}

float fetch_area_shadow(uint light_id, float bias, vec4 ls_coords) {
    if (ls_coords.w <= 0.0) {
        return 0.0;
    }
    const vec2 texelSize = 1.0 / textureSize(sampler2DArray(AreaShadowMaps, ShadowSampler), 0).xy;
    const vec3 flip_correction = vec3(0.5, -0.5, 1);
    const vec3 coords = ls_coords.xyz / ls_coords.w * flip_correction;
    const vec2 s_uv = coords.xy + 0.5;
    float depth = coords.z - BIAS;
    const vec3 p = vec3(s_uv, light_id);

    depth = linearizeDepth(depth, AreaTransforms[light_id].PosRange.w);
    const vec2 moments = texture(sampler2DArray(AreaShadowMaps, ShadowSampler), p).xy;

    if (depth < moments.x) {
        return 1.0;
    }

    float variance = moments.y - (moments.x * moments.x);
    if (variance <= VARIANCE_MIN) {
        return 0.0;
    }

    const float d = depth - moments.x;
    const float p_max = variance / (variance + d * d);
    return p_max;
}

float fetch_spot_shadow(uint light_id, float bias, vec4 ls_coords) {
    if (ls_coords.w <= 0.0) {
        return 0.0;
    }
    const vec2 texelSize = 1.0 / textureSize(sampler2DArray(SpotShadowMaps, ShadowSampler), 0).xy;
    const vec3 flip_correction = vec3(0.5, -0.5, 1);
    const vec3 coords = ls_coords.xyz / ls_coords.w * flip_correction;
    const vec2 s_uv = coords.xy + 0.5;
    float depth = coords.z - BIAS;
    const vec3 p = vec3(s_uv, light_id);

    depth = linearizeDepth(depth, SpotTransforms[light_id].PosRange.w);
    const vec2 moments = texture(sampler2DArray(SpotShadowMaps, ShadowSampler), p).xy;

    if (depth < moments.x) {
        return 1.0;
    }

    float variance = moments.y - (moments.x * moments.x);
    if (variance <= VARIANCE_MIN) {
        return 0.0;
    }

    const float d = depth - moments.x;
    const float p_max = variance / (variance + d * d);
    return p_max;
}

float fetch_dir_shadow(uint light_id, float bias, vec4 ls_coords) {
    if (ls_coords.w <= 0.0) {
        return 1.0;
    }

    const vec2 texelSize = 1.0 / textureSize(sampler2DArray(DirShadowMaps, ShadowSampler), 0).xy;
    const vec3 flip_correction = vec3(0.5, -0.5, 1);
    const vec3 coords = ls_coords.xyz / ls_coords.w * flip_correction;
    const vec2 s_uv = coords.xy + 0.5;
    float depth = coords.z - BIAS;
    const vec3 p = vec3(s_uv, light_id);
    if (depth > 1.0) {
        return 0.0;
    }

    const vec2 moments = texture(sampler2DArray(DirShadowMaps, ShadowSampler), p).xy;
    if (depth < moments.x) {
        return 1.0;
    }

    float variance = moments.y - (moments.x * moments.x);
    if (variance <= VARIANCE_MIN) {
        return 0.0;
    }

    const float d = depth - moments.x;
    const float p_max = variance / (variance + d * d);
    return p_max;
}
