#version 460
#pragma shader_stage(fragment)
#extension GL_EXT_shader_explicit_arithmetic_types_int64 : enable
#extension GL_EXT_buffer_reference : require
#extension GL_EXT_nonuniform_qualifier : enable

#include "bindless_structs.glsl"
#include "graphics_bindings.glsl"

layout (location = 0) in vec4 in_Position;
layout (location = 1) in vec4 in_Normal;
layout (location = 2) in vec4 in_Color0;
layout (location = 3) in vec2 in_TexCoord0;
layout (location = 4) in flat uint in_DrawID;

layout (location = 0) out vec4 outColor;

layout (push_constant) uniform Const {
    NodeRef nodes;
};

vec4 GetAlbedo(vec4 albedoFactor, uint texHandle, vec4 in_Color0, vec2 uv)
{
    return albedoFactor * in_Color0 * ((texHandle != INVALID_HANDLE) ? texture(textures[texHandle], uv) : vec4(1.0f));
}

float GetOcclusion(uint texHandle, vec2 uv)
{
    return texHandle != INVALID_HANDLE ? texture(textures[texHandle], uv).r : 1.0f;
}

vec3 GetEmissive(vec3 emissionFactor, uint texHandle, vec2 uv)
{
    return emissionFactor * (texHandle != INVALID_HANDLE ? texture(textures[texHandle], uv).rgb : vec3(1.0f));
}

vec3 GetNormal(uint texHandle, vec3 position, vec3 normal, vec2 uv)
{
    vec3 N = normalize(normal);
    
    if (texHandle == INVALID_HANDLE)
    {
        return N;
    }
        
    vec3 tanSpaceNormal = texture(textures[texHandle], uv).xyz * 2.0f - 1.0f;

    vec3 q1 = dFdx(position);
    vec3 q2 = dFdy(position);
    vec2 st1 = dFdx(uv);
    vec2 st2 = dFdy(uv);
    
    vec3 T = normalize(q1 * st2.y - q2 * st1.y).xyz;
    vec3 B = -normalize(cross(N, T));
    mat3 TBN = mat3(T, B, N); // Construction is Col by Col

    return normalize(TBN * tanSpaceNormal);
}

vec2 GetMetalRough(float metalFactor, float roughFactor, uint texHandle, vec2 uv)
{
    return vec2(metalFactor, roughFactor) * (texHandle != INVALID_HANDLE ? texture(textures[texHandle], uv).bg : vec2(1.0f)); // Metal is B, Rough is G.
}

vec3 SampleIrradiance(vec3 direction)
{
    uint texHandle = lights.m_DiffuseIrradianceHandle;
    return ((texHandle != INVALID_HANDLE) ? texture(textureCubes[texHandle], direction).rgb : vec3(0.04f));
}

vec3 SamplePrefiltered(vec3 direction, float roughness)
{
    const float MAX_MIP_LEVEL = 5.0f;
    float mipLevel = MAX_MIP_LEVEL * roughness;

    uint texHandle = lights.m_PrefilterHandle;

    return (texHandle != INVALID_HANDLE) ? textureLod(textureCubes[texHandle], direction, mipLevel).rgb : vec3(0.0f);
}

vec2 SampleBrdfLut(float ndotv, float roughness) 
{
    return lights.m_BrdfLutHandle != INVALID_HANDLE ? texture(textures[lights.m_BrdfLutHandle], vec2(ndotv, roughness)).rg : 0.0f.xx;
}


float TrowbridgeReitzGGX(vec3 normal, vec3 halfway, float roughness)
{
    float coeff = roughness * roughness;
    float coeff2 = coeff * coeff;
    float ndoth = max(dot(normal, halfway), 0.0f);
    float ndoth2 = ndoth * ndoth;
	
    float numerator = coeff2;
    float denominator = ndoth2 * (coeff2 - 1.0f) + 1.0f;
    denominator = PI * denominator * denominator;
	
    return numerator / denominator;
}

float GeometrySchlickGGX(float ndotv, float roughness)
{
    float r = roughness + 1.0f;
    float k = (r * r) / 8.0f;

    float numerator = ndotv;
    float denominator = ndotv * (1.0f - k) + k;
	
    return numerator / denominator;
}

