#version 450

const int MAX_LIGHTS = 4;

struct Light {
    vec4 color;       // rgb = 颜色，a = 强度
    vec4 direction;   // 方向光：xyz = 方向（已归一）
    vec4 position;    // 点光：xyz = 世界位置
    vec4 params;      // x = 类型(0=方向,1=点)，y = 衰减半径
};

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 = 5) in vec3 fragTangent;

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;
    vec4 cameraPos;
} ubo;

// 材质纹理槽（binding = Material::Slot）
layout(set = 1, binding = 0) uniform sampler2D albedoMap;
layout(set = 1, binding = 1) uniform sampler2D normalMap;
layout(set = 1, binding = 2) uniform sampler2D specularMap;
layout(set = 1, binding = 3) uniform sampler2D emissiveMap;

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;
}

// 法线贴图 → 世界空间法线（TBN 构建）。
// 无切线数据（零向量）或平法线贴图时退回几何法线。
vec3 MappedNormal(vec3 geomN, vec3 T, vec2 uv) {
    if (dot(T, T) < 1e-4) return geomN;

    vec3 N = geomN;
    vec3 nT = normalize(T - N * dot(N, T));   // 正交化
    vec3 B = cross(N, nT);
    vec3 sampled = texture(normalMap, uv).xyz * 2.0 - 1.0;
    if (dot(sampled, sampled) < 1e-4) return geomN;

    return normalize(nT * sampled.x + B * sampled.y + N * sampled.z);
}

void main() {
    vec3 albedo = texture(albedoMap, fragUV).rgb;
    float specIntensity = texture(specularMap, fragUV).r;
    vec3 emissive = texture(emissiveMap, fragUV).rgb;

    vec3 N = MappedNormal(normalize(fragNormal), fragTangent, fragUV);
    vec3 V = normalize(ubo.cameraPos.xyz - fragWorldPos);

    vec3 result = ubo.ambientColor.rgb * albedo;

    for (int i = 0; i < MAX_LIGHTS; ++i) {
        // 按类型分支：方向光无衰减；点光方向指向光源 + 平方距离衰减
        vec3 L;
        float atten = 1.0;
        if (ubo.lights[i].params.x > 0.5) {
            vec3 toLight = ubo.lights[i].position.xyz - fragWorldPos;
            float d = length(toLight);
            if (d >= ubo.lights[i].params.y) continue;   // 衰减半径外，跳过
            L = toLight / d;
            float a = clamp(1.0 - (d * d) /
                            (ubo.lights[i].params.y * ubo.lights[i].params.y),
                            0.0, 1.0);
            atten = a * a;   // 平方衰减，边缘更柔和
        } else {
            L = normalize(-ubo.lights[i].direction.xyz);
        }

        float diff = max(dot(N, L), 0.0);
        // Blinn-Phong 高光（半程向量；specular 图做强度系数）
        vec3 H = normalize(L + V);
        float spec = pow(max(dot(N, H), 0.0), 32.0) * specIntensity;

        // 阴影只来自第 0 盏方向光（CSM）
        float shade = (i == 0 && ubo.lights[0].params.x < 0.5)
            ? ShadowFactor(N, fragViewDepth) : 1.0;

        result += ubo.lights[i].color.rgb * ubo.lights[i].color.a *
                  (albedo * diff + vec3(spec)) * atten * shade;
    }

    // 自发光不经光照/阴影直接叠加
    result += emissive;

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