// 约定：row_major 内存布局
//      mul(rowVector, matrix)
//      CPU 侧矩阵不转置

#pragma pack_matrix(row_major)
#include "include/pbr_brdf.hlsli"

cbuffer CBScene : register(b0) {
    row_major float4x4 gView;
    row_major float4x4 gProjection;
    float4 gCameraPosition;
    float4 gLightDirection;          // xyz: 光线从光源射向场景的方向
    float4 gLightColorIntensity;     // rgb: 线性颜色, w: 强度
    float4 gAmbientColorIntensity;   // rgb: 线性颜色, w: 强度；IBL 前的临时环境项
    float4 gRenderParams;            // x: exposure，其余保留
};

cbuffer CBPerMat : register(b1) {
    float4 gBaseColorFactor = float4 (1.0f,1.0f,1.0f,1.0f);
    float4 gEmissiveFactor = float4 (1.0f,1.0f,1.0f,1.0f);  // xyz 有效
    float4 gPbrParams = float4 (1.0f,1.0f,1.0f,1.0f);       // x metallic, y roughness, z normalScale, w AO strength
};

cbuffer CBPerObj : register(b2) {
    row_major float4x4 gWorld;
    row_major float4x4 gWorldInvTranspose;
};

cbuffer CBShadow : register(b3) {
    row_major float4x4 gLightViewProjection;
    float4 gShadowParams;            // xy texelSize, z minBias, w slopeBias
};

Texture2D gBaseColorTex : register(t0); // sRGB SRV
Texture2D gNormalTex : register(t1);    // Linear SRV
Texture2D gOrmTex : register(t2);       // Linear SRV: R=AO, G=Roughness, B=Metallic
Texture2D gEmissiveTex : register(t3);  // sRGB SRV
Texture2D<float> gShadowMap : register(t8);

SamplerState gMaterialSampler : register(s0);
SamplerComparisonState gShadowSampler : register(s2);

struct VSInput {
    float3 positionOS : POSITION;
    float3 normalOS : NORMAL;
    float3 tangentOS : TANGENT;
    float2 uv0 : TEXCOORD0;
};

struct VSOutput {
    float4 positionCS : SV_POSITION;
    float3 positionWS : POSITION1;
    float3 normalWS : NORMAL0;
    float3 tangentWS : TANGENT0;
    float2 uv0 : TEXCOORD0;
    float4 shadowPosition : TEXCOORD1;
};

VSOutput VSMain(VSInput input) {
    VSOutput output;

    const float4 positionWS = mul(float4(input.positionOS, 1.0f), gWorld);
    const float4 positionVS = mul(positionWS, gView);

    output.positionCS = mul(positionVS, gProjection);
    output.positionWS = positionWS.xyz;

    // 法线必须乘 world 的逆转置矩阵。
    const float3 normalWS = SafeNormalize(
        mul(float4(input.normalOS, 0.0f), gWorldInvTranspose).xyz,
        float3(0.0f, 1.0f, 0.0f));

    // 切线按普通方向乘 world，再对法线做 Gram-Schmidt 正交化。
    float3 tangentWS = mul(float4(input.tangentOS, 0.0f), gWorld).xyz;
    tangentWS -= normalWS * dot(normalWS, tangentWS);
    tangentWS = SafeNormalize(tangentWS, BuildFallbackTangent(normalWS));

    output.normalWS = normalWS;
    output.tangentWS = tangentWS;
    output.uv0 = input.uv0;
    //沿法线外扩后，获取shadowPos
    float4 posWS_bigger =float4( (positionWS.xyz + normalWS.xyz * 0.1f),1.0f);
//    output.shadowPosition = mul(positionWS, gLightViewProjection);
    output.shadowPosition = mul(posWS_bigger, gLightViewProjection);
    return output;
}

