#version 330 core

in  vec3 v_normal;
in  vec2 v_tex;
in  vec3 v_position; // view-space position

out vec4 frag_color;

// Material inputs
uniform sampler2D u_base_color_map;      // sRGB
uniform int       u_has_base_color_map;
uniform vec4      u_base_color_factor;

uniform sampler2D u_metallic_roughness_map; // linear, (r=occlusion in glTF separate, g=roughness, b=metallic)
uniform int       u_has_metallic_roughness_map;
uniform float     u_metallic_factor;
uniform float     u_roughness_factor;

uniform sampler2D u_occlusion_map; // linear (r)
uniform int       u_has_occlusion_map;

uniform sampler2D u_emissive_map; // sRGB
uniform int       u_has_emissive_map;
uniform vec3      u_emissive_factor;

// Directional light (in view space)
uniform vec3  u_dir_light_dir;       // direction from light towards the scene (L)
uniform vec3  u_dir_light_color;
uniform float u_dir_light_intensity;

// Point lights (in view space)
#define MAX_POINT_LIGHTS 4
uniform int   u_point_light_count;
uniform vec3  u_point_light_pos[MAX_POINT_LIGHTS];
uniform vec3  u_point_light_color[MAX_POINT_LIGHTS];
uniform float u_point_light_intensity[MAX_POINT_LIGHTS];
uniform float u_point_light_range[MAX_POINT_LIGHTS];

const float PI = 3.14159265359;

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

float distribution_ggx(float NdotH, float alpha) {
    float a2 = alpha * alpha;
    float denom = (NdotH * NdotH) * (a2 - 1.0) + 1.0;
    return a2 / (PI * denom * denom + 1e-6);
}

float geometry_schlick_ggx(float NdotV, float k) {
    return NdotV / (NdotV * (1.0 - k) + k + 1e-6);
}

float geometry_smith(float NdotV, float NdotL, float alpha) {
    // Schlick-GGX with k from alpha for direct lighting
    float k = (alpha + 1.0);
    k = (k * k) / 8.0;
    float ggx1 = geometry_schlick_ggx(NdotV, k);
    float ggx2 = geometry_schlick_ggx(NdotL, k);
    return ggx1 * ggx2;
}

struct PBRMaterial {
    vec3 baseColor;
    float metallic;
    float roughness;
    float ao;
    vec3 emissive;
};

PBRMaterial getMaterial() {
    // Base color (sampled as sRGB then assumed linear by GL; still gamma-correct final output)
    vec3 baseTex = (u_has_base_color_map != 0) ? texture(u_base_color_map, v_tex).rgb : vec3(1.0);
    vec3 baseColor = baseTex * u_base_color_factor.rgb;

    float metallicTex = (u_has_metallic_roughness_map != 0) ? texture(u_metallic_roughness_map, v_tex).b : 1.0;
    float roughnessTex = (u_has_metallic_roughness_map != 0) ? texture(u_metallic_roughness_map, v_tex).g : 1.0;

    float metallic  = clamp(u_metallic_factor  * metallicTex,  0.0, 1.0);
    float roughness = clamp(u_roughness_factor * roughnessTex, 0.04, 1.0);

    float ao = (u_has_occlusion_map != 0) ? texture(u_occlusion_map, v_tex).r : 1.0;

    vec3 emissiveTex = (u_has_emissive_map != 0) ? texture(u_emissive_map, v_tex).rgb : vec3(0.0);
    vec3 emissive = emissiveTex * u_emissive_factor;

    PBRMaterial m;
    m.baseColor = baseColor;
    m.metallic = metallic;
    m.roughness = roughness;
    m.ao = ao;
    m.emissive = emissive;
    return m;
}

vec3 evaluatePBR(vec3 N, vec3 V, vec3 L, vec3 lightColor, float lightIntensity, PBRMaterial m) {
    vec3 H = normalize(V + L);
    float NdotV = max(dot(N, V), 0.0);
    float NdotL = max(dot(N, L), 0.0);
    float NdotH = max(dot(N, H), 0.0);
    float HdotV = max(dot(H, V), 0.0);

    if (NdotL <= 0.0 || NdotV <= 0.0) return vec3(0.0);

    float alpha = m.roughness * m.roughness;

    vec3 F0 = mix(vec3(0.04), m.baseColor, m.metallic);
    vec3  F = fresnel_schlick(HdotV, F0);
    float D = distribution_ggx(NdotH, alpha);
    float G = geometry_smith(NdotV, NdotL, alpha);

    vec3  numerator   = F * D * G;
    float denominator = max(4.0 * NdotV * NdotL, 1e-6);
    vec3  specular    = numerator / denominator;

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

    vec3 diffuse = (m.baseColor / PI) * kD;

    vec3 radiance = lightColor * lightIntensity;

    return (diffuse + specular) * radiance * NdotL;
}

void main() {
    PBRMaterial mat = getMaterial();

    vec3 N = normalize(v_normal);
    vec3 V = normalize(-v_position); // camera at origin in view-space

    vec3 color = vec3(0.0);

    // Directional light contribution
    vec3 Ld = normalize(u_dir_light_dir);
    color += evaluatePBR(N, V, Ld, u_dir_light_color, u_dir_light_intensity, mat);

    // Point lights
    for (int i = 0; i < u_point_light_count && i < MAX_POINT_LIGHTS; ++i) {
        vec3 toLight = u_point_light_pos[i] - v_position;
        float dist   = length(toLight);
        vec3  L      = toLight / max(dist, 1e-4);

        float cutoff = 1.0 - smoothstep(0.9 * u_point_light_range[i], u_point_light_range[i], dist);
        float attenuation = (1.0 / max(dist * dist, 1e-4)) * cutoff;

        vec3 contrib = evaluatePBR(N, V, L, u_point_light_color[i], u_point_light_intensity[i] * attenuation, mat);
        color += contrib;
    }

    // Simple ambient and AO
    vec3 ambient = 0.03 * mat.baseColor * (1.0 - mat.metallic) * mat.ao;
    color += ambient + mat.emissive;

    // Gamma correction to sRGB
    color = pow(color, vec3(1.0 / 2.2));

    frag_color = vec4(color, u_base_color_factor.a);
}