#include "graphics_structs.hlsli"

struct FS_Input
{
    float4 InPosition   : POSITION;
    float4 InNormal     : NORMAL;
    float4 InColor      : COLOR0;
    float2 InUV0        : TEXCOORD0;
};

struct FS_Output
{
    float4 ColorTarget : SV_Target0;
};

float4 GetAlbedo(float2 UV, float4 InColor)
{
    float4 AlbedoFactor = (float4) MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].AlbedoFactor;
    uint AlbedoTexId = MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].AlbedoTex;
    
    return AlbedoFactor * InColor * (AlbedoTexId != INVALID_HANDLE ? Textures[AlbedoTexId].Sample(ImmutableSamplers[AlbedoTexId], UV) : 1.0f.xxxx);
}

float GetOcclusion(float2 UV)
{
    uint OcclusionTexId = MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].OcclusionTex;
    
    return OcclusionTexId != INVALID_HANDLE ? Textures[OcclusionTexId].Sample(ImmutableSamplers[OcclusionTexId], UV).r : 1.0f;
}

float3 GetEmissive(float2 UV)
{
    float3 EmissionFactor = (float3) MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].EmissionFactor;
    uint EmissionTexId = MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].EmissionTex;
    
    return EmissionFactor * (EmissionTexId != INVALID_HANDLE ? Textures[EmissionTexId].Sample(ImmutableSamplers[EmissionTexId], UV).rgb : 1.0f.xxx);
}

float3 GetNormal(float3 Position, float3 Normal, float2 UV)
{
    float3 N = normalize(Normal);
    
    uint NormalTexId = MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].NormalTex;
    if (NormalTexId == INVALID_HANDLE)
    {
        return N;
    }
        
    float3 TangentSpaceNormal = Textures[NormalTexId].Sample(ImmutableSamplers[NormalTexId], UV).xyz * 2.0f - 1.0f;

    float3 q1 = ddx(Position);
    float3 q2 = ddy(Position);
    float2 st1 = ddx(UV);
    float2 st2 = ddy(UV);
    
    float3 T = normalize(q1 * st2.y - q2 * st1.y).xyz;
    float3 B = -normalize(cross(N, T));
    float3x3 TBN = float3x3(T, B, N); // Construction is Row by Row.

    return normalize(mul(TangentSpaceNormal, TBN)); // Post multiple to avoid transpose.
}

float2 GetMetalRough(float2 UV)
{
    float2 MetalRoughFactors = float2(MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].MetalFactor, MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].RoughFactor);
    uint MetalRoughTexId = MaterialsBuffer[PushConstant.MaterialBufferHandle][PushConstant.MaterialIdx].MetalRoughTex;
    
    return MetalRoughFactors * (MetalRoughTexId != INVALID_HANDLE ? Textures[MetalRoughTexId].Sample(ImmutableSamplers[MetalRoughTexId], UV).bg : 1.0f.xx); // Metal is B, Rough is G.
}

float3 SampleIrradiance(float3 Direction)
{
    if (PushConstant.DiffuseIrradianceHandle != INVALID_HANDLE)
    {
        return TextureCubes[PushConstant.DiffuseIrradianceHandle].Sample(ImmutableSamplers[PushConstant.DiffuseIrradianceHandle], Direction).rgb;
    }
    return 0.04f.xxx;
}

