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

#pragma once

#include <Atom/Features/PBR/LightingUtils.azsli>
#include <Atom/Features/PBR/Microfacet/Fresnel.azsli>

// compute final probe specular using the probe cubemap and the roughness, normals, and specularF0 for the surface
bool ComputeProbeSpecular(float2 screenCoords, float3 positionWS, float4x4 obbTransformInverse, float3 outerObbHalfLengths, uint sampleIndex, out float3 specular)
{
    // make sure the pixel belongs to this probe volume
    // this is necessary since it could have the correct stencil value but actually reside 
    // in another volume that's in between the camera and the volume we're rendering
    if (!ObbContainsPoint(obbTransformInverse, outerObbHalfLengths, positionWS))
    {
        return false;
    }

    // retrieve normal from the encoded buffer written by the forward pass
    float4 encodedNormal = PassSrg::m_normal.Load(screenCoords, sampleIndex);
    float3 normal = DecodeNormalSignedOctahedron(encodedNormal.rgb);
    bool multiScatterCompensationEnabled = (encodedNormal.a > 0.0f);

    float3 dirToCamera = normalize(ViewSrg::m_worldPosition.xyz - positionWS);
    float NdotV = dot(normal, dirToCamera);
    NdotV = max(NdotV, 0.01f);    // [GFX TODO][ATOM-4466] This is a current band-aid for specular noise at grazing angles.

    // retrieve specularF0 and roughness
    float4 specularF0Sample = PassSrg::m_specularF0.Load(screenCoords, sampleIndex);
    float3 specularF0 = specularF0Sample.rgb;
    float roughness = specularF0Sample.a;

    // sample BRDF map (using smoothness)
    float smoothness = 1.0f - roughness;
    float2 brdfUV = float2(saturate(NdotV), smoothness);
    float2 brdf = PassSrg::m_brdfMap.Sample(PassSrg::LinearSampler, brdfUV).rg;

    float3 reflectDir = reflect(-dirToCamera, normal);

    // compute parallax corrected reflection vector, if necessary
    float3 localReflectDir = reflectDir;
    if (ObjectSrg::m_useParallaxCorrection)
    {
        localReflectDir = ApplyParallaxCorrectionOBB(
            ObjectSrg::GetWorldMatrixInverse(),
            ObjectSrg::m_outerObbHalfLengths,
            positionWS,
            reflectDir);
    }

    // sample reflection cubemap with the appropriate roughness mip
    float3 probeSpecular = ObjectSrg::m_reflectionCubeMap.SampleLevel(SceneSrg::m_samplerEnv, GetCubemapCoords(localReflectDir), GetRoughnessMip(roughness)).rgb;

    // compute final specular amount
    float3 multiScatterCompensation = GetMultiScatterCompensation(specularF0, brdf, multiScatterCompensationEnabled);
    specular = probeSpecular * multiScatterCompensation * (specularF0.xyz * brdf.x + brdf.y);

    return true;
}