#version 450 core

layout(binding = 0) uniform sampler2D heightTexture;

in TE_OUT {
    vec3 position;
    vec2 texCoord;
    vec3 normal;
} fs_in;

layout(std140, binding = 0) uniform SceneUniforms
{
    mat4 PV;	//camera projection * view matrix
    vec4 eye_w;	//world-space eye position
};

layout(location = 0) out vec4 FragColor;
layout(location = 1) out vec4 FragPosition;

float getHeight(vec2 uv) {
    return texture(heightTexture, uv).r * 0.025f;
}

vec3 calculateNormalsFromHeightTexture() {
    // Calculate height at base_uv
    float height = getHeight(fs_in.texCoord);

    // Calculate normals using height differences around the fragment
    float eps = 1.0f / textureSize(heightTexture, 0).x;
    float hRight = getHeight(fs_in.texCoord + vec2(eps, 0.0));
    float hLeft = getHeight(fs_in.texCoord - vec2(eps, 0.0));
    float hUp = getHeight(fs_in.texCoord + vec2(0.0, eps));
    float hDown = getHeight(fs_in.texCoord - vec2(0.0, eps));

    float heightTL = getHeight(fs_in.texCoord + vec2(-eps,  eps));
    float heightT  = getHeight(fs_in.texCoord + vec2( 0.0,  eps));
    float heightTR = getHeight(fs_in.texCoord + vec2( eps,  eps));
    float heightL  = getHeight(fs_in.texCoord + vec2(-eps,  0.0));
    float heightR  = getHeight(fs_in.texCoord + vec2( eps,  0.0));
    float heightBL = getHeight(fs_in.texCoord + vec2(-eps, -eps));
    float heightB  = getHeight(fs_in.texCoord + vec2( 0.0, -eps));
    float heightBR = getHeight(fs_in.texCoord + vec2( eps, -eps));

    float sobelX = (heightTL + 2.0 * heightL + heightBL) - (heightTR + 2.0 * heightR + heightBR);
    float sobelZ = (heightTL + 2.0 * heightT + heightTR) - (heightBL + 2.0 * heightB + heightBR);

    vec3 tangentX = normalize(vec3(eps, 0.0, sobelX));
    vec3 tangentZ = normalize(vec3(0.0, sobelZ, eps));

    // Final normal for the fragment
    return normalize(cross(tangentZ, tangentX));
}

void main() {
    //FragColor = vec4(fragNormal, 1.0f);
    float shininess = 1.0f;
    vec3 lightColor = vec3(0.5f);
    vec3 objectColor = vec3(texture(heightTexture, fs_in.texCoord).r);

    // Ambient
    vec3 ambient = vec3(0.15f) * vec3(0.15f);

    // Diffuse
    //vec3 norm = normalize(normal);
    vec3 norm = calculateNormalsFromHeightTexture();
    vec3 lightDir = normalize(vec3(1.0f, 1.0f, -1.0f) - fs_in.position);
    float diff = max(dot(norm, lightDir), 0.0);
    vec3 diffuse = diff * vec3(texture(heightTexture, fs_in.texCoord).r);

    // Specular (Phong)
    vec3 viewDir = normalize(eye_w.xyz - fs_in.position);
    vec3 reflectDir = reflect(-lightDir, norm);
    float spec = pow(max(dot(viewDir, reflectDir), 0.0), shininess);
    vec3 specular = spec * lightColor;

    // Combine results
    vec3 result = ambient + diffuse;
    FragColor = vec4(fs_in.normal, 1.0);
    FragPosition = vec4(fs_in.position, 1.0);

    //FragColor = vec4(vec3(texture(heightTexture, fs_in.texCoord).r), 1.0);
    //FragColor = vec4((fs_in.position * 0.5 + 0.5), 1.0);
}