#version 460 core

in vec4 vertex_color;
in vec2 texture_coord;
in vec3 normal;
in vec3 frag_pos;

out vec4 color;

// Must match the same as in math.h
const int MAX_POINT_LIGHTS = 3;
const int MAX_SPOT_LIGHTS = 3;

struct Light {
    vec3 color;
    float ambient_intensity;
    float diffuse_intensity;
};

struct Directional_Light {
    Light light;
    vec3 direction;
};

struct Point_Light {
    Light light;
    vec3 position;
    float constant;
    float linear;
    float exponent;
};

struct Spot_Light {
    Point_Light point_light;

    vec3 direction;
    float edge;
};

struct Material {
    float shininess;
    float specular_intensity;
};


uniform Directional_Light sun;
uniform Point_Light point_lights[MAX_POINT_LIGHTS];
uniform Spot_Light spot_lights[MAX_SPOT_LIGHTS];

uniform int point_light_count;
uniform int spot_light_count;
uniform sampler2D tex;
uniform Material material;
uniform vec3 eye_position;

vec4 calculate_light_by_direction(Light light, vec3 direction) {
    vec4 ambient_color = vec4(light.color, 1.0f) * light.ambient_intensity;

    float diffuse_factor = max(dot(normalize(normal), normalize(direction)), 0.0f);
    vec4 diffuse_color = vec4(light.color * light.diffuse_intensity * diffuse_factor, 1.0f);

    vec4 specular_color = vec4(0, 0, 0, 0);

    if (diffuse_factor > 0.0f) {
        vec3 frag_to_eye = normalize(eye_position - frag_pos);
        vec3 reflected_vertex = normalize(reflect(direction, normalize(normal)));
        float specular_factor = dot(frag_to_eye, reflected_vertex);

        if (specular_factor > 0.0f) {
            specular_factor = pow(specular_factor, material.shininess);
            specular_color = vec4(light.color * material.specular_intensity * specular_factor, 1.0f);
        }
    }

    return ambient_color + diffuse_color + specular_color;
}

vec4 calculate_point_light(Point_Light point_light) {
    vec3 direction = frag_pos - point_light.position;
    float distance = length(direction);
    direction = normalize(direction);

    vec4 light_color = calculate_light_by_direction(point_light.light, direction);
    float attentuation = point_light.exponent * distance * distance + point_light.linear * distance + point_light.constant;

    return (light_color / attentuation);
}

vec4 calculate_point_lights() {
    vec4 final_color = vec4(0, 0, 0, 0);

    for (int i=0; i<point_light_count; i++) {
        final_color += calculate_point_light(point_lights[i]);
    }

    return final_color;
}

vec4 calculate_spot_light(Spot_Light spot_light) {
    vec3 ray_direction = normalize(frag_pos - spot_light.point_light.position);
    float spot_light_factor = dot(ray_direction, spot_light.direction);

    if (spot_light_factor > spot_light.edge) {
        vec4 light_color = calculate_point_light(spot_light.point_light);

        return light_color * (1.0f - (1.0f - spot_light_factor)*(1.0f/(1.0f - spot_light.edge)));
    }

    return vec4(0, 0, 0, 0);
}

vec4 calculate_spot_lights() {
    vec4 final_color = vec4(0, 0, 0, 0);

    for (int i=0; i<spot_light_count; i++) {
        final_color += calculate_spot_light(spot_lights[i]);
    }

    return final_color;
}

void main() {
    vec4 final_color = calculate_light_by_direction(sun.light, sun.direction);

    final_color += calculate_point_lights();
    final_color += calculate_spot_lights();

    color = texture(tex, texture_coord) * final_color;
}