#version 450
#extension GL_EXT_nonuniform_qualifier : require

struct DirectionalLightStruct
{
  vec4 directionWorld;
  vec3 direction;
  uint isActive;
  uint cullingMask;
  vec4 diffuseColor;
  mat4 pvMatrix;
  uint shadowData;
};

struct AmbientLightStruct
{
  vec4 ambientColor;
  float strength;
  uint isActive;
  uint cullingMask;
};

layout(local_size_x = 16, local_size_y = 16) in;
layout(set = 3, binding = 0, rgba32f) uniform image2D positions;
layout(set = 3, binding = 1, rgba32f) uniform image2D normals;
layout(set = 3, binding = 2, rgba8) uniform image2D albedo;
layout(set = 3, binding = 3, rgba32ui) uniform uimage2D lightLayerData;
layout(set = 3, binding = 4, r8) uniform image2D ssaoBlurredImage;
layout(set = 3, binding = 5, rgba8) uniform image2D positionWorldSpace;
layout(set = 3, binding = 6, rgba8) uniform image2D targetImage;
layout(set = 3, binding = 7, rgba8) uniform image2D objectVFXImage;

layout (set = 4, binding = 0) uniform sampler2D shadowMaps[]; // for textures (global)

layout(set = 1, binding = 0) uniform LightCounts
{
  uint directionalLights;
  uint pointLights;
  uint spotLights;
  uint ambientLights;

} lightCounts;

layout(std430, set = 1, binding = 1) buffer DirectionalLightData
{
  DirectionalLightStruct dLightData[];
} DirLightData;

layout(std430, set = 1, binding = 4) buffer AmbientLightData
{
  AmbientLightStruct aLightData[];
} AmbLightData;

float LinStep (float val, float low, float high)
{
  return clamp ((val - low)/(high - low), 0.0f, 1.0f);
}

float CalcShadowValue (sampler2D shadowMap, vec4 worldSpaceFragPos, mat4 lightPV, vec3 worldNormal, vec3 lightDir)
{
  // clip space for fragment from light view space
  vec4 fragPosLightPOV = lightPV * worldSpaceFragPos;

  // Perform perspective division and convert to 0 to 1 range
  vec3 converted = (fragPosLightPOV.xyz / fragPosLightPOV.w) * vec3(0.5f) + vec3(0.5f);

  vec2 moments = texture(shadowMap, converted.xy).xy;

  if (converted.x < 0.0f || converted.x > 1.0f || converted.y < 0.0f || converted.y > 1.0f)
    return 1.0f;

  float returnVal = 0.0f;

  float worldNormalDotLight = dot (normalize (worldNormal), normalize(lightDir));
  
  if (worldNormalDotLight <= 0.0f)
    return 0.7f;

  // if (worldNormalDotLight <= 0.01f)
  //   return 0.7f;

  if (fragPosLightPOV.z > moments.x && fragPosLightPOV.w > 0.0f)
  {
    float p = step (fragPosLightPOV.z, moments.x);
    float variance = max (moments.y - (moments.x * moments.x), 0.00002f);

    float d = fragPosLightPOV.z - moments.x;
    float pMax = LinStep (variance / (variance + (d * d)), 0.9f, 1.0f);
   
    returnVal = min (max (p, pMax) + 0.7f, 1.0f);

    return returnVal;
    
  }
  else if (fragPosLightPOV.z > 1.0f)
  {
    return 0.0f;
  }

  // return min (worldNormalDotLight + 0.7f, 1.0f);
  return 1.0f;

}

void main()
{ 
  // convenient variables
  ivec2 globalThread = ivec2(gl_GlobalInvocationID);

  // Get the diffuse color of the pixel
  vec3 pixelDiffuse = imageLoad (albedo, globalThread).rgb;

  // Get position of fragment in world space
  vec4 positionWorld = vec4 (imageLoad (positionWorldSpace, globalThread).rgb, 1.0f);

  // Get position of fragment in view spacee
  vec3 positionView = imageLoad (positions, globalThread).rgb;

  // normal of fragment
  vec3 normalView = imageLoad(normals, globalThread).rgb;

  uvec4 lightLayerAndNormal = imageLoad (lightLayerData, globalThread);

  // light layer index
  uint lightLayer =  lightLayerAndNormal.x;
  
  // Normals are stored in 2 32-bit uints (only first 48 bits are used) where they can be unpacked in 3 floats so we unpack them here.
  vec3 worldNormal = vec3 (unpackHalf2x16 (lightLayerAndNormal.y).xy, unpackHalf2x16 (lightLayerAndNormal.z).x);

  vec3 fragColor = vec3 (0.0f);

  vec4 shadowMapColor = vec4 (1.0f);

  // Shadow multiplier
  float shadowValue = 1.0f;

  for (int i = 0; i < lightCounts.ambientLights; ++i)
  {
    if ((lightLayer & AmbLightData.aLightData[i].cullingMask) != 0)
    {
      // Just do some add
      fragColor += pixelDiffuse.rgb * AmbLightData.aLightData[i].ambientColor.rgb * vec3 (AmbLightData.aLightData[i].strength);
    }
  }

  for (int i = 0; i < lightCounts.directionalLights; ++i)
  {
    if ((lightLayer & DirLightData.dLightData[i].cullingMask) != 0)
    {
      // get normalized direction of light 
      vec3 dLightNormalized = normalize (DirLightData.dLightData[i].direction);

      // Get diffuse strength
      float diffuseStrength = max (0, dot (-dLightNormalized, normalView));

      // Calculate the fragment color
      fragColor += DirLightData.dLightData[i].diffuseColor.rgb * diffuseStrength.rrr * pixelDiffuse;

      // If the shadow map is enabled (test the bit)
      if ((DirLightData.dLightData[i].shadowData & uint(1)) == 1)
      {
        uint shadowMapIndex = (DirLightData.dLightData[i].shadowData >> 8);
        shadowValue = min (shadowValue, CalcShadowValue (shadowMaps[nonuniformEXT(shadowMapIndex)], positionWorld, DirLightData.dLightData[i].pvMatrix, worldNormal, DirLightData.dLightData[i].directionWorld.xyz));
      }
    }
  }

  // calculate shadow map here
  if (shadowValue != 0.0f)
    fragColor.rgb *= shadowValue;

  float ssaoVal = imageLoad (ssaoBlurredImage, globalThread).r;
  fragColor *= ssaoVal;

  vec4 objectVFXColor = imageLoad (objectVFXImage, globalThread);
  fragColor += objectVFXColor.rgb * objectVFXColor.a;

  // store result into result image
  imageStore(targetImage, ivec2(gl_GlobalInvocationID.xy), vec4(fragColor.rgb, 1.0f));
}