#version 330 core
out vec4 color;

in vec3 FragPos;
in vec3 Normal;

uniform vec3 viewPos;

// Point Lights
struct PointLight {
    vec3 position;
    vec3 color;

    float constant;
    float linear;
    float quadratic;
};

#define MAX_POINT_LIGHTS 256
uniform int numPointLights;
uniform PointLight pointLights[MAX_POINT_LIGHTS];

// Directional Lights
struct DirLight {
    vec3 direction;
    vec3 color;
};

#define MAX_DIR_LIGHTS 10
uniform int numDirLights;
uniform DirLight dirLights[MAX_DIR_LIGHTS];

// Spotlights
struct SpotLight {
    vec3 position;
    vec3 direction;
    vec3 color;

    float cutOff;       // Cosine of the inner cutoff angle
    float outerCutOff;  // Cosine of the outer cutoff angle

    float constant;
    float linear;
    float quadratic;
};

#define MAX_SPOT_LIGHTS 64
uniform int numSpotLights;
uniform SpotLight spotLights[MAX_SPOT_LIGHTS];

uniform vec3 objectColor; // Fallback color
uniform sampler2D texture1; // Texture sampler
uniform bool useTexture; // Flag to determine if a texture is used

in vec2 TexCoords;

// Function to calculate lighting for a Point Light
vec3 CalcPointLight(PointLight light, vec3 normal, vec3 viewDir, vec3 fragPos)
{
    // Ambient
    float ambientStrength = 0.1;
    vec3 ambient = ambientStrength * light.color;

    // Diffuse
    vec3 lightDir = normalize(light.position - fragPos);
    float diff = max(dot(normal, lightDir), 0.0);
    vec3 diffuse = diff * light.color;

    // Specular
    float specularStrength = 0.5;
    vec3 reflectDir = reflect(-lightDir, normal);
    float spec = pow(max(dot(viewDir, reflectDir), 0.0), 32);
    vec3 specular = specularStrength * spec * light.color;

    // Attenuation
    float distance = length(light.position - fragPos);
    float attenuation = 1.0 / (light.constant + light.linear * distance + 
                               light.quadratic * (distance * distance));

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

// Function to calculate lighting for a Directional Light
vec3 CalcDirLight(DirLight light, vec3 normal, vec3 viewDir)
{
    // Ambient
    float ambientStrength = 0.1;
    vec3 ambient = ambientStrength * light.color;

    // Diffuse
    vec3 lightDir = normalize(-light.direction); // Directional light direction
    float diff = max(dot(normal, lightDir), 0.0);
    vec3 diffuse = diff * light.color;

    // Specular
    float specularStrength = 0.5;
    vec3 reflectDir = reflect(-lightDir, normal);
    float spec = pow(max(dot(viewDir, reflectDir), 0.0), 32);
    vec3 specular = specularStrength * spec * light.color;

    return (ambient + diffuse + specular);
}

// Function to calculate lighting for a Spotlight
vec3 CalcSpotLight(SpotLight light, vec3 normal, vec3 viewDir, vec3 fragPos)
{
    // Ambient
    float ambientStrength = 0.1;
    vec3 ambient = ambientStrength * light.color;

    // Diffuse
    vec3 lightDir = normalize(light.position - fragPos);
    float diff = max(dot(normal, lightDir), 0.0);
    vec3 diffuse = diff * light.color;

    // Specular
    float specularStrength = 0.5;
    vec3 reflectDir = reflect(-lightDir, normal);
    float spec = pow(max(dot(viewDir, reflectDir), 0.0), 32);
    vec3 specular = specularStrength * spec * light.color;

    // Attenuation
    float distance = length(light.position - fragPos);
    float attenuation = 1.0 / (light.constant + light.linear * distance + 
                               light.quadratic * (distance * distance));

    // Spotlight intensity
    vec3 spotDir = normalize(light.direction);
    float theta = dot(lightDir, spotDir);
    float epsilon = light.cutOff - light.outerCutOff;
    float intensity = clamp((theta - light.outerCutOff) / epsilon, 0.0, 1.0);

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

void main()
{
    vec3 norm = normalize(Normal);
    vec3 viewDir = normalize(viewPos - FragPos);
    vec3 result = vec3(0.0);

    // Calculate Point Lights
    for (int i = 0; i < numPointLights; ++i)
    {
        result += CalcPointLight(pointLights[i], norm, viewDir, FragPos);
    }

    // Calculate Directional Lights
    for (int i = 0; i < numDirLights; ++i)
    {
        result += CalcDirLight(dirLights[i], norm, viewDir);
    }

    // Calculate Spotlights
    for (int i = 0; i < numSpotLights; ++i)
    {
        result += CalcSpotLight(spotLights[i], norm, viewDir, FragPos);
    }

    // Combine lighting with object color or texture
    vec3 baseColor = useTexture ? texture(texture1, TexCoords).rgb : objectColor;
    color = vec4(result * baseColor, 1.0);
}