#version 450

#extension GL_GOOGLE_include_directive : require
#include "input_structures.glsl"

layout (location = 0) in vec3 inNormal;
layout (location = 1) in vec3 inColor;
layout (location = 2) in vec2 inUV;
layout (location = 3) in vec3 inWorldPos;
layout (location = 4) in vec4 inTangent;

layout (location = 0) out vec4 outFragColor;

const float PI = 3.14159265359;

// IBL bindings (set=3): specular equirect 2D + BRDF LUT + SH UBO
layout(set=3, binding=0) uniform sampler2D iblSpec2D;
layout(set=3, binding=1) uniform sampler2D iblBRDF;
layout(std140, set=3, binding=2) uniform IBL_SH { vec4 sh[9]; } iblSH;

vec3 sh_eval_irradiance(vec3 n)
{
    float x=n.x, y=n.y, z=n.z;
    const float c0=0.2820947918; const float c1=0.4886025119; const float c2=1.0925484306; const float c3=0.3153915653; const float c4=0.5462742153;
    float Y[9];
    Y[0]=c0; Y[1]=c1*y; Y[2]=c1*z; Y[3]=c1*x; Y[4]=c2*x*y; Y[5]=c2*y*z; Y[6]=c3*(3.0*z*z-1.0); Y[7]=c2*x*z; Y[8]=c4*(x*x-y*y);
    vec3 r=vec3(0.0); for (int i=0;i<9;++i) r += iblSH.sh[i].rgb * Y[i]; return r;
}

vec2 dir_to_equirect(vec3 d)
{
    d = normalize(d);
    float phi = atan(d.z, d.x);
    float theta = acos(clamp(d.y, -1.0, 1.0));
    return vec2(phi * (0.15915494309) + 0.5, theta * (0.31830988618));
}

vec3 fresnelSchlick(float cosTheta, vec3 F0)
{
    return F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0);
}

float DistributionGGX(vec3 N, vec3 H, float roughness)
{
    float a      = roughness * roughness;
    float a2     = a * a;
    float NdotH  = max(dot(N, H), 0.0);
    float NdotH2 = NdotH * NdotH;

    float num   = a2;
    float denom = (NdotH2 * (a2 - 1.0) + 1.0);
    denom = PI * denom * denom;

    return num / max(denom, 0.001);
}

float GeometrySchlickGGX(float NdotV, float roughness)
{
    float r = (roughness + 1.0);
    float k = (r * r) / 8.0;

    float denom = NdotV * (1.0 - k) + k;
    return NdotV / denom;
}

float GeometrySmith(vec3 N, vec3 V, vec3 L, float roughness)
{
    float ggx2 = GeometrySchlickGGX(max(dot(N, V), 0.0), roughness);
    float ggx1 = GeometrySchlickGGX(max(dot(N, L), 0.0), roughness);
    return ggx1 * ggx2;
}

void main()
{
    // Base color with material factor and texture
    vec4 baseTex = texture(colorTex, inUV);
    vec3 albedo = inColor * baseTex.rgb * materialData.colorFactors.rgb;
    // glTF: metallicRoughnessTexture uses G=roughness, B=metallic
    vec2 mrTex = texture(metalRoughTex, inUV).gb;
    float roughness = clamp(mrTex.x * materialData.metal_rough_factors.y, 0.04, 1.0);
    float metallic  = clamp(mrTex.y * materialData.metal_rough_factors.x, 0.0, 1.0);

    // Normal mapping path for forward/transparent pipeline
    // Expect UNORM normal map; support BC5 (RG) by reconstructing Z from XY.
    vec2 enc = texture(normalMap, inUV).xy * 2.0 - 1.0;
    float normalScale = max(materialData.extra[0].x, 0.0);
    enc *= normalScale;
    float z2 = 1.0 - dot(enc, enc);
    float nz = z2 > 0.0 ? sqrt(z2) : 0.0;
    vec3 Nm = vec3(enc, nz);
    vec3 Nn = normalize(inNormal);
    vec3 T = normalize(inTangent.xyz);
    vec3 B = normalize(cross(Nn, T)) * inTangent.w;
    vec3 N = normalize(T * Nm.x + B * Nm.y + Nn * Nm.z);
    vec3 camPos = vec3(inverse(sceneData.view)[3]);
    vec3 V = normalize(camPos - inWorldPos);
    vec3 L = normalize(-sceneData.sunlightDirection.xyz);
    vec3 H = normalize(V + L);

    vec3 F0 = mix(vec3(0.04), albedo, metallic);
    vec3 F  = fresnelSchlick(max(dot(H, V), 0.0), F0);
    float NDF = DistributionGGX(N, H, roughness);
    float G   = GeometrySmith(N, V, L, roughness);

    vec3 numerator    = NDF * G * F;
    float denom       = 4.0 * max(dot(N, V), 0.0) * max(dot(N, L), 0.0);
    vec3 specular     = numerator / max(denom, 0.001);

    vec3 kS = F;
    vec3 kD = vec3(1.0) - kS;
    kD *= 1.0 - metallic;

    float NdotL = max(dot(N, L), 0.0);
    vec3 irradiance = sceneData.sunlightColor.rgb * sceneData.sunlightColor.a * NdotL;

    vec3 color = (kD * albedo / PI + specular) * irradiance;

    // IBL: specular from equirect 2D mips; diffuse from SH
    vec3 R = reflect(-V, N);
    float levels = float(textureQueryLevels(iblSpec2D));
    float lod = clamp(roughness * max(levels - 1.0, 0.0), 0.0, max(levels - 1.0, 0.0));
    vec2 uv = dir_to_equirect(R);
    vec3 prefiltered = textureLod(iblSpec2D, uv, lod).rgb;
    vec2 brdf = texture(iblBRDF, vec2(max(dot(N, V), 0.0), roughness)).rg;
    vec3 specIBL = prefiltered * (F0 * brdf.x + brdf.y);
    vec3 diffIBL = (1.0 - metallic) * albedo * sh_eval_irradiance(N);
    color += diffIBL + specIBL;

    // Alpha from baseColor texture and factor (glTF spec)
    float alpha = clamp(baseTex.a * materialData.colorFactors.a, 0.0, 1.0);
    outFragColor = vec4(color, alpha);
}