/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

// GlobalSrg::m_floatBuffer
// PerSceneSrg::m_textureArray
// PerSceneSrg::m_sampler
#include "BindlessPrototypeSrg.azsli"

struct BindlessMaterial0
{
    uint4 materialIndex;

    float4 m_diffuseColor;
};

struct BindlessMaterial1
{
    uint4 materialIndex;

    float4 m_diffuseColor;
    uint m_diffuseTextureIndex;
};

struct BindlessMaterial2
{
    uint4 materialIndex;

    float4 m_diffuseColor;
    uint m_diffuseTextureIndex;
    uint m_normalTextureIndex;
    uint m_specularTextureIndex;
};

struct PerObject
{
    float4x4 m_localToWorldMatrix;
    float4 rotation;
};

ShaderResourceGroupSemantic PerSubMesh
{
    FrequencyId = 1;
};

// [TODO ATOM-2769] When the inline feature is complete, use InlineConstant instead of sending the data
// via ConstantBuffer.
ShaderResourceGroup HandleSrg : PerSubMesh
{
    uint m_perViewHandle;
    uint m_perObjectHandle;
    uint m_materialHandle;
    uint m_lightHandle;
};

struct VertexInput
{
    float3 m_position : POSITION;
    float3 m_normal : NORMAL;
    float3 m_tangent : TANGENT;
    float3 m_bitangent : BITANGENT;
    float2 m_uv : UV0;
};

struct VertexOutput
{
    float4 m_position : SV_Position;
    float3 m_normal : NORMAL;
    float3 m_tangent : TANGENT;
    float3 m_bitangent : BITANGENT;
    float2 m_uv : UV0;
};

VertexOutput MainVS(VertexInput vsInput)
{
    VertexOutput OUT;

    // Read the PerObject data from the FloatBuffer
    uint offset = 0;
    PerObject perObject;
    {
        ReadFromFloatBuffer(perObject.m_localToWorldMatrix, HandleSrg::m_perObjectHandle, offset);
        ReadFromFloatBuffer(perObject.rotation, HandleSrg::m_perObjectHandle, offset);
    }

    // Read the world matrix from the FloatBuffer
    float4x4 worldToClipMatrix;
    offset = 0;
    {
        ReadFromFloatBuffer(worldToClipMatrix, HandleSrg::m_perViewHandle, offset);
    }

    const float4 worldPosition = mul(perObject.m_localToWorldMatrix,  float4(vsInput.m_position, 1.0));
    OUT.m_position = mul(worldToClipMatrix, worldPosition);
    OUT.m_uv = vsInput.m_uv;
    OUT.m_normal = normalize(vsInput.m_normal);

    return OUT;
}

struct PixelOutput
{
    float4 m_color : SV_Target0;
};

PixelOutput MainPS(VertexOutput psInput)
{
    PixelOutput OUT;
    uint offset = 0;
    
    // Read the material index to identify the material type
    uint4 materialIndex;
    {
        ReadFromFloatBuffer(materialIndex, HandleSrg::m_materialHandle, offset);
    }
    
    // Read the material data from the FloatBuffer depending ont he material index
    if(materialIndex.x == 0) // Albedo material
    {
        BindlessMaterial0 bindlessMaterial0; 
        ReadFromFloatBuffer(bindlessMaterial0.m_diffuseColor, HandleSrg::m_materialHandle, offset);
        
        OUT.m_color = float4(bindlessMaterial0.m_diffuseColor.xyz, 1.0);
    }
    else if(materialIndex.x == 1) // Texture sample material
    {
        BindlessMaterial1 bindlessMaterial1; 
        ReadFromFloatBuffer(bindlessMaterial1.m_diffuseColor, HandleSrg::m_materialHandle, offset);
        ReadFromFloatBuffer(bindlessMaterial1.m_diffuseTextureIndex, HandleSrg::m_materialHandle, offset);

        Texture2D texture = ImageSrg::m_textureArray[bindlessMaterial1.m_diffuseTextureIndex % 8]; // % 8 for wrap-around texture index as specified in ImageSrg.m_textureArray
        OUT.m_color = texture.Sample(ImageSrg::m_sampler, psInput.m_uv);
    }
    else if(materialIndex.x == 2) // Shaded material
    {
        float4 color;
        BindlessMaterial2 bindlessMaterial2; 
        ReadFromFloatBuffer(bindlessMaterial2.m_diffuseColor, HandleSrg::m_materialHandle, offset);
        ReadFromFloatBuffer(bindlessMaterial2.m_diffuseTextureIndex, HandleSrg::m_materialHandle, offset);
        ReadFromFloatBuffer(bindlessMaterial2.m_normalTextureIndex, HandleSrg::m_materialHandle, offset);
        ReadFromFloatBuffer(bindlessMaterial2.m_specularTextureIndex, HandleSrg::m_materialHandle, offset);

        Texture2D texture = ImageSrg::m_textureArray[bindlessMaterial2.m_diffuseTextureIndex % 8]; // % 8 for wrap-around texture index as specified in ImageSrg.m_textureArray
        color = texture.Sample(ImageSrg::m_sampler, psInput.m_uv);

        float3 lightDir;
        uint lightOffset = 0;
        ReadFromFloatBuffer(lightDir, HandleSrg::m_lightHandle, lightOffset); 
        lightDir = normalize(-lightDir);

        color *= dot(lightDir, psInput.m_normal) * 8.0;
        OUT.m_color = color;
    }
    else
    {
        OUT.m_color = float4(1.0, 1.0, 1.0, 1.0);
    }

    return OUT;
}