float GeometrySmith(float ndotv, float ndotl, float roughness)
{
    float ggx1 = GeometrySchlickGGX(ndotv, roughness);
    float ggx2 = GeometrySchlickGGX(ndotl, roughness);
	
    return ggx1 * ggx2;
}

// https://en.wikipedia.org/wiki/Schlick%27s_approximation
// TODO: Possibly needs fixing. unreal vs LearnOpenGL impl.
vec3 FresnelSchlick(float cosine, vec3 f0)
{
    return f0 + (1.0f - f0) * pow(clamp(1.0f - cosine, 0.0f, 1.0f), 5.0f); // Clamp to avoid artifacts.
}

// Sebastian Lagarde
vec3 FresnelSchlickRoughness(float cosine, vec3 f0, float roughness)
{
    return f0 + (max((1.0f - roughness).xxx, f0) - f0) * pow(clamp(1.0f - cosine, 0.0f, 1.0f), 5.0f); // Clamp to avoid artifacts.
}

vec3 GetPBRContrib(vec3 albedo, vec3 lightColor, vec3 viewDir, vec3 normal, float metallic, float roughness, vec3 f0, vec3 lightDir, float lightDistance)
{
    float attenuation = 1.0f / (lightDistance * lightDistance); // TODO: Controlled Attenuation
    vec3 halfway = normalize(viewDir + lightDir);
        
    float cosineFactor = max(dot(halfway, viewDir), 0.0f);
    float ndotv = max(dot(normal, viewDir), 0.0f);
    float ndotl = max(dot(normal, lightDir), 0.0f);
        
    vec3 radiance = lightColor * attenuation;
    
    float normalDistribution = TrowbridgeReitzGGX(normal, halfway, roughness);
    float geometry = GeometrySmith(ndotv, ndotl, roughness);
    vec3 fresnel = FresnelSchlickRoughness(cosineFactor, f0, roughness);
        
    vec3 numerator = (normalDistribution * geometry) * fresnel;
    float denominator = 4.0f * ndotv * ndotl;
    vec3 specular = numerator / (denominator + 0.00001f);
        
    vec3 kSpecular = fresnel;
    vec3 kDiffuse = 1.0f.xxx - kSpecular;
        
    kDiffuse *= 1.0f - metallic;
        
    return ndotl * radiance * (kDiffuse * albedo / PI + specular);
}

vec3 GetPointLightInfluence(vec3 albedo, vec2 metalRough, vec3 position, vec3 normal)
{
    if (lights.m_LightHandle == INVALID_HANDLE) 
        return 0.0f.xxx;
    
    uint offset = IndexerOffset(lights.m_PointLightIndexer);
    uint count = IndexerCount(lights.m_PointLightIndexer);
        
    vec3 viewDir = normalize(camera.m_Position.xyz - position);
    
    float metallic = metalRough.r;
    float roughness = metalRough.g;
    
    // Dielectric F_0 based on LearnOpenGL.
    // TODO: Cite    
    vec3 f0 = vec3(0.04f);
    f0 = mix(f0, albedo, metallic);
    
    vec3 contrib = vec3(0.0f);
    for (uint i = 0; i < count; ++i)
    {
        PointLight light = pointLightBuffers[lights.m_LightHandle].lights[i + offset];
        
        if (light.m_Range < 0.0f)
            continue;
        
        vec3 lightDir = vec3(light.m_Position) - position;
        float lightDistance = length(lightDir);
        
        if (lightDistance > light.m_Range)
            continue;
        
        lightDir /= lightDistance; // Normalization
        
        // Color Unpack
        float r = (light.m_Color & 0xFF000000) >> 24;
        float g = (light.m_Color & 0x00FF0000) >> 16;
        float b = (light.m_Color & 0x0000FF00) >> 8;
        
        vec3 lightColor = light.m_Intensity * vec3(r, g, b) * 0.00392156862f; // 0.00392156862 = 1/255
 
        contrib += GetPBRContrib(albedo, lightColor, viewDir, normal, metallic, roughness, f0, lightDir, lightDistance);
    }
    
    return contrib;
}

