#version 450

struct Light {
    vec4 color;
    vec4 dir;
    vec4 pos;
    vec4 params;
};

layout(location = 0) in float iHeight;
layout(location = 1) in vec3 iNormal;
layout(location = 2) in vec3 iWorldPos;
layout(location = 3) in vec3 iColor;
layout(location = 4) in vec4 iLightFragPos;

layout(location = 0) out vec4 oColor;

layout(set = 0, binding = 0) uniform uGlobal {
    mat4 view;
    mat4 proj;
    vec4 camPos;
};

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

layout(std430, set = 2, binding = 0) readonly buffer Lights {
    Light lights[];
};

layout(set = 3, binding = 0) uniform uFog {
    vec4 fogColor;
    float fogEnd;
    float fogDensity;
};

const float PI = 3.14159265359f;
//vec3 albedo = vec3(0.01f);
vec3 albedo = iColor;
float roughness = 0.5f;
float metallic = 0.0f;

float alpha = pow(2.0f, metallic * 8.0f);
float kd = roughness;
float ks = 1.0f - kd;
float ka = 0.05f;

// fog code inspired by https://youtu.be/oQksg57qsRA?si=iFEi0f9zL1c-w_Gb
float calculateFog() {
    float distanceCameraToFragment = length(iWorldPos - camPos.xyz);
    float distanceRatio = 4.0f * distanceCameraToFragment / fogEnd;
    float fogFactor = exp(-distanceRatio * distanceRatio * fogDensity * fogDensity);

    return fogFactor;
}

// PBR inspired by
// - https://learnopengl.com/PBR/Theory
// - https://rtarun9.github.io/blogs/physically_based_rendering/

// approximates the reflected light
vec3 fresnelSchlick(float vDotH, vec3 F0) {
    return F0 + (1.0f - F0) * pow(clamp(1.0f - vDotH, 0.0f, 1.0f), 5.0f);
}

// based on roughness approximates alignment of the surface's microfacets
float normalDistribution(vec3 normal, vec3 halfWayVector) {
    float alpha = roughness * roughness;
    float nDotH = max(dot(normal, halfWayVector), 0.0f);

    return (alpha * alpha) / (PI * pow((nDotH * nDotH) * (alpha * alpha - 1.0f) + 1.0f, 2.0f));
}

// approximates the relative surface area where micro surface-details overshadow each other
float geometrySchlick(float x) {
    // k is a remapping of alpha, differs depending on if image based lighting is used or not (we aren't using IBL)
    float r = roughness + 1.0f;
    float k = (r * r) / 8.0f;

    return x / (x * (1.0f - k) + k);
}

// take view (geometry obstruction) and light (geometry shadowing) directions into account
float geometrySmith(vec3 normal, vec3 viewDirection, vec3 lightDirection) {
    float nDotV = max(dot(normal, viewDirection), 0.0f);
    float nDotL = max(dot(normal, lightDirection), 0.0f);

    return geometrySchlick(nDotV) * geometrySchlick(nDotL);
}

// PBR lighting calcs (btw calcs stands for calculations, I'm just using slang)
vec3 cookTorrenceBRDF(vec3 normal, vec3 viewDirection, vec3 lightDirection, vec3 halfWayVector) {
    vec3 F0 = mix(vec3(0.04f), albedo, metallic); // F0 = base reflectivity
    vec3 F = fresnelSchlick(max(dot(halfWayVector, viewDirection), 0.0f), F0);
    float NDF = normalDistribution(normal, halfWayVector);
    float G = geometrySmith(normal, viewDirection, lightDirection);

    vec3 specular = (NDF * G * F) / (4.0f * max(dot(normal, viewDirection), 0.0f) * max(dot(normal, lightDirection), 0.0f) + 0.0001f); // last constant to prevent divide by zero

    vec3 kd = (vec3(1.0f) - F) * (1.0f - metallic);
    vec3 diffuse = kd * (albedo / PI);

    return diffuse + specular;
}

vec3 calculateDirLight(int lightIndex) {
    vec3 normal = normalize(iNormal);
    vec3 viewDirection = normalize(camPos.xyz - iWorldPos);
    vec3 lightDirection = normalize(-lights[lightIndex].dir.xyz);
    vec3 halfWayVector = normalize(lightDirection + viewDirection);

    vec3 BDRF = cookTorrenceBRDF(normal, viewDirection, lightDirection, halfWayVector);
    vec3 radiance = lights[lightIndex].color.rgb;
    float nDotL = max(dot(normal, lightDirection), 0.0f);

    return BDRF * radiance * nDotL;
}

vec3 calculatePointLight(int lightIndex) {
    vec3 lightPos = lights[lightIndex].pos.xyz;

    vec3 normal = normalize(iNormal);
    vec3 viewDirection = normalize(camPos.xyz - iWorldPos);
    vec3 lightDirection = normalize(lightPos - iWorldPos);
    vec3 halfWayVector = normalize(lightDirection + viewDirection);

    float distance = length(lightPos - iWorldPos);
    float attenuation = 1.0f / (1.0f + lights[lightIndex].params.x * distance + lights[lightIndex].params.y * (distance * distance));

    vec3 BDRF = cookTorrenceBRDF(normal, viewDirection, lightDirection, halfWayVector);
    vec3 radiance = lights[lightIndex].color.rgb * attenuation;
    float nDotL = max(dot(normal, lightDirection), 0.0f);

    return BDRF * radiance * nDotL;
}

