#include "lighting.glsl"

in VS_OUT {
    vec3 FragPos;
    vec3 Normal;
    vec2 TexCoords;
    mat3 TBN;
} vs_in;

layout(location = 0) out vec4 outColor;
layout(location = 1) out vec4 brightColor;

struct Light {
    vec4 position;
    vec3 color;
    int type;
    vec3 direction;
    float radius;
    mat4 lightspace;
};

layout(std140, binding = 0) uniform UniformSceneData {
    mat4 projection, view, model;
    vec4 cameraPosition;
} sceneData;

layout(std140, binding = 1) uniform UniformMaterialData {
    vec4 albedoTint;
    vec3 emission;
    float metallic, roughness;
} materialData;

layout(std140, binding = 2) uniform UniformLightData {
    Light lights[MAX_LIGHTS];
    int numlights;
} lightData;

const mat4 biasMat = mat4(
	0.5, 0.0, 0.0, 0.0,
	0.0, 0.5, 0.0, 0.0,
	0.0, 0.0, 1.0, 0.0,
	0.5, 0.5, 0.0, 1.0 );

layout(binding = 3) uniform sampler2D albedo_tex;
layout(binding = 4) uniform sampler2D normal_tex;
layout(binding = 5) uniform sampler2D metallic_tex;
layout(binding = 6) uniform sampler2D roughness_tex;
layout(binding = 7) uniform sampler2D directional_shadow;
layout(binding = 8) uniform samplerCube point_shadow;

float PointShadowCalculation(Light light)
{
    if(light.position.w == 0.0f)
        return 1.0f;

    vec3 fragToLight = vs_in.FragPos - light.position.xyz;
    float closestDepth = texture(point_shadow, fragToLight).r;
    closestDepth *= 25.0; //far-plane

    float currentDepth = length(fragToLight);

    float bias = 0.05;
    float shadow = currentDepth -  bias > closestDepth ? 0.0 : 1.0;

    return shadow;
}

float ShadowCalculation(Light light)
{
    if(light.position.w == 0.0f)
        return 1.0f;

    vec4 fragPosLightSpace = light.lightspace * vec4(vs_in.FragPos, 1.0);

    float bias = 0.001;

    vec3 projCoords = fragPosLightSpace.xyz / fragPosLightSpace.w;
    projCoords = projCoords * 0.5 + 0.5;
    float closestDepth = texture(directional_shadow, projCoords.xy).r;
    float currentDepth = projCoords.z;
    float shadow = currentDepth - bias > closestDepth ? 0.0 : 1.0;

    return shadow;
}

vec4 GetAlbedo() {
    if(textureSize(albedo_tex, 0).x <= 1)
        return vec4(materialData.albedoTint.rgb, 1.0);

    return vec4(texture(albedo_tex, vs_in.TexCoords).rgb * materialData.albedoTint.rgb, texture(albedo_tex, vs_in.TexCoords).a);
}

vec3 GetNormal() {
    if(textureSize(normal_tex, 0).x <= 1)
        return normalize(vs_in.Normal);

    vec3 normal = texture(normal_tex, vs_in.TexCoords).rgb;
    normal = 2.0 * normal - 1.0;

    normal = vs_in.TBN * normal;

    return normalize(normal);
}

float GetMetallic() {
    if(textureSize(metallic_tex, 0).x <= 1)
        return materialData.metallic;

    return texture(metallic_tex, vs_in.TexCoords).r;
}

float GetRoughness() {
    if(textureSize(roughness_tex, 0).x <= 1)
        return materialData.roughness;

    return texture(roughness_tex, vs_in.TexCoords).r;
}

vec3 CalculatePointLight(Light light, vec3 kD, vec3 F) {
    vec3 N = GetNormal();
    vec3 V = normalize(sceneData.cameraPosition.xyz - vs_in.FragPos);

    vec3 L = normalize(light.position.xyz - vs_in.FragPos);
    vec3 H = normalize(V + L);
    vec3 radiance = (light.color * (CalculateAttenuation(light.position.xyz, light.radius, vs_in.FragPos) * light.radius)) * PointShadowCalculation(light);

    float NDF = DistributionGGX(N, H, GetRoughness());
    float G = GeometrySmith(N, V, L, GetRoughness());

    vec3 nominator = NDF * G * F;
    float denominator = 4 * max(dot(V, N), 0.0) * max(dot(L, N), 0.0) + 0.001;
    vec3 brdf = nominator / denominator;

    // add to outgoing radiance Lo
    float NdotL = max(dot(N, L), 0.0);

    return (kD * GetAlbedo().rgb / PI + brdf) * radiance * NdotL;
}

vec3 CalculateDirectionalLight(Light light, vec3 kD, vec3 F) {
    vec3 N = GetNormal();
    vec3 V = normalize(sceneData.cameraPosition.xyz - vs_in.FragPos);

    vec3 L = normalize(-light.direction);
    vec3 H = normalize(V + L);
    vec3 radiance = light.color * ShadowCalculation(light);

    float NDF = DistributionGGX(N, H, GetRoughness());
    float G = GeometrySmith(N, V, L, GetRoughness());

    vec3 nominator = NDF * G * F;
    float denominator = 4 * max(dot(V, N), 0.0) * max(dot(L, N), 0.0) + 0.001;
    vec3 brdf = nominator / denominator;

    // add to outgoing radiance Lo
    float NdotL = max(dot(N, L), 0.0);

    return (kD * GetAlbedo().rgb / PI + brdf) * radiance * NdotL;
}

void main() {
    if(GetAlbedo().a < 0.1)
        discard;

    vec3 V = normalize(sceneData.cameraPosition.xyz - vs_in.FragPos);

    vec3 F0 = vec3(0.04);
    F0 = mix(F0, GetAlbedo().rgb, GetMetallic());
    vec3 F = fresnelSchlickRoughness(max(dot(GetNormal(), V), 0.0), F0, GetRoughness());

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

    vec3 result = vec3(0);
    for(int i = 0; i < lightData.numlights; i++) {
        switch(lightData.lights[i].type) {
            case 0:
                result += CalculateDirectionalLight(lightData.lights[i], kD, F);
                break;
            case 1:
                result += CalculatePointLight(lightData.lights[i], kD, F);
                break;
        }
    }

    vec3 ambient = vec3(0.01) * GetAlbedo().rgb;

    outColor = vec4(ambient + result + materialData.emission, GetAlbedo().a);

    float brightness = dot(outColor.rgb, vec3(0.2126, 0.7152, 0.0722));
    if(brightness > 0.9)
        brightColor = vec4(outColor.rgb, 1.0);
}