#type vertex
#version 410 core

layout(location = 0) in vec3 v_position;
layout(location = 1) in vec3 v_normal;
layout(location = 2) in int v_ao;

out vec3 normalWS;
out vec3 normalVS;
out vec3 positionVS;
out vec3 positionWS;
out vec3 vertexPosition;
out float ao;

uniform mat4 u_worldMatrix;
uniform mat4 u_viewMatrix;

void main()
{
    vec4 worldPos = u_worldMatrix * vec4(v_position, 1.0);
    gl_Position = u_viewMatrix * worldPos;

    positionWS = worldPos.xyz;
    positionVS = gl_Position.xyz;
    vertexPosition = v_position;

    mat3 worldNormalMatrix = transpose(inverse(mat3(u_worldMatrix)));
    normalWS = normalize(worldNormalMatrix * v_normal);

    mat3 viewNormalMatrix = transpose(inverse(mat3(u_viewMatrix * u_worldMatrix)));
    normalVS = normalize(viewNormalMatrix * v_normal);

    ao = float(v_ao) / 255;
}

#type fragment
#version 410 core

layout(location = 0) out vec4 fragColor;
layout(location = 1) out int objectID;

in vec3 normalWS;
in vec3 normalVS;
in vec3 positionVS;
in vec3 positionWS;
in vec3 vertexPosition;
in float ao;


uniform vec3 u_cameraPos;
uniform int u_objectID;
uniform sampler2D u_texture;          // Albedo
uniform sampler2D u_texture_ambient;  // Ambient occlusion
uniform sampler2D u_texture_normal;   // Normal map
uniform sampler2D u_texture_roughness; // Roughness map

vec2 getFaceUV(vec3 pos, int face)
{
    float scaleFactor = 1.2;
    if(face == 0) return vec2(-fract(pos.z / scaleFactor), fract(pos.y / scaleFactor));
    if(face == 1) return vec2(fract(pos.z / scaleFactor), fract(pos.y / scaleFactor));
    if(face == 2) return vec2(fract(pos.x / scaleFactor), -fract(pos.z / scaleFactor));
    if(face == 3) return vec2(fract(pos.x / scaleFactor), fract(pos.z / scaleFactor));
    if(face == 4) return vec2(fract(pos.x / scaleFactor), fract(pos.y / scaleFactor));
    return vec2(-fract(pos.x / scaleFactor), fract(pos.y / scaleFactor));
}

void main()
{
    vec3 lightDir = normalize(vec3(3, 4, 1));

    // Cube face weights
    float faceWeights[6];
    faceWeights[0] = clamp(dot(vec3(-1,0,0), normalWS), 0.0, 1.0);
    faceWeights[1] = clamp(dot(vec3(1,0,0), normalWS), 0.0, 1.0);
    faceWeights[2] = clamp(dot(vec3(0,-1,0), normalWS), 0.0, 1.0);
    faceWeights[3] = clamp(dot(vec3(0,1,0), normalWS), 0.0, 1.0);
    faceWeights[4] = clamp(dot(vec3(0,0,-1), normalWS), 0.0, 1.0);
    faceWeights[5] = clamp(dot(vec3(0,0,1), normalWS), 0.0, 1.0);

    // Soft-hard blend: sharpen dominant face but allow slight interpolation
    float blendSharpness = 8.0; // higher = sharper
    float sumWeights = 0.0;
    for(int i = 0; i < 6; i++) {
        faceWeights[i] = pow(faceWeights[i], blendSharpness);
        sumWeights += faceWeights[i];
    }
    for(int i = 0; i < 6; i++)
        faceWeights[i] /= sumWeights;

    vec3 bitangents[6] = vec3[6](vec3(0,1,0), vec3(0,1,0), vec3(0,0,-1), vec3(0,0,1), vec3(0,1,0), vec3(0,1,0));
    vec3 tangents[6] = vec3[6](vec3(0, 0, -1), vec3(0, 0, 1), vec3(1, 0, 0), vec3(1, 0, 0), vec3(1, 0, 0), vec3(1, 0, 0));

    vec3 blendedColor = vec3(0.0);
    vec3 blendedNormal = vec3(0.0);
    float blendedAmbient = 0.0;
    float blendedPerceptualRoughness = 0.0;

    for(int i = 0; i < 6; i++)
    {
        if(faceWeights[i] == 0.0) continue;

        vec2 uv = getFaceUV(vertexPosition, i);

        vec3 texColor = texture(u_texture, uv).rgb;
        vec3 texNormal = texture(u_texture_normal, uv).rgb * 2.0 - 1.0;
        float texAmbient = texture(u_texture_ambient, uv).r;
        float texRoughness = texture(u_texture_roughness, uv).r;

        mat3 TBN = mat3(tangents[i], bitangents[i], normalize(normalWS));
        vec3 worldNormal = normalize(TBN * texNormal);

        blendedColor += texColor * faceWeights[i];
        blendedNormal += worldNormal * faceWeights[i];
        blendedAmbient += texAmbient * faceWeights[i];
        blendedPerceptualRoughness += texRoughness * faceWeights[i];
    }

    blendedNormal = normalize(blendedNormal);
    float blendedRoughness = sqrt(blendedPerceptualRoughness);

    // View direction (camera at origin in view space)
    vec3 viewDir = normalize(-vertexPosition); 

    // Schlick Fresnel
    vec3 F0 = vec3(0.04); // base reflectivity for dielectrics
    float cosTheta = clamp(dot(viewDir, blendedNormal), 0.0, 1.0);
    vec3 fresnel = F0 + (1.0 - F0) * pow(1.0 - cosTheta, 5.0);


    vec3 finalColor = blendedAmbient * blendedColor;

    vec3 N = normalize(blendedNormal);
    vec3 L = normalize(lightDir);
    vec3 V = normalize(u_cameraPos - positionWS);

    vec3 H = normalize(L + V);
    float roughnessClamped = clamp(1 - blendedRoughness, 0.001, 1.0);
    float specularExponent = (2.0 / (roughnessClamped * roughnessClamped)) - 2.0;

    float NdotH = max(dot(N, H), 0.0);
    float specular = pow(NdotH, specularExponent);

    float NdotL = max(dot(N, L), 0.0);
    vec3 diffuse = blendedColor * NdotL * blendedAmbient;

    blendedAmbient += 0.2;
    vec3 ambient = blendedColor * blendedAmbient;
    float add = clamp(dot(vec3(0, 0, 1), normalize(normalVS)) + 0.2,0, 1) * NdotL;
    finalColor = diffuse + specular * 0.3 + ambient * 0.2 + add;

    // Add specular
    /*
    finalColor = mix(
        finalColor,
        vec3(0.7, 0.9, 0.3) * NdotL,
        clamp(dot(blendedNormal, vec3(0, 1, 0)) * 10 - 3, 0, 1)
    );*/
    finalColor = mix(
        finalColor,
        vec3(0.8, 0.8, 0.8) * NdotL + vec3(0.3, 0.3, 0.3),
        clamp(dot(blendedNormal, vec3(0, 1, 0)) * 10 - 3, 0, 1)
    );

    fragColor = vec4(finalColor, 1.0);

    objectID = u_objectID;
}