#include "bindless_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(int MaterialIdx, float2 UV)
{
    uint AlbedoTexId = MaterialsBuffer[PushConstant.MaterialBufferHandle][MaterialIdx].AlbedoTex;
    if (AlbedoTexId == INVALID_HANDLE)
    {
        return (float4) MaterialsBuffer[PushConstant.MaterialBufferHandle][MaterialIdx].AlbedoFactor;
    }
    else
    {
        return Textures[AlbedoTexId].Sample(ImmutableSamplers[AlbedoTexId], UV);
    }
}

float2 GetMetalRough(int MaterialIdx, float2 UV)
{
    uint MetalRoughTexId = MaterialsBuffer[PushConstant.MaterialBufferHandle][MaterialIdx].MetalRoughTex;
    if (MetalRoughTexId == INVALID_HANDLE)
    {
        return float2(MaterialsBuffer[PushConstant.MaterialBufferHandle][MaterialIdx].MetalFactor, MaterialsBuffer[PushConstant.MaterialBufferHandle][MaterialIdx].RoughFactor);
    }
    else
    {
        return Textures[MetalRoughTexId].Sample(ImmutableSamplers[MetalRoughTexId], UV).bg; // Metal in Blue, Roughness in Green 
    }
}

float3 GetDirectionalLightInfluence(float3 Normal)
{
    if (PushConstant.LightHandle == INVALID_HANDLE) 
        return float3(0.0f, 0.0f, 0.0f);
    
    uint Count = IndexerCount(PushConstant.DirectionLightIndexer);
    
    float3 Contrib = 0.0f;
    for (uint i = 0; i < Count; ++i)
    {
        Light Light = LightBuffer[PushConstant.LightHandle][i];
        float3 LightDir = - (float3) Light.Position; // Position is actually direction for directionalLight; LightDir is Direction towards the light (-direction)
        float DiffuseFactor = max(dot(Normal, LightDir), 0.0f);
    
        float R = (Light.Color & 0xFF000000) >> 24;
        float G = (Light.Color & 0x00FF0000) >> 16;
        float B = (Light.Color & 0x0000FF00) >> 8;
        
        float3 Color = float3(R, G, B) * 0.00392156862f; // 0.00392156862 = 1/255
        
        float3 Diffuse = DiffuseFactor * Color;
        
        Contrib += (Light.Range < 0 ? float3(0.0f, 0.0f, 0.0f) : Diffuse);
    }
    
    return Contrib;
}

float3 GetPointLightInfluence(float3 Position, float3 Normal)
{
    if (PushConstant.LightHandle == INVALID_HANDLE) 
        return float3(0.0f, 0.0f, 0.0f);
    
    uint Offset = IndexerOffset(PushConstant.PointLightIndexer);
    uint Count = IndexerCount(PushConstant.PointLightIndexer);
    
    float3 Contrib = 0.0f;
    for (uint i = 0; i < Count; ++i)
    {
        Light Light = LightBuffer[PushConstant.LightHandle][i + Offset];
        float3 LightDir = normalize(((float3)Light.Position) - Position);
        float DiffuseFactor = max(dot(Normal, LightDir), 0.0f);
    
        float R = (Light.Color & 0xFF000000) >> 24;
        float G = (Light.Color & 0x00FF0000) >> 16;
        float B = (Light.Color & 0x0000FF00) >> 8;
        
        float3 Color = float3(R, G, B) * 0.00392156862f; // 0.00392156862 = 1/255
        
        float3 Diffuse = DiffuseFactor * Color;
        
        Contrib += (Light.Range < 0 ? float3(0.0f, 0.0f, 0.0f) : Diffuse);
    }
    
    return Contrib;
}

FS_Output main(FS_Input StageInput)
{
    float3 Ambient = float3(0.02f, 0.02f, 0.02f);
    
    float3 Normal = normalize(StageInput.InNormal.xyz);
    float3 Position = StageInput.InPosition.xyz;
    
    float4 ObjColor = PushConstant.MaterialIdx < 0 ? StageInput.InColor : GetAlbedo(PushConstant.MaterialIdx, StageInput.InUV0);
    float2 MetalRough = PushConstant.MaterialIdx < 0 ? float2(0.0f, 0.0f) : GetMetalRough(PushConstant.MaterialIdx, StageInput.InUV0);
    
    float3 Diffuse = GetDirectionalLightInfluence(Normal) + GetPointLightInfluence(Position, Normal);
    
    FS_Output Output;
    Output.ColorTarget = float4(ObjColor.rgb * (Diffuse + Ambient), ObjColor.a);
    return Output;
}