#version 450
#extension GL_ARB_separate_shader_objects : enable

#define MAX_LIGHTS 4

struct AmbientPointLight
{
    vec3 color;
    vec3 pos;
};

struct PointLight
{
    vec3 color;
    float constant;
    vec3 pos;
    float linear;
    float quadratic;
};

struct SpotLight
{
    vec3 color;
    float innerCutOff;
    vec3 pos;
    float outerCutOff;
    vec3 dir;
    float constant;
    float linear;
    float quadratic;
};

layout(std140, binding = 2) uniform Data
{
    vec3 ambient;
    bool preMultiply;
    vec4 multiplier;
    uint ambientPointLightCount;
    uint pointLightCount;
    uint spotLightCount;
    AmbientPointLight ambientLights[MAX_LIGHTS];
    PointLight pointLights[MAX_LIGHTS];
    SpotLight spotLights[MAX_LIGHTS];
} data;

layout(binding = 3) uniform sampler2D colorMap;
layout(binding = 4) uniform sampler2D normalMap;
layout(binding = 5) uniform sampler2D specularMap;

layout(location = 0) in vec4 fVertPos;
layout(location = 1) in vec3 fTanPos;
layout(location = 2) in vec2 fUV;
layout(location = 3) in vec3 fViewPos;
layout(location = 4) in mat3 fTBN;

layout(location = 0) out vec4 result;

vec3 CalcAmbientPointLight(vec3 normal, vec3 vertexPos, vec3 viewDir, AmbientPointLight light)
{
    vec3 lightDir = normalize(light.pos - vertexPos);

    vec3 diffuse = light.color * max(dot(normal, lightDir), 0.0f);

    vec3 specular = light.color * pow(max(dot(normal, normalize(lightDir + viewDir)), 0.0), texture(specularMap, fUV).r);

    return diffuse + specular;
}

vec3 CalcPointLight(vec3 normal, vec3 vertexPos, vec3 viewDir, PointLight light)
{
    vec3 lightDir = normalize(light.pos - vertexPos);

    float dist = length(light.pos - vertexPos);

    float attenuation = 1.0f / (light.constant + light.linear * dist + light.quadratic * (dist * dist));

    vec3 diffuse = light.color * max(dot(normal, lightDir), 0.0f);

    vec3 specular = light.color * pow(max(dot(normal, normalize(lightDir + viewDir)), 0.0), texture(specularMap, fUV).r);

    return (diffuse + specular) * attenuation;
}

vec3 CalcSpotLight(vec3 normal, vec3 vertexPos, vec3 viewDir, SpotLight light)
{
    vec3 lightDir = normalize(light.pos - vertexPos);

    float dist = length(light.pos - vertexPos);

    float attenuation = 1.0f / (light.constant + light.linear * dist + light.quadratic * (dist * dist));

    float theta = dot(lightDir, normalize(-light.dir));
    float epsilon = light.innerCutOff - light.outerCutOff;
    float intensity = clamp((theta - light.outerCutOff) / epsilon, 0.0f, 1.0f);

    vec3 diffuse = light.color * max(dot(normal, lightDir), 0.0f);

    vec3 specular = light.color * pow(max(dot(normal, normalize(lightDir + viewDir)), 0.0), texture(specularMap, fUV).r);

    return (diffuse + specular) * intensity * attenuation;
}

void main()
{
    vec3 viewDir = normalize(fViewPos - fTanPos);
    vec3 nNorm = normalize(texture(normalMap, fUV).rgb * 2.0f - 1.0f);

    vec3 final = data.ambient;

    //Ambient Point Lights
    for (uint i = 0; i < data.ambientPointLightCount; ++i)
        final += CalcAmbientPointLight(nNorm, fVertPos.xyz, viewDir, data.ambientLights[i]);

    //Point Lights
    for (uint i = 0; i < data.pointLightCount; ++i)
        final += CalcPointLight(nNorm, fVertPos.xyz, viewDir, data.pointLights[i]);

    //Spot Lights
    for (uint i = 0; i < data.spotLightCount; ++i)
        final += CalcSpotLight(nNorm, fVertPos.xyz, viewDir, data.spotLights[i]);

    vec4 tColor = texture(colorMap, fUV);

    result = vec4(tColor.rgb * final * data.multiplier.rgb, tColor.a * data.multiplier.a);

    if (data.preMultiply)
        result = vec4(result.rgb * result.a, result.a);
}