/*
 * 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>
#include <Atom/Features/PBR/AlphaUtils.azsli>
#include <Atom/Features/PBR/DefaultObjectSrg.azsli>
#include <Atom/Features/PBR/ForwardPassSrg.azsli>
#include <Atom/Features/PBR/ForwardPassOutput.azsli>
#include <Atom/Features/ColorManagement/TransformColor.azsli>
#include <Atom/Features/SrgSemantics.azsli>
#include <Atom/Features/PBR/Lighting/StandardLighting.azsli>
#include <Atom/Features/PBR/Decals.azsli>

ShaderResourceGroup MaterialSrg : SRG_PerMaterial
{
    float3 m_emissiveColor;
    float m_emissiveIntensity;
    Texture2D m_emissiveMap;
    
    Sampler m_sampler
    {
       AddressU = Wrap;
       AddressV = Wrap;
       MinFilter = Linear;
       MagFilter = Linear;
       MipFilter = Linear;
    };
}

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

struct VSOutput
{
    precise float4 m_position : SV_Position;
    float3 m_normal: NORMAL;
    float3 m_tangent : TANGENT; 
    float3 m_bitangent : BITANGENT; 
    float3 m_worldPosition : UV0;
    float2 m_uv : UV1;
    float3 m_shadowCoords[ViewSrg::MaxCascadeCount] : UV3;
};

#include <Atom/Features/Vertex/VertexHelper.azsli>
option bool o_emissive_useTexture; 

VSOutput MainVS(VSInput IN)
{
    VSOutput OUT;
    float3 worldPosition = mul(ObjectSrg::GetWorldMatrix(), float4(IN.m_position, 1.0)).xyz;
    VertexHelper(IN, OUT, worldPosition);
    OUT.m_uv = IN.m_uv;
    return OUT;
}

ForwardPassOutput MainPS(VSOutput IN)
{
    // ------- Surface -------

    Surface surface;
    
    // Position, Normal, Roughness
    surface.position = IN.m_worldPosition.xyz;
    surface.normal = surface.vertexNormal = normalize(IN.m_normal);
    surface.roughnessLinear = 1.0f;
    surface.CalculateRoughnessA();

    // Albedo, SpecularF0
    const float3 baseColor = float3(0.05f, 0.05f, 0.05f);
    const float metallic = 0.0f;
    const float specularF0Factor = 0.5f;
    surface.SetAlbedoAndSpecularF0(baseColor, specularF0Factor, metallic);

    // Clear Coat
    surface.clearCoat.InitializeToZero();

    // ------- LightingData -------

    LightingData lightingData;

    // Light iterator
    lightingData.tileIterator.Init(IN.m_position, PassSrg::m_lightListRemapped, PassSrg::m_tileLightData);
    lightingData.Init(surface.position, surface.normal, surface.roughnessLinear);

    // Shadow, Occlusion
    lightingData.shadowCoords = IN.m_shadowCoords;

    // Emissive
    float3 emissive = MaterialSrg::m_emissiveColor.rgb * MaterialSrg::m_emissiveIntensity;
    if (o_emissive_useTexture)
    {
        float4 sampledValue = MaterialSrg::m_emissiveMap.Sample(MaterialSrg::m_sampler, IN.m_uv);
        emissive *= TransformColor(sampledValue.rgb, ColorSpaceId::LinearSRGB, ColorSpaceId::ACEScg);
    }
    lightingData.emissiveLighting = emissive;

    // Diffuse and Specular response
    lightingData.specularResponse = FresnelSchlickWithRoughness(lightingData.NdotV, surface.specularF0, surface.roughnessLinear);
    lightingData.diffuseResponse = 1.0f - lightingData.specularResponse;

    const float alpha = 1.0f;

    // ------- Lighting Calculation -------

    // Apply Decals
    ApplyDecals(lightingData.tileIterator, surface);

    // Apply Direct Lighting
    ApplyDirectLighting(surface, lightingData);

    // Apply Image Based Lighting (IBL)
    ApplyIBL(surface, lightingData);

    // Finalize Lighting
    lightingData.FinalizeLighting();

    PbrLightingOutput lightingOutput = GetPbrLightingOutput(surface, lightingData, alpha);

    ForwardPassOutput OUT;

    OUT.m_diffuseColor = lightingOutput.m_diffuseColor;
    OUT.m_diffuseColor.w = -1; // Subsurface scattering is disabled
    OUT.m_specularColor = lightingOutput.m_specularColor;
    OUT.m_specularF0 = lightingOutput.m_specularF0;
    OUT.m_albedo = lightingOutput.m_albedo;
    OUT.m_normal = lightingOutput.m_normal;
    return OUT;
}