float SampleDirectionalShadow(float4 shadowPosition, float3 n, float3 l) {
    float finalRet=1.0f;
    if (shadowPosition.w > 0.0f){
        const float3 ndc = shadowPosition.xyz / shadowPosition.w;
        const float2 uv = float2(ndc.x * 0.5f + 0.5f, -ndc.y * 0.5f + 0.5f);
        if (!any(uv < 0.0f) && !any(uv > 1.0f) &&! (ndc.z <= 0.0f) && !(ndc.z >= 1.0f)) {
            const float bias = max(
                gShadowParams.z,
                gShadowParams.w * (1.0f - saturate(dot(n, l))));

            float visibility = 0.0f;
            [unroll]
            for (int y = -1; y <= 1; ++y) {
                [unroll]
                for (int x = -1; x <= 1; ++x) {
                    const float2 offset = float2(x, y) * gShadowParams.xy;
                    visibility += gShadowMap.SampleCmpLevelZero(
                        gShadowSampler, uv + offset, ndc.z - bias);
                }
            }
            finalRet = visibility/9.0f;
        }
    }
    return finalRet;

    /* 编译器 warn -> error 多return路径
    if (shadowPosition.w <= 0.0f) return 1.0f;

    const float3 ndc = shadowPosition.xyz / shadowPosition.w;
    const float2 uv = float2(ndc.x * 0.5f + 0.5f, -ndc.y * 0.5f + 0.5f);

    if (any(uv < 0.0f) || any(uv > 1.0f) || ndc.z <= 0.0f || ndc.z >= 1.0f) {
        return 1.0f;
    }

    const float bias = max(
        gShadowParams.z,
        gShadowParams.w * (1.0f - saturate(dot(n, l))));

    float visibility = 0.0f;
    [unroll]
    for (int y = -1; y <= 1; ++y) {
        [unroll]
        for (int x = -1; x <= 1; ++x) {
            const float2 offset = float2(x, y) * gShadowParams.xy;
            visibility += gShadowMap.SampleCmpLevelZero(
                gShadowSampler, uv + offset, ndc.z - bias);
        }
    }
    return visibility / 9.0f;
    */
}

float4 PSMain(VSOutput input) : SV_TARGET {
    const float4 baseSample = gBaseColorTex.Sample(gMaterialSampler, input.uv0);
    const float4 baseColor = baseSample * gBaseColorFactor;

    const float3 orm = gOrmTex.Sample(gMaterialSampler, input.uv0).rgb;
    const float ao = lerp(1.0f, orm.r, saturate(gPbrParams.w));
    const float roughness = clamp(orm.g * gPbrParams.y, 0.045f, 1.0f);
    const float metallic = saturate(orm.b * gPbrParams.x);

    const float3 normalTS = DecodeNormalMap(gNormalTex.Sample(gMaterialSampler, input.uv0).xyz, gPbrParams.z);

    const float3 geometricNormal = SafeNormalize(input.normalWS, float3(0.0f, 1.0f, 0.0f));
    const float3 tangent = SafeNormalize(input.tangentWS, BuildFallbackTangent(geometricNormal));

    // 当前 VertexData 只有 float3 tangent，暂按 handedness=+1 构造 B。
    // 生产版本应把 tangent 改为 float4，并乘 tangent.w 处理镜像 UV。
    const float3 bitangent = SafeNormalize( cross(geometricNormal, tangent),
                                 cross(geometricNormal, BuildFallbackTangent(geometricNormal)));
    const float3x3 tangentToWorld = float3x3(tangent, bitangent, geometricNormal);
    const float3 n = SafeNormalize(mul(normalTS, tangentToWorld), geometricNormal);

    const float3 v = SafeNormalize(gCameraPosition.xyz - input.positionWS, geometricNormal);
    // Scene.lightDirection 表示光线传播方向，因此表面指向光源的 L 取反。
    const float3 l = SafeNormalize(-gLightDirection.xyz, float3(0.0f, 1.0f, 0.0f));
    const float3 h = SafeNormalize(v + l, n);

    const float nDotL = saturate(dot(n, l));
    const float nDotV = saturate(dot(n, v));
    const float hDotV = saturate(dot(h, v));

    const float3 f0 = lerp(float3(0.04f, 0.04f, 0.04f), baseColor.rgb, metallic);
    const float d = DistributionGGX(n, h, roughness);
    const float g = GeometrySmith(n, v, l, roughness);
    const float3 f = FresnelSchlick(hDotV, f0);

    const float3 specular = (d * g * f) / max(4.0f * nDotV * nDotL, PBR_EPSILON);

    const float3 kS = f;
    const float3 kD = (1.0f - kS) * (1.0f - metallic);
    const float3 diffuse = kD * baseColor.rgb / PBR_PI;

    const float visibility = SampleDirectionalShadow(input.shadowPosition, n, l);
    const float3 radiance = gLightColorIntensity.rgb * gLightColorIntensity.w;
    const float3 directLighting = (diffuse + specular) * radiance * nDotL * visibility;

    // [PBR-TEMP] 尚未接入 IBL 的低成本环境近似。
    const float3 ambientBase = (baseColor.rgb * (1.0f - metallic) + f0 * 0.25f) * ao;
    const float3 ambientLighting = ambientBase * gAmbientColorIntensity.rgb * gAmbientColorIntensity.w;

    const float3 emissive = gEmissiveTex.Sample(gMaterialSampler, input.uv0).rgb * gEmissiveFactor.rgb;

    float3 color = directLighting + ambientLighting + emissive;
    color *= max(gRenderParams.x, 0.0f);
    color = ToneMapACES(color);

    // Graphic 当前使用 UNORM back buffer，因此在 PS 中完成线性 -> 显示编码。
    color = pow(max(color, 0.0f), 1.0f / 2.2f);
    return float4(color, baseColor.a);
}