/*
 * 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
 *
 */

#include <viewsrg.srgi>
#define UvSetCount 2
#include <Atom/RPI/ShaderResourceGroups/DefaultObjectSrg.azsli>
#include <Atom/RPI/ShaderResourceGroups/DefaultDrawSrg.azsli>
#include <Atom/RPI/TangentSpace.azsli>

// These enums need to be kept in sync with the enum values in DebugVertexStreams_IncompatibleEnums.lua
option enum class DebugVertexStream { Normals, Tangents, Bitangents, Uvs, TangentW } o_debugVertexStream = DebugVertexStream::Normals;
option enum class TangentOptions { UseVertexData, UseSurfaceGradient} o_tangentOptions = TangentOptions::UseVertexData;
option enum class BitangentOptions { UseVertexData, UseSurfaceGradient, ReconstructBitangent} o_bitangentOptions = BitangentOptions::UseVertexData;
option enum class ColorDisplayMode {ColorSpace, UnitSpace} o_colorDisplayMode = ColorDisplayMode::ColorSpace;

ShaderResourceGroup MaterialSrg : SRG_PerMaterial
{
    uint m_uvIndex;
}

struct VSInput
{
    // Base fields (required by the template azsli file)...
    float3 m_position : POSITION;
    float3 m_normal : NORMAL;
    float4 m_tangent : TANGENT; 
    float3 m_bitangent : BITANGENT; 
 
    // Extended fields (only referenced in this azsl file)...
    float2 m_uv0 : UV0;
    float2 m_uv1 : UV1;
};

struct VSOutput
{
    // Base fields (required by the template azsli file)...
    // "centroid" is needed for SV_Depth to compile
    precise linear centroid float4 m_position : SV_Position;
    float3 m_normal: NORMAL;
    float4 m_tangent : TANGENT; 
    float3 m_bitangent : BITANGENT; 
    float3 m_worldPosition : UV0;

    // Extended fields (only referenced in this azsl file)...
    float2 m_uv[UvSetCount] : UV1;
};

VSOutput MainVS(VSInput IN)
{
    VSOutput OUT;    
    OUT.m_worldPosition = mul(ObjectSrg::GetWorldMatrix(), float4(IN.m_position, 1.0)).xyz;
    OUT.m_position = mul(ViewSrg::m_viewProjectionMatrix, float4(OUT.m_worldPosition, 1.0));

    // Only UV0 is supported
    OUT.m_uv[0] = IN.m_uv0;
    OUT.m_uv[1] = IN.m_uv1;    
    
    float4x4 objectToWorld = ObjectSrg::GetWorldMatrix();
    float3x3 objectToWorldIT = ObjectSrg::GetWorldMatrixInverseTranspose();

    ConstructTBN(IN.m_normal, IN.m_tangent, IN.m_bitangent, objectToWorld, objectToWorldIT, OUT.m_normal, OUT.m_tangent.xyz, OUT.m_bitangent);

    OUT.m_tangent.w = IN.m_tangent.w;
    return OUT;
}

struct PixelOutput
{
    float4 m_color : SV_Target0;
};


float3 OffsetColor(float3 color)
{
    if(o_colorDisplayMode == ColorDisplayMode::ColorSpace)
    {
        // Represent a vector in the (-1, -1, -1) to (1, 1, 1) range as a color in the (0, 0, 0) to (1, 1, 1) range
        // Color key
        // + x-axis: Light Coral
        // - x-axis: Teal
        // + y-axis: Bright Green
        // - y-axis: Dark Magenta
        // + z-axis: Medium Slate Blue
        // - z-axis: Olive
        return normalize(color) * 0.5 + 0.5;
    }
    else
    {
        // Use the normalized color, with any negative values represented as black
        // + x-axis: Red
        // - x-axis: Black
        // + y-axis: Green
        // - y-axis: Black
        // + z-axis: Blue
        // - z-axis: Black
        return normalize(color);
    }
}

PixelOutput MainPS(VSOutput IN)
{
    PixelOutput OUT;
    
    float3 tangents[UvSetCount] = { IN.m_tangent.xyz, IN.m_tangent.xyz };
    float3 bitangents[UvSetCount] = { IN.m_bitangent.xyz, IN.m_bitangent.xyz };
    
    if(o_bitangentOptions == BitangentOptions::ReconstructBitangent && MaterialSrg::m_uvIndex == 0)
    {
        bitangents[MaterialSrg::m_uvIndex] = cross(IN.m_normal.xyz, IN.m_tangent.xyz) * sign(IN.m_tangent.w);
    }
    else if((o_debugVertexStream == DebugVertexStream::Tangents && o_tangentOptions == TangentOptions::UseSurfaceGradient)
    || (o_debugVertexStream == DebugVertexStream::Bitangents && o_bitangentOptions == BitangentOptions::UseSurfaceGradient)
    || MaterialSrg::m_uvIndex > 0)
    {
        const bool isBackface = false;
        SurfaceGradientNormalMapping_Init(IN.m_normal, IN.m_worldPosition, isBackface);
        SurfaceGradientNormalMapping_GenerateTB(IN.m_uv[MaterialSrg::m_uvIndex], tangents[MaterialSrg::m_uvIndex], bitangents[MaterialSrg::m_uvIndex]);
    }


    float3 outColor = float3(1.0, 1.0, 1.0);
    switch(o_debugVertexStream)
        {
        case DebugVertexStream::Normals:
            outColor = OffsetColor(IN.m_normal);
            break;
        case DebugVertexStream::Tangents:
            outColor = OffsetColor(tangents[MaterialSrg::m_uvIndex]);
            break;
        case DebugVertexStream::Bitangents:
            outColor = OffsetColor(bitangents[MaterialSrg::m_uvIndex]);
            break;
        case DebugVertexStream::Uvs:
            // Assume a tiled uv visualization, where anything greater than 1 wraps back around to 0
            outColor = float3(frac(IN.m_uv[MaterialSrg::m_uvIndex].x), frac(IN.m_uv[MaterialSrg::m_uvIndex].y), 0.0f);
            break;            
        case DebugVertexStream::TangentW:
            float red = IN.m_tangent.w >= 0.0f ? 1.0f : 0.0f;
            float green = IN.m_tangent.w <= 0.0f ? 1.0f : 0.0f;
            float blue = IN.m_tangent.w == 0.0f ? 1.0f : 0.0f;
            outColor = float3(red, green, blue);
            break;
        }

    OUT.m_color.rgb = outColor;
	OUT.m_color.a = 1.0;

    return OUT;
}