#version 450

layout(set = 0, binding = 0) uniform PerFrame {
    mat4 view;
    mat4 projection;
    mat4 lightSpaceMatrix;
    vec4 lightDir;
    vec4 lightColor;
    vec4 ambientColor;
    vec4 viewPos;
    vec4 fogColor;
    vec4 fogParams;
    vec4 shadowParams;
};

layout(set = 1, binding = 0) uniform sampler2D uTexture;

layout(set = 1, binding = 2) uniform M2Material {
    int hasTexture;
    int alphaTest;
    int colorKeyBlack;
    float colorKeyThreshold;
    int unlit;
    int blendMode;
    float fadeAlpha;
    float interiorDarken;
    float specularIntensity;
};

layout(set = 0, binding = 1) uniform sampler2DShadow uShadowMap;

layout(location = 0) in vec3 FragPos;
layout(location = 1) in vec3 Normal;
layout(location = 2) in vec2 TexCoord;
layout(location = 3) flat in vec3 InstanceOrigin;
layout(location = 4) in float ModelHeight;

layout(location = 0) out vec4 outColor;

// 4x4 Bayer dither matrix (normalized to 0..1)
float bayerDither4x4(ivec2 p) {
    int idx = (p.x & 3) + (p.y & 3) * 4;
    float m[16] = float[16](
         0.0/16.0,  8.0/16.0,  2.0/16.0, 10.0/16.0,
        12.0/16.0,  4.0/16.0, 14.0/16.0,  6.0/16.0,
         3.0/16.0, 11.0/16.0,  1.0/16.0,  9.0/16.0,
        15.0/16.0,  7.0/16.0, 13.0/16.0,  5.0/16.0
    );
    return m[idx];
}