vec3 calculateSpotLight(int lightIndex) {
    vec3 lightPos = lights[lightIndex].pos.xyz;
    vec3 lightDirection = normalize(lightPos - iWorldPos);
    vec3 spotDirection = -normalize(lights[lightIndex].dir.xyz);

    vec3 result = vec3(0.0);

    float theta = dot(lightDirection, spotDirection);
    if (theta > lights[lightIndex].params.z) {
        vec3 normal = normalize(iNormal);
        vec3 viewDirection = normalize(camPos.xyz - iWorldPos);
        vec3 halfWayVector = normalize(lightDirection + viewDirection);

        float distance = length(lightPos - iWorldPos);
        //        float attenuation = 1.0f / (distance * distance);
        float attenuation = 1.0f / (1.0f + lights[lightIndex].params.x * distance + lights[lightIndex].params.y * (distance * distance));
        float intensity = smoothstep(lights[lightIndex].params.w, lights[lightIndex].params.z, theta);
        intensity *= intensity;

        vec3 BDRF = cookTorrenceBRDF(normal, viewDirection, lightDirection, halfWayVector);
        vec3 radiance = lights[lightIndex].color.rgb * attenuation;
        float nDotL = max(dot(normal, lightDirection), 0.0f);

        result = BDRF * radiance * nDotL * intensity;
    }

    return result;
}

float calculateShadow() {
    // perspective divide (needed in case of perspective projection)
    vec3 projCoords = iLightFragPos.xyz / iLightFragPos.w;
    // for some reason Vulkans NDC are in [-1, 1] for x and y, but in [0, 1] for z, so we convert xy to [0, 1]
    projCoords.xy = projCoords.xy * 0.5f + 0.5f;

    vec2 shadowMapUV = projCoords.xy;
    float currentDepth = projCoords.z;

    // return if fragment is outside shadow map's projection,
    // as in behind the far plane, left/right/top/bottom are taken care of by the sampler
    if (currentDepth > 1.0f) { return 0.0f; }

    // Percentage-Close Filtering (PCF), basically sampling neighboring texels and averaging the result
    float shadow = 0.0f;
    vec2 texelSize = 1.0f / textureSize(shadowMap, 0);

    for (int x = -1; x <= 1; x++) {
        for(int y = -1; y <= 1; y++) {
            vec2 offset = vec2(x, y) * texelSize;
            float pcfDepth = texture(shadowMap, shadowMapUV + offset).r;
            shadow += currentDepth > pcfDepth ? 1.0f : 0.0f;
        }
    }

    shadow /= 9.0f;

    return shadow;
}

void main() {
    vec3 normal = normalize(iNormal);
    //normal = normalize(cross(dFdx(iWorldPos), dFdy(iWorldPos))); // flat overwrite

    //vec3 planetColor = vec3(0.2, 0.2, 0.2);
    //vec3 floraColor = vec3(0.1, 0.3, 0.1);
    vec3 planetColor = vec3(0.2); // * iColor;
    vec3 floraColor = iColor;

    // slope styles
    float slope = clamp(dot(normal, vec3(0, 1, 0)), 0, 1);
    slope = pow(slope, 2);

    vec3 color = mix(planetColor, floraColor, slope);
    int levels = 40;

    //height lines
    float f = fract(iHeight * levels);
    float line = step(f, 0.05);
    color = mix(color, color * 1.3, line);

    // height faces
    float l = floor(iHeight * levels) / levels;
    color = mix(color, color * 0.1, 1 - l);

    // cell shading
    float steps = 4;
    //diffuse = floor(diffuse * steps) / steps;

    vec3 lighting = vec3(0.0f);

    for (int i = 0; i < lights.length(); i++) {
        if (lights[i].dir.w == 0.0f) {
            lighting += calculateDirLight(i);
        } else if (lights[i].dir.w == 1.0f) {
            lighting += calculatePointLight(i);
        } else if (lights[i].dir.w == 2.0f) {
            lighting += calculateSpotLight(i);
        }
    }

    vec3 ambient = vec3(0.03f) * albedo;
    float shadow = calculateShadow();

    color = (ambient + (1.0 - shadow) * lighting) * color;

    color = color / (color + vec3(1.0f)); // HDR tonemapping
    color = pow(color, vec3(1.0f / 2.2f)); // gamma correction

    float fogFactor = calculateFog();
    color = mix(fogColor.rgb, color, fogFactor);

    oColor = vec4(color, 1.0f);

//    vec3 projCoords = iLightFragPos.xyz / iLightFragPos.w;
//    vec4 debugColor = vec4(1.0f);
//    if (projCoords.z > 1.0) debugColor = vec4(0.0f, 1.0f, 0.0f, 1.0f);        // Green = outside far plane
//    else if (projCoords.z < 0.0) debugColor = vec4(0.0f, 0.0f, 1.0f, 1.0f);   // Blue = before near plane
//    else debugColor = vec4(projCoords.z, 0.0f, 0.0f, 1.0f);                   // Red = inside frustum
//
//    oColor = debugColor;
}
