#version 450 core
// Purpose: Game scene lit-model fragment shader; handles CanYouDrone material, lighting, fog, and shadows.

in vec3 vWorldPos;
in vec3 vWorldNormal;
in vec2 vUv;

uniform sampler2D uBaseColorTex;
uniform vec4 uBaseColorFactor;
uniform int uHasTexture;

uniform vec3 uCameraPos;

uniform vec3 uDirLightDirection;
uniform vec3 uDirLightColor;
uniform float uDirLightIntensity;
uniform vec3 uAmbientLightColor;
uniform float uAmbientLightIntensity;

const int PointLightCount = 16;

uniform int uPointLightCount;
uniform vec3 uPointLightPositions[PointLightCount];
uniform vec3 uPointLightColors[PointLightCount];
uniform float uPointLightIntensities[PointLightCount];
uniform float uPointLightRanges[PointLightCount];
uniform samplerCube uPointShadowMaps[PointLightCount];
uniform int uPointShadowEnabled;
uniform float uPointShadowFar[PointLightCount];
uniform float uPointShadowBias[PointLightCount];
uniform float uPointShadowNormalOffset[PointLightCount];
uniform int uPointShadowEnabledByLight[PointLightCount];

uniform mat4 uLightVP[3];
uniform sampler2D uShadowMaps[3];
uniform float uShadowDepthRange[3];
uniform float uShadowBiasMin[3];
uniform float uShadowBiasMax[3];
uniform float uShadowNormalOffset[3];

out vec4 FragColor;

vec4 LightClipPos(int cascadeIndex, vec3 worldPos)
{
    if (cascadeIndex == 0) return uLightVP[0] * vec4(worldPos, 1.0);
    if (cascadeIndex == 1) return uLightVP[1] * vec4(worldPos, 1.0);
    return uLightVP[2] * vec4(worldPos, 1.0);
}

float ShadowDepth(int cascadeIndex, vec2 uv)
{
    if (cascadeIndex == 0) return texture(uShadowMaps[0], uv).r;
    if (cascadeIndex == 1) return texture(uShadowMaps[1], uv).r;
    return texture(uShadowMaps[2], uv).r;
}

vec2 ShadowTexelSize(int cascadeIndex)
{
    if (cascadeIndex == 0) return 1.0 / vec2(textureSize(uShadowMaps[0], 0));
    if (cascadeIndex == 1) return 1.0 / vec2(textureSize(uShadowMaps[1], 0));
    return 1.0 / vec2(textureSize(uShadowMaps[2], 0));
}

vec3 ProjectShadowPos(int cascadeIndex, vec3 worldPos)
{
    vec4 clip = LightClipPos(cascadeIndex, worldPos);
    vec3 proj = clip.xyz / clip.w;
    return proj * 0.5 + 0.5;
}

float SampleShadowCascade(int cascadeIndex, vec3 proj, vec3 normal, vec3 lightDir)
{
    float currentDepth = proj.z;
    float slope = 1.0 - clamp(dot(normal, lightDir), 0.0, 1.0);
    float biasWorld = mix(uShadowBiasMin[cascadeIndex], uShadowBiasMax[cascadeIndex], slope);
    float bias = biasWorld / max(uShadowDepthRange[cascadeIndex], 0.001);

    vec2 texelSize = ShadowTexelSize(cascadeIndex);

    float lit = 0.0;
    float weightSum = 0.0;
    for (int y = -2; y <= 2; ++y)
    {
        for (int x = -2; x <= 2; ++x)
        {
            float weight = (3.0 - abs(float(x))) * (3.0 - abs(float(y)));
            float sampledDepth = ShadowDepth(cascadeIndex, proj.xy + vec2(x, y) * texelSize);
            lit += (((currentDepth - bias) <= sampledDepth) ? 1.0 : 0.0) * weight;
            weightSum += weight;
        }
    }

    return lit / weightSum;
}