void main() {
    vec4 texColor = hasTexture != 0 ? texture(uTexture, TexCoord) : vec4(1.0);

    bool isFoliage = (alphaTest == 2);

    float alphaCutoff = 0.5;
    if (alphaTest == 2) {
        // Vegetation cutout: lower threshold to preserve leaf coverage at grazing angles.
        alphaCutoff = 0.33;
    } else if (alphaTest == 3) {
        // Ground detail clutter (grass/small cards) needs softer clipping.
        alphaCutoff = 0.20;
    } else if (alphaTest != 0) {
        alphaCutoff = 0.35;
    }
    if (alphaTest == 2) {
        float alpha = texColor.a;
        float softBand = 0.15;
        if (alpha < (alphaCutoff - softBand)) discard;
        if (alpha < alphaCutoff) {
            ivec2 p = ivec2(gl_FragCoord.xy);
            float threshold = bayerDither4x4(p);
            float keep = clamp((alpha - (alphaCutoff - softBand)) / softBand, 0.0, 1.0);
            if (threshold > keep) discard;
        }
    } else if (alphaTest != 0 && texColor.a < alphaCutoff) {
        discard;
    }
    if (colorKeyBlack != 0) {
        float lum = dot(texColor.rgb, vec3(0.299, 0.587, 0.114));
        if (lum < colorKeyThreshold) discard;
    }
    if (blendMode == 1 && texColor.a < 0.004) discard;

    // Per-instance color variation (foliage only)
    if (isFoliage) {
        float hash = fract(sin(dot(InstanceOrigin.xy, vec2(127.1, 311.7))) * 43758.5453);
        float hueShiftR = 1.0 + (hash - 0.5) * 0.16;       // ±8% red
        float hueShiftB = 1.0 + (fract(hash * 7.13) - 0.5) * 0.16; // ±8% blue
        float brightness = 0.85 + hash * 0.30;               // 85–115%
        texColor.rgb *= vec3(hueShiftR, 1.0, hueShiftB) * brightness;
    }

    vec3 norm = normalize(Normal);
    bool foliageTwoSided = (alphaTest == 2);
    if (!foliageTwoSided && !gl_FrontFacing) norm = -norm;

    // Detail normal perturbation (foliage only) — UV-based only so wind doesn't cause flicker
    if (isFoliage) {
        float nx = sin(TexCoord.x * 12.0 + TexCoord.y * 5.3) * 0.10;
        float ny = sin(TexCoord.y * 14.0 + TexCoord.x * 4.7) * 0.10;
        norm = normalize(norm + vec3(nx, ny, 0.0));
    }

    vec3 ldir = normalize(-lightDir.xyz);
        float nDotL = dot(norm, ldir);
        float diff = foliageTwoSided ? abs(nDotL) : max(nDotL, 0.0);

    vec3 result;
    if (unlit != 0) {
        result = texColor.rgb;
    } else {
        vec3 viewDir = normalize(viewPos.xyz - FragPos);

        // Foliage: no specular, no shadow map — both flicker on swaying thin cards
        float spec = 0.0;
        float shadow = 1.0;
        if (isFoliage) {
            // Use a fixed gentle shadow from the shadow system strength
            if (shadowParams.x > 0.5) {
                shadow = mix(1.0, 0.75, shadowParams.y);
            }
        } else {
            vec3 halfDir = normalize(ldir + viewDir);
            spec = pow(max(dot(norm, halfDir), 0.0), 32.0) * specularIntensity;

            if (shadowParams.x > 0.5) {
                vec4 lsPos = lightSpaceMatrix * vec4(FragPos, 1.0);
                vec3 proj = lsPos.xyz / lsPos.w;
                proj.xy = proj.xy * 0.5 + 0.5;
                if (proj.x >= 0.0 && proj.x <= 1.0 &&
                    proj.y >= 0.0 && proj.y <= 1.0 &&
                    proj.z >= 0.0 && proj.z <= 1.0) {
                    float bias = max(0.0005 * (1.0 - abs(dot(norm, ldir))), 0.00005);
                    shadow = texture(uShadowMap, vec3(proj.xy, proj.z - bias));
                }
                shadow = mix(1.0, shadow, shadowParams.y);
            }
        }

        // Leaf subsurface scattering (foliage only) — uses stable normal, no FragPos dependency
        vec3 sss = vec3(0.0);
        if (isFoliage) {
            float backLit = max(-nDotL, 0.0);
            float viewDotLight = max(dot(viewDir, -ldir), 0.0);
            float sssAmount = backLit * pow(viewDotLight, 4.0) * 0.35 * texColor.a;
            sss = sssAmount * vec3(1.0, 0.9, 0.5) * lightColor.rgb;
        }

        result = ambientColor.rgb * texColor.rgb
               + shadow * (diff * lightColor.rgb * texColor.rgb + spec * lightColor.rgb)
               + sss;

        if (interiorDarken > 0.0) {
            result *= mix(1.0, 0.5, interiorDarken);
        }
    }

    // Canopy ambient occlusion (foliage only)
    if (isFoliage) {
        float normalizedHeight = clamp(ModelHeight / 18.0, 0.0, 1.0);
        float aoFactor = mix(0.55, 1.0, smoothstep(0.0, 0.6, normalizedHeight));
        result *= aoFactor;
    }

    float dist = length(viewPos.xyz - FragPos);
    float fogFactor = clamp((fogParams.y - dist) / (fogParams.y - fogParams.x), 0.0, 1.0);
    result = mix(fogColor.rgb, result, fogFactor);

    float outAlpha = texColor.a * fadeAlpha;
    // Cutout materials should not remain partially transparent after discard,
    // otherwise foliage cards look view-dependent.
    if (alphaTest != 0 || colorKeyBlack != 0) {
        outAlpha = fadeAlpha;
    }
    // Foliage cutout should stay opaque after alpha discard to avoid
    // view-angle translucency artifacts.
    if (alphaTest == 2 || alphaTest == 3) {
        outAlpha = 1.0 * fadeAlpha;
    }
    outColor = vec4(result, outAlpha);
}