vec3 GetDirectionalLightInfluence(vec3 albedo, vec2 metalRough, vec3 position, vec3 normal)
{
    if (lights.m_LightHandle == INVALID_HANDLE) 
        return 0.0f.xxx;
    
    uint count = IndexerCount(lights.m_DirectionalLightIndexer);
    
    vec3 viewDir = normalize(camera.m_Position.xyz - position);
    
    float metallic = metalRough.r;
    float roughness = metalRough.g;
    
    // Dielectric F_0 based on LearnOpenGL.
    // TODO: Cite    
    vec3 f0 = vec3(0.04f);
    f0 = mix(f0, albedo, metallic);
    
    vec3 contrib = vec3(0.0f);
    for (uint i = 0; i < count; ++i)
    {
        DirectionalLight light = directionalLightBuffers[lights.m_LightHandle].lights[i];
        
        if (light.m_Validity_ < 0.0f)
            continue;
        
        vec3 lightDir = -normalize(light.m_Direction);
        
        // Color Unpack
        float r = (light.m_Color & 0xFF000000) >> 24;
        float g = (light.m_Color & 0x00FF0000) >> 16;
        float b = (light.m_Color & 0x0000FF00) >> 8;
        
        vec3 lightColor = light.m_Intensity * vec3(r, g, b) * 0.00392156862f; // 0.00392156862 = 1/255
 
        contrib += GetPBRContrib(albedo, lightColor, viewDir, normal, metallic, roughness, f0, lightDir, 1.0f);
    }
    
    return contrib;
}

vec3 GetAmbientInfluence(vec3 albedo, float metal, float roughness, vec3 Position, vec3 Normal, float Occlusion)
{
    vec3 ViewDir = normalize(camera.m_Position.xyz - Position);
    float CosineFactor = max(dot(Normal, ViewDir), 0.0f);       // Normal instead of Halfway since there's no halfway in ambient.

    vec3 F_0 = 0.04f.xxx;
    F_0 = mix(F_0, albedo, metal);
    vec3 K_Specular = FresnelSchlickRoughness(CosineFactor, F_0, roughness);
    vec3 K_Diffuse = 1.0f.xxx - K_Specular;
        
    K_Diffuse *= 1.0f - metal; // Metals don't have diffuse/refractions.
     
    vec3 ReflectionDir = reflect(-ViewDir, Normal);
    
    float NdotV = max(dot(Normal, ViewDir), 0.0f);
    vec3 PrefilteredColor = SamplePrefiltered(ReflectionDir, roughness).rgb;
    vec2 EnvBRDF = SampleBrdfLut(NdotV, roughness);
    vec3 Specular = PrefilteredColor * (K_Specular * EnvBRDF.x + EnvBRDF.y);
    
    vec3 DiffuseIrradiance = albedo * SampleIrradiance(Normal);
//#ifdef _DEBUG
//    if ((PushConstant.DebugFlags & USE_DIFFUSE_BIT) == 0) {
//        DiffuseIrradiance = 0.0f.xxx;
//    }
//    if ((PushConstant.DebugFlags & USE_SPECULAR_BIT) == 0) {
//        Specular = 0.0f.xxx;
//    }
//#endif
    
    return (K_Diffuse * DiffuseIrradiance + Specular) * Occlusion;
}

void main() {

    Material material = nodes.nodes[nonuniformEXT(in_DrawID)].pMaterial.material;
    
    vec4 albedoA = GetAlbedo(material.m_AlbedoFactor, material.m_AlbedoTex, in_Color0, in_TexCoord0);
    vec3 albedo = albedoA.rgb;
    float alpha = albedoA.a;
    
    vec3 normal = GetNormal(material.m_NormalTex, in_Position.xyz, normalize(in_Normal.xyz), in_TexCoord0);

    vec3 emission = GetEmissive(material.m_EmissionFactor, material.m_EmissionTex, in_TexCoord0);
    
    vec2 metalRough = GetMetalRough(material.m_MetalFactor, material.m_RoughFactor, material.m_MetalRoughTex, in_TexCoord0);
    float occlusion = GetOcclusion(material.m_OcclusionTex, in_TexCoord0);

    vec3 pointLumin = GetPointLightInfluence(albedo, metalRough, in_Position.xyz, normal);
    vec3 directionalLumin = GetDirectionalLightInfluence(albedo, metalRough, in_Position.xyz, normal);
    vec3 ambientLumin = GetAmbientInfluence(albedo, metalRough.r, metalRough.g, in_Position.xyz, normal, occlusion);
    
    vec3 viewDir = normalize(camera.m_Position.xyz - in_Position.xyz);
    outColor = vec4(Uncharted2Tonemap(pointLumin + directionalLumin + ambientLumin + emission), alpha);
}