bool IsInsideCascade(vec3 proj)
{
    return
        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 CascadeEdgeFade(vec3 proj)
{
    float edge = min(min(proj.x, 1.0 - proj.x), min(proj.y, 1.0 - proj.y));
    return smoothstep(0.04, 0.12, edge);
}

float ComputeShadowPCF(vec3 normal, vec3 lightDir)
{
    for (int i = 0; i < 3; ++i)
    {
        float slope = 1.0 - clamp(dot(normal, lightDir), 0.0, 1.0);
        vec3 offsetWorldPos = vWorldPos + normal * (uShadowNormalOffset[i] * (0.35 + slope));
        vec3 proj = ProjectShadowPos(i, offsetWorldPos);

        if (IsInsideCascade(proj))
        {
            float cascadeShadow = SampleShadowCascade(i, proj, normal, lightDir);

            if (i < 2)
            {
                vec3 nextOffsetWorldPos = vWorldPos + normal * (uShadowNormalOffset[i + 1] * (0.35 + slope));
                vec3 nextProj = ProjectShadowPos(i + 1, nextOffsetWorldPos);

                if (IsInsideCascade(nextProj))
                {
                    float nextShadow = SampleShadowCascade(i + 1, nextProj, normal, lightDir);
                    cascadeShadow = mix(nextShadow, cascadeShadow, CascadeEdgeFade(proj));
                }
            }

            return cascadeShadow;
        }
    }

    return 1.0;
}

float PointShadowDepth(int lightIndex, vec3 sampleDir)
{
    if (lightIndex == 0) return texture(uPointShadowMaps[0], sampleDir).r;
    if (lightIndex == 1) return texture(uPointShadowMaps[1], sampleDir).r;
    if (lightIndex == 2) return texture(uPointShadowMaps[2], sampleDir).r;
    if (lightIndex == 3) return texture(uPointShadowMaps[3], sampleDir).r;
    return texture(uPointShadowMaps[4], sampleDir).r;
}

float ComputePointShadow(int lightIndex, vec3 normal, vec3 lightDir)
{
    if (uPointShadowEnabled == 0 || uPointShadowEnabledByLight[lightIndex] == 0)
    {
        return 1.0;
    }

    float slope = 1.0 - clamp(dot(normal, lightDir), 0.0, 1.0);
    vec3 offsetWorldPos = vWorldPos + normal * (uPointShadowNormalOffset[lightIndex] * (0.35 + slope));
    vec3 fromLight = offsetWorldPos - uPointLightPositions[lightIndex];
    float currentDepth = length(fromLight);

    if (currentDepth >= uPointShadowFar[lightIndex])
    {
        return 1.0;
    }

    vec2 sampleOffsets[9] = vec2[](
        vec2( 0.0,  0.0),
        vec2( 1.0,  0.0), vec2(-1.0,  0.0),
        vec2( 0.0,  1.0), vec2( 0.0, -1.0),
        vec2( 0.7,  0.7), vec2(-0.7,  0.7),
        vec2( 0.7, -0.7), vec2(-0.7, -0.7)
    );

    vec3 sampleDir = fromLight / max(currentDepth, 0.001);
    vec3 basisSeed = (abs(sampleDir.y) < 0.98) ? vec3(0.0, 1.0, 0.0) : vec3(1.0, 0.0, 0.0);
    vec3 tangentRight = normalize(cross(basisSeed, sampleDir));
    vec3 tangentUp = normalize(cross(sampleDir, tangentRight));

    float diskRadius = mix(0.0015, 0.012, clamp(currentDepth / max(uPointShadowFar[lightIndex], 0.001), 0.0, 1.0));
    float bias = uPointShadowBias[lightIndex] * (0.35 + slope);

    float lit = 0.0;
    for (int i = 0; i < 9; ++i)
    {
        vec3 lookupDir = normalize(sampleDir +
            tangentRight * (sampleOffsets[i].x * diskRadius) +
            tangentUp * (sampleOffsets[i].y * diskRadius));
        float closestDepth = PointShadowDepth(lightIndex, lookupDir) * uPointShadowFar[lightIndex];
        lit += ((currentDepth - bias) <= closestDepth) ? 1.0 : 0.0;
    }

    return lit / 9.0;
}

void main()
{
    vec4 albedo4 = uBaseColorFactor;
    if (uHasTexture != 0)
    albedo4 *= texture(uBaseColorTex, vUv);

    vec3 albedo = albedo4.rgb;
    vec3 N = normalize(vWorldNormal);
    vec3 V = normalize(uCameraPos - vWorldPos);

    vec3 ambient = albedo * uAmbientLightColor * uAmbientLightIntensity;

    vec3 Ld = normalize(-uDirLightDirection);
    float ndl = max(dot(N, Ld), 0.0);
    float shadow = ComputeShadowPCF(N, Ld);

    vec3 dirDiffuse = albedo * uDirLightColor * (uDirLightIntensity * ndl * shadow);

    vec3 Hd = normalize(Ld + V);
    float ndh = max(dot(N, Hd), 0.0);
    float dirSpec = pow(ndh, 32.0) * step(0.0001, ndl);
    vec3 dirSpecular = uDirLightColor * (0.15 * uDirLightIntensity * dirSpec * shadow);

    vec3 pointDiffuse = vec3(0.0);
    vec3 pointSpecular = vec3(0.0);
    for (int i = 0; i < PointLightCount; ++i)
    {
        if (i >= uPointLightCount)
        {
            break;
        }

        vec3 LpVec = uPointLightPositions[i] - vWorldPos;
        float dist = length(LpVec);
        vec3 Lp = (dist > 0.0001) ? (LpVec / dist) : vec3(0.0, 1.0, 0.0);

        float atten = clamp(1.0 - dist / max(uPointLightRanges[i], 0.001), 0.0, 1.0);
        atten *= atten;

        float npl = max(dot(N, Lp), 0.0);
        if (uPointLightIntensities[i] <= 0.0 || atten <= 0.0 || npl <= 0.0)
        {
            continue;
        }

        float pointShadow = ComputePointShadow(i, N, Lp);
        pointDiffuse += albedo * uPointLightColors[i] * (uPointLightIntensities[i] * npl * atten * pointShadow);

        vec3 Hp = normalize(Lp + V);
        float nph = max(dot(N, Hp), 0.0);
        float pointSpec = pow(nph, 24.0) * step(0.0001, npl);
        pointSpecular += uPointLightColors[i] * (0.12 * uPointLightIntensities[i] * pointSpec * atten * pointShadow);
    }

    vec3 color = ambient + dirDiffuse + dirSpecular + pointDiffuse + pointSpecular;
    FragColor = vec4(color, albedo4.a);
}
