#version 450
/*
 * Copyright 2023  TU Wien, Institute of Visual Computing & Human-Centered Technology.
 * This file is part of the GCG Lab Framework and must not be redistributed.
 *
 * Original version created by Lukas Gersthofer and Bernhard Steiner.
 * Vulkan edition created by Johannes Unterguggenberger (junt@cg.tuwien.ac.at).
 *
 * 2026: Adapted for the PTVC framework by Balázs Kovács.
 */

const float POINT_LIGHT_INTENSITY = 7.5;
const float SHADOW_MAP_FAR_PLANE = 25.0;

const float FOG_START = 15.0;
const float FOG_END = 30.0;
const vec3 FOG_COLOR = vec3(0.0, 0.0, 0.0);

layout(set = 0, binding = 0) uniform sampler2D uPosition;
layout(set = 0, binding = 1) uniform sampler2D uNormal;
layout(set = 0, binding = 2) uniform sampler2D uAlbedo;
layout(set = 0, binding = 3) uniform sampler2D uMaterialProperties;

layout(set = 1, binding = 0) uniform UniformBuffer {
    mat4 view;
    mat4 proj;
    mat4 viewInverse;
    mat4 projInverse;
    vec4 position;
    float nearPlane;
    float farPlane;
    int isDebugCam;
} ub_data;

layout(set = 1, binding = 1) uniform DirectionalLight {
    vec4 color;
    vec4 direction;
} dl_data;

layout(set = 1, binding = 2) uniform PointLight {
    vec4 color;
    vec4 position;
    vec4 attenuation;
} pl_data;

layout(set = 2, binding = 0) uniform samplerCube uShadowMap;

layout(set = 3, binding = 0) uniform sampler2D uSSAO;

layout(location = 0) in vec2 inUV;
layout(location = 0) out vec4 out_color;

vec3 phong(vec3 n, vec3 l, vec3 v, vec3 diffuseC, float diffuseF, vec3 specularC, float specularF, float alpha, bool attenuate, vec3 attenuation, float ao) {
    float d = length(l);
    l = normalize(l);
    float att = 1.0;
    if (attenuate) {
        att = 1.0f / (attenuation.x + d * attenuation.y + d * d * attenuation.z);
    }
    vec3 h = normalize(l + v);
    return ao * (diffuseF * diffuseC * max(0, dot(n, l)) + specularF * specularC * pow(max(0, dot(n, h)), alpha)) * att;
}

float shadowCalculation(vec3 lightToFrag) {
    float currentDepth = length(lightToFrag);

    if (currentDepth > SHADOW_MAP_FAR_PLANE) {
        return 0.0;
    }

    float shadow = 0.0;
    float bias = 0.05;
    float samples = 4.0;
    float offset = 0.1;
    float diskRadius = 0.05;

    float closestDepthSum = 0.0;

    for (float x = -offset; x < offset; x += offset / (samples * 0.5)) {
        for (float y = -offset; y < offset; y += offset / (samples * 0.5)) {
            for (float z = -offset; z < offset; z += offset / (samples * 0.5)) {
                float closestDepth = texture(uShadowMap, lightToFrag + vec3(x, y, z) * diskRadius).r;
                closestDepth *= SHADOW_MAP_FAR_PLANE;
                closestDepthSum += closestDepth;

                if (currentDepth - bias > closestDepth) {
                    shadow += 1.0;
                }
            }
        }
    }
    shadow /= (samples * samples * samples);
    float closestDepthAvg = closestDepthSum / (samples * samples * samples);

    float shadowFade = 1.0 - smoothstep(0.0, SHADOW_MAP_FAR_PLANE * 0.5, currentDepth - closestDepthAvg);
    shadow *= shadowFade;

    return shadow;
}

// Computes the reflection direction for an incident vector I about normal N,
// and clamps the reflection to a maximum of 180, i.e. the reflection vector
// will always lie within the hemisphere around normal N.
// Aside from clamping, this function produces the same result as GLSL's reflect function.
vec3 clampedReflect(vec3 I, vec3 N)
{
    return I - 2.0 * min(dot(N, I), 0.0) * N;
}