float TrowbridgeReitzGGX(float3 Normal, float3 Halfway, float Roughness)
{
    float Coeff = Roughness * Roughness;
    float Alpha2 = Coeff * Coeff;
    float NdotH = max(dot(Normal, Halfway), 0.0f);
    float NdotH2 = NdotH * NdotH;
	
    float Numerator = Alpha2;
    float Denominator = NdotH2 * (Alpha2 - 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
float3 FresnelSchlick(float cosine, float3 F_0)
{
    return F_0 + (1.0f - F_0) * pow(clamp(1.0f - cosine, 0.0f, 1.0f), 5.0f); // Clamp to avoid artifacts.
}

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

float3 GetPBRContrib(float3 Albedo, float3 LightColor, float3 ViewDir, float3 Normal, float Metallic, float Roughness, float3 F_0, float3 LightDir, float LightDistance)
{
    float Attenuation = 1.0f / (LightDistance * LightDistance); // TODO: Controlled Attenuation
    float3 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);
        
    float3 Radiance = LightColor * Attenuation;
    
    float NormalDistribution = TrowbridgeReitzGGX(Normal, Halfway, Roughness);
    float Geometry = GeometrySmith(NdotV, NdotL, Roughness);
    float3 Fresnel = FresnelSchlickRoughness(CosineFactor, F_0, Roughness);
        
    float3 Numerator = (NormalDistribution * Geometry) * Fresnel;
    float Denominator = 4.0f * NdotV * NdotL;
    float3 Specular = Numerator / (Denominator + 0.00001f);
        
    float3 K_Specular = Fresnel;
    float3 K_Diffuse = 1.0f.xxx - K_Specular;
        
    K_Diffuse *= 1.0f - Metallic;
        
    return NdotL * Radiance * (K_Diffuse * Albedo / PI + Specular);
}

float3 GetPointLightInfluence(float3 Albedo, float2 MetalRough, float3 Position, float3 Normal)
{
    if (PushConstant.LightHandle == INVALID_HANDLE) 
        return 0.0f.xxx;
    
    uint Offset = IndexerOffset(PushConstant.PointLightIndexer);
    uint Count = IndexerCount(PushConstant.PointLightIndexer);
        
    float3 ViewDir = normalize(Camera.Position.xyz - Position);
    
    float Metallic = MetalRough.r;
    float Roughness = MetalRough.g;
    
    // Dielectric F_0 based on LearnOpenGL.
    // TODO: Cite    
    float3 F_0 = 0.04f.xxx;
    F_0 = lerp(F_0, Albedo, Metallic);
    
    float3 Contrib = 0.0f;
    for (uint i = 0; i < Count; ++i)
    {
        PointLight Light = PointLightBuffer[PushConstant.LightHandle][i + Offset];
        
        if (Light.Range < 0.0f)
            continue;
        
        float3 LightDir = float3(Light.Position) - Position;
        float LightDistance = length(LightDir);
        
        if (LightDistance > Light.Range)
            continue;
        
        LightDir /= LightDistance; // Normalization
        
        // Color Unpack
        float R = (Light.Color & 0xFF000000) >> 24;
        float G = (Light.Color & 0x00FF0000) >> 16;
        float B = (Light.Color & 0x0000FF00) >> 8;
        
        float3 LightColor = Light.Intensity * float3(R, G, B) * 0.00392156862f; // 0.00392156862 = 1/255
 
        Contrib += GetPBRContrib(Albedo, LightColor, ViewDir, Normal, Metallic, Roughness, F_0, LightDir, LightDistance);
    }
    
    return Contrib;
}


float3 GetDirectionalLightInfluence(float3 Albedo, float2 MetalRough, float3 Position, float3 Normal)
{
    if (PushConstant.LightHandle == INVALID_HANDLE) 
        return 0.0f.xxx;
    
    uint Count = IndexerCount(PushConstant.DirectionalLightIndexer);
    
    float3 ViewDir = normalize(Camera.Position.xyz - Position);
    
    float Metallic = MetalRough.r;
    float Roughness = MetalRough.g;
    
    // Dielectric F_0 based on LearnOpenGL.
    // TODO: Cite    
    float3 F_0 = 0.04f.xxx;
    F_0 = lerp(F_0, Albedo, Metallic);
    
    float3 Contrib = 0.0f;
    for (uint i = 0; i < Count; ++i)
    {
        DirectionalLight Light = DirectionalLightBuffer[PushConstant.LightHandle][i];
        
        if (Light.Validity_ < 0.0f)
            continue;
        
        float3 LightDir = -normalize(float3(Light.Direction));
        
        // Color Unpack
        float R = (Light.Color & 0xFF000000) >> 24;
        float G = (Light.Color & 0x00FF0000) >> 16;
        float B = (Light.Color & 0x0000FF00) >> 8;
        
        float3 LightColor = Light.Intensity * float3(R, G, B) * 0.00392156862f; // 0.00392156862 = 1/255
 
        Contrib += GetPBRContrib(Albedo, LightColor, ViewDir, Normal, Metallic, Roughness, F_0, LightDir, 1.0f);
    }
    
    return Contrib;
}

float3 GetAmbientInfluence(float3 Albedo, float2 MetalRough, float3 Position, float3 Normal, float Occlusion)
{
    float3 ViewDir = normalize(Camera.Position.xyz - Position);
    float CosineFactor = max(dot(Normal, ViewDir), 0.0f);       // TODO: Is Normal or is it Halfway?
    
    float3 F_0 = 0.04f.xxx;
    F_0 = lerp(F_0, Albedo, MetalRough.r);
    float3 K_Specular = FresnelSchlickRoughness(CosineFactor, F_0, MetalRough.g);
    float3 K_Diffuse = 1.0f.xxx - K_Specular;
        
    K_Diffuse *= 1.0f - MetalRough.r;   // Metals don't have diffuse/refractions.
    
    float3 DiffuseIrradiance = K_Diffuse * Albedo * SampleIrradiance(Normal) * Occlusion;

    return DiffuseIrradiance;
}

FS_Output main(FS_Input StageInput)
{
    FS_Output Output;
    
    // TODO: This should be invalid on the CPU side.
    if (PushConstant.MaterialIdx < 0)
    {
        Output.ColorTarget = float4(1.0f, 0.0f, 1.0f, 1.0f);
        return Output;
    }

    float3 Position = StageInput.InPosition.xyz;
    float3 Normal = GetNormal(Position, StageInput.InNormal.xyz, StageInput.InUV0);
    
    float4 AlbedoAlpha = GetAlbedo(StageInput.InUV0, StageInput.InColor);
    float3 Albedo = AlbedoAlpha.rgb;
    float Alpha = AlbedoAlpha.a;
    float2 MetalRough = GetMetalRough(StageInput.InUV0);
    float3 Emission = GetEmissive(StageInput.InUV0);
    float Occlusion = GetOcclusion(StageInput.InUV0);
    
    float3 DirectionalLightLum = GetDirectionalLightInfluence(Albedo, MetalRough, Position, Normal);
    float3 PointLighLum = GetPointLightInfluence(Albedo, MetalRough, Position, Normal);
    float3 AmbientLum = GetAmbientInfluence(Albedo, MetalRough, Position, Normal, Occlusion);
    float3 Color = DirectionalLightLum + PointLighLum + AmbientLum;
    
    Output.ColorTarget = float4(Uncharted2Tonemap(Color + Emission), Alpha);
    return Output;
}