#version 450

const int MAX_LIGHTS = 4;

struct Light {
    vec4 direction;
    vec4 color;
};

layout(location = 0) in vec2 fragUV;
layout(location = 1) in vec3 fragNormal;
layout(location = 2) in vec3 fragWorldPos;
layout(location = 3) in float fragViewDepth;
layout(location = 4) in vec2 fragVelocity;

layout(location = 0) out vec4 outColor;
layout(location = 1) out vec2 outVelocity;   // TAA velocity 附件

layout(set = 0, binding = 0) uniform UniformBufferObject {
    mat4 view;
    mat4 proj;
    mat4 lightSpace[3];     // CSM 各级联光空间 VP
    vec4 cascadeSplits;     // xyz = 级联 0..2 远边界（视图深度）
    vec4 ambientColor;
    Light lights[MAX_LIGHTS];
    mat4 projNoJitter;
    mat4 prevView;
    mat4 prevProj;
} ubo;

layout(set = 1, binding = 0) uniform sampler2D texSampler;
layout(set = 2, binding = 0) uniform sampler2D shadowMap0;   // CSM 级联 0（最近）
layout(set = 2, binding = 1) uniform sampler2D shadowMap1;   // 级联 1
layout(set = 2, binding = 2) uniform sampler2D shadowMap2;   // 级联 2（最远）

// 第 0 盏灯（KeyLight）的 CSM 阴影因子：0.3 影中 / 1.0 受光。
// 级联按视图深度选择；采样器均 clamp-to-border 白色：
// 光视锥外自动 1.0（无阴影），省出界分支。
float ShadowFactor(vec3 N, float viewDepth) {
    // 级联选择：第一个 viewDepth < split 的级联
    int cascade = 2;
    for (int i = 0; i < 3; ++i) {
        if (viewDepth < ubo.cascadeSplits[i]) { cascade = i; break; }
    }

    vec4 ls = ubo.lightSpace[cascade] * vec4(fragWorldPos, 1.0);
    vec3 proj = ls.xyz / ls.w;
    // NDC [-1,1] → UV [0,1]（ortho 的 w=1）
    vec2 uv = proj.xy * 0.5 + 0.5;

    float closest;
    if (cascade == 0)      closest = texture(shadowMap0, uv).r;
    else if (cascade == 1) closest = texture(shadowMap1, uv).r;
    else                   closest = texture(shadowMap2, uv).r;

    // 斜率缩放 bias：掠射面用更大的 bias 压 acne
    vec3 L = normalize(-ubo.lights[0].direction.xyz);
    float bias = max(0.004 * (1.0 - dot(N, L)), 0.0012);

    return (proj.z - bias > closest) ? 0.3 : 1.0;
}

void main() {
    vec3 albedo = texture(texSampler, fragUV).rgb;
    vec3 N = normalize(fragNormal);

    vec3 result = ubo.ambientColor.rgb * albedo;

    for (int i = 0; i < MAX_LIGHTS; ++i) {
        vec3 L = normalize(-ubo.lights[i].direction.xyz);
        float diff = max(dot(N, L), 0.0);
        float shade = (i == 0) ? ShadowFactor(N, fragViewDepth) : 1.0;
        result += ubo.lights[i].color.rgb * ubo.lights[i].color.a * diff * albedo * shade;
    }

    outColor = vec4(result, 1.0);
    outVelocity = fragVelocity;
}