// Gets the reflected color value from a a certain position, from
// a certain direction INSIDE of a cornell box of size 3 which is
// positioned at the origin.
// positionWS:  Position inside the cornell box for which to get the
//              reflected color value for from -directionWS.
// directionWS: Outgoing direction vector (from positionWS towards the
//              outside) for which to get the reflected color value for.
vec3 getCornellBoxReflectionColor(vec3 positionWS, vec3 directionWS) {
    vec3 P0 = positionWS;
    vec3 V = normalize(directionWS);

    const float boxSize = 1.5;
    vec4[5] planes = {
            vec4(-1.0, 0.0, 0.0, -boxSize), // left
            vec4(1.0, 0.0, 0.0, -boxSize), // right
            vec4(0.0, 1.0, 0.0, -boxSize), // top
            vec4(0.0, -1.0, 0.0, -boxSize), // bottom
            vec4(0.0, 0.0, -1.0, -boxSize) // back
        };
    vec3[5] colors = {
            vec3(0.49, 0.06, 0.22), // left
            vec3(0.0, 0.13, 0.31), // right
            vec3(0.96, 0.93, 0.85), // top
            vec3(0.64, 0.64, 0.64), // bottom
            vec3(0.76, 0.74, 0.68) // back
        };

    for (int i = 0; i < 5; ++i) {
        vec3 N = planes[i].xyz;
        float d = planes[i].w;
        float denom = dot(V, N);
        if (denom <= 0) continue;
        float t = -(dot(P0, N) + d) / denom;
        vec3 P = P0 + t * V;
        float q = boxSize + 0.01;
        if (P.x > -q && P.x < q && P.y > -q && P.y < q && P.z > -q && P.z < q) {
            return colors[i];
        }
    }
    return vec3(0.0, 0.0, 0.0);
}

vec3 fresnelSchlick(float cosTheta, vec3 F0) {
    return F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0);
}

void main() {
    vec3 fragPosView = texture(uPosition, inUV).xyz;
    vec3 normalView = normalize(texture(uNormal, inUV).xyz);
    vec3 albedo = texture(uAlbedo, inUV).rgb;
    vec4 materialProperties = texture(uMaterialProperties, inUV);
    float ao = texture(uSSAO, inUV).r;

    vec3 worldPos = (ub_data.viewInverse * vec4(fragPosView, 1.0)).xyz;
    vec3 normal = normalize(mat3(ub_data.viewInverse) * normalView);

    vec3 n = normalize(normal);
    vec3 v = normalize(worldPos - ub_data.position.xyz);
    vec3 R = normalize(clampedReflect(v, n));
    vec3 reflectionColor = getCornellBoxReflectionColor(worldPos, R);
    vec3 F0 = vec3(0.1); // <-- some kind of plastic
    vec3 reflectivity = fresnelSchlick(dot(n, -v), F0);

    // Start with ambient illumination contribution:
    vec3 color = albedo * materialProperties.x;
    float diffuseF = materialProperties.y;
    float specularF = materialProperties.z;
    float specularA = materialProperties.w;

    if (specularF == 0.0) {
        ao = 0.0;
    }

    // Add directional light's contribution:
    color += phong(
            n,
            -dl_data.direction.xyz,
            -v,
            dl_data.color.rgb * albedo, diffuseF,
            dl_data.color.rgb, specularF, specularA,
            false, vec3(1.0),
            ao
        );

    vec3 lightDir = pl_data.position.xyz - worldPos;

    // Add point light's contribution
    color += POINT_LIGHT_INTENSITY * phong(
                n,
                lightDir,
                -v,
                pl_data.color.rgb * albedo, diffuseF,
                pl_data.color.rgb, specularF, specularA,
                ub_data.isDebugCam == 0, pl_data.attenuation.xyz,
                ao
            );

    float shadow = 0.0;
    if (ub_data.isDebugCam == 0) {
        shadow = shadowCalculation(-lightDir);
    }

    color *= (1.0 - shadow);

    float dist = length(worldPos - ub_data.position.xyz);
    float fogFactor = clamp((dist - FOG_START) / (FOG_END - FOG_START), 0.0, 1.0);
    vec3 finalColor = mix(color, FOG_COLOR, fogFactor);

    if (dist > FOG_END) {
        finalColor = FOG_COLOR;
    }

    // Write color for the current fragment:
    out_color = vec4(finalColor, 1.0);

    // no fresnel >:(
    if (false)
    {
        vec3 mixedColor = mix(color, reflectionColor, reflectivity);
        out_color = vec4(mixedColor, 1.0);
    }
}
