// -*- mode: c -*-
#extension GL_NV_shadow_samplers_cube : enable

struct Material {
  vec3 baseColorFactor;
  int baseColorTexcoord;
  float metallicFactor;
  float roughnessFactor;
  int metallicRoughnessTexcoord;
  int normalTexcoord;
  int occlusionTexcoord;
  vec3 emissiveFactor;
  int emissiveTexcoord;
  int alphaMode;
  float alphaCutoff;
};

struct Light {
  bool enabled;
  int type;
  vec3 position;
  vec3 direction;
  vec4 color;
  float intensity;
  float cutOff;
};

#ifdef GLSL120
varying vec3 fragWorldPos;
varying vec3 fragNormal;
varying vec3 fragTangent;
varying vec2 fragTexcoord0;
varying vec2 fragTexcoord1;
varying vec4 fragColor0;
#else
in vec3 fragWorldPos;
in vec3 fragNormal;
in vec3 fragTangent;
in vec2 fragTexcoord0;
in vec2 fragTexcoord1;
in vec4 fragColor0;
#endif

#ifdef GLSL330
out vec4 fragColor;
#endif

#define MAX_LIGHTS 4

uniform Material material;
uniform Light lights[MAX_LIGHTS];
uniform bool vertexColored;
uniform vec3 cameraPosition;
uniform samplerCube skybox;
uniform sampler2D baseColorTexture;
uniform sampler2D metallicRoughnessTexture;
uniform sampler2D normalTexture;
uniform sampler2D occlusionTexture;
uniform sampler2D emissiveTexture;

const float PI = 3.14159265358979323846;
const float GAMMA = 2.2;

#ifndef GLSL330
// Compatibility shim for older GLSL versions.
vec4 texture(sampler2D tex, vec2 coord) {
  return texture2D(tex, coord);
}
vec4 texture(samplerCube tex, vec3 coord) {
  return textureCube(tex, coord);
}
#endif

float posDot(vec3 v1, vec3 v2) {
  return max(dot(v1, v2), 0.0);
}

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

float distributionGGX(vec3 normal, vec3 halfAngle, float roughness)
{
    float a = roughness * roughness * roughness * roughness;
    float ndoth = posDot(normal, halfAngle);
    float denominator = ndoth * ndoth * (a - 1.0) + 1.0;

    return a / (PI * denominator * denominator);
}

float geometrySchlickGGX(float ndotv, float roughness)
{
    float r = roughness + 1.0;
    float k = (r * r) / 8.0;

    return ndotv / (ndotv * (1.0 - k) + k);
}

float geometrySmith(vec3 normal, vec3 viewDirection, vec3 lightDirection,
                    float roughness)
{
    return geometrySchlickGGX(posDot(normal, viewDirection), roughness) *
      geometrySchlickGGX(posDot(normal, lightDirection), roughness);
}

vec4 applyAlpha(vec4 color) {
  // Apply alpha mode.
  if(material.alphaMode == 0) { // opaque
    return vec4(color.rgb, 1.0);
  } else if(material.alphaMode == 1) { // mask
    if(color.a >= material.alphaCutoff) {
      return vec4(color.rgb, 1.0);
    } else {
      discard;
    }
  } else if(material.alphaMode == 2) { // blend
    if(color.a <= 0.005) {
      discard;
    } else {
      return color;
    }
  }
}

vec2 texcoord(int i) {
  return i == 0 ? fragTexcoord0: fragTexcoord1;
}

vec4 sRGBtoLinear(vec4 srgb) {
  return vec4(pow(srgb.rgb, vec3(GAMMA)), srgb.a);
}

vec3 gammaCorrect(vec3 color) {
  return pow(color, vec3(1.0 / GAMMA));
}

vec3 toneMap(vec3 color) {
  return color / (color + vec3(1.0));
}

float materialMetallic() {
  float m = material.metallicFactor;
  m *= texture(metallicRoughnessTexture,
               texcoord(material.metallicRoughnessTexcoord)).b;
  return m;
}

float materialRoughness() {
  float r = material.roughnessFactor;

  r *= texture(metallicRoughnessTexture,
               texcoord(material.metallicRoughnessTexcoord)).g;

  return r;
}

vec4 materialAlbedo() {
  vec4 color = vec4(1.0, 1.0, 1.0, 1.0);
  vec4 texColor = texture(baseColorTexture,
                          texcoord(material.baseColorTexcoord));
  color = sRGBtoLinear(texColor);

  color *= vec4(material.baseColorFactor, 1.0);

  if(vertexColored) {
    color *= fragColor0;
  }

  return color;
}

vec3 materialEmissive() {
  vec3 color = vec3(0.0);

  vec4 texColor = texture(emissiveTexture,
                          texcoord(material.emissiveTexcoord));
  color = sRGBtoLinear(texColor).rgb;

  return color * material.emissiveFactor;
}

vec3 materialOcclusion() {
  return vec3(texture(occlusionTexture,
                      texcoord(material.occlusionTexcoord)).r);
}

vec3 materialNormal() {
  // See: https://github.com/SaschaWillems/Vulkan-glTF-PBR/blob/master/data/shaders/pbr_khr.frag
  // See: http://www.thetenthplanet.de/archives/1180
  vec2 uv = texcoord(material.normalTexcoord);
  vec3 tangentNormal = texture(normalTexture, uv).xyz * 2.0 - 1.0;
  vec3 q1 = dFdx(fragWorldPos);
  vec3 q2 = dFdy(fragWorldPos);
  vec2 st1 = dFdx(uv);
  vec2 st2 = dFdy(uv);
  vec3 N = normalize(fragNormal);
  vec3 T = normalize(q1 * st2.t - q2 * st1.t);
  vec3 B = -normalize(cross(N, T));
  mat3 TBN = mat3(T, B, N);

  return normalize(TBN * tangentNormal);
}

vec3 lightDirection(Light light) {
  if(light.type == 0 || light.type == 2) { // point and spot lights
    return normalize(light.position - fragWorldPos);
  } else if(light.type == 1) { // directional light
    return normalize(-light.direction);
  }

  return vec3(0.0); // should never be reached.
}

vec3 lightAttenuate(Light light) {
  float distance = length(light.position - fragWorldPos);
  float attenuation = 1.0 / (distance * distance);
  return light.color.rgb * light.intensity * attenuation;
}

vec3 lightRadiance(Light light, vec3 direction) {
  if(light.type == 0) { // point light
    return lightAttenuate(light);
  } else if(light.type == 1) { // directional light
    return light.color.rgb * light.intensity;
  } else if(light.type == 2) { // spot light
    float theta = dot(direction, normalize(-light.direction));
    // Spot lights only shine light in a specific conical area.
    // They have no effect outside of that area.
    if(theta > light.cutOff) {
      // Feather out the light as it approaches the edge of its cone.
      float intensity = (theta - light.cutOff) / (1.0 - light.cutOff);
      return lightAttenuate(light) * intensity;
    } else {
      return vec3(0.0);
    }
  }

  return vec3(0.0); // should never be reached.
}

// Useful resources I learned a lot from:
// https://learnopengl.com/PBR/Theory
// http://graphicrants.blogspot.com/2013/08/specular-brdf-reference.html
void main(void) {
  // The unit vector pointing from the fragment position to the viewer
  // position.
  vec3 viewDirection = normalize(cameraPosition - fragWorldPos);
  float metallic = materialMetallic();
  float roughness = materialRoughness();
  vec3 normal = materialNormal();
  vec3 reflection = reflect(-viewDirection, normal);
  // The "raw" albedo has an alpha channel which we need to preserve
  // so that we can apply the desired alpha blending method at the
  // end, but it is completely ignored for lighting calculations.
  vec4 rawAlbedo = materialAlbedo();
  vec3 albedo = rawAlbedo.rgb;
  // The ambient occlusion factor affects the degree to which ambient
  // lighting has an affect on the fragment.  Each element of the
  // vector is in the range [0, 1].
  vec3 ambientOcclusion = materialOcclusion();
  // The color that the fragment relects at zero incidence.  In other
  // words, the color reflected when looking directly at the fragment.
  // A material that is fully metallic will fully express its albedo
  // color.  A material that isn't metallic at all will have a gray
  // initial color.  The grayscale value of 0.04 apparently achieves a
  // good result so that's what we do.  Any other metallic value is
  // somewhere between both colors, as calculated by a linear
  // interpolation.
  vec3 baseColor = mix(vec3(0.04), albedo, metallic);
  // We will accumulate the results of all lights (ambient, emissive,
  // and user-specified direct lights) into this color vector.
  vec3 color = vec3(0.0);

  // Apply direct lighting.
  for(int i = 0; i < MAX_LIGHTS; ++i)
  {
    Light light = lights[i];

    if(!light.enabled) {
      continue;
    }

    // The unit vector pointing from the fragment position to the
    // light position.
    vec3 direction = lightDirection(light);
    // The dot product of the surface normal and the light direction
    // gives a value in the range [0, 1] that tells us how directly
    // light is hitting the fragment. 0 means the normal and light
    // direction form a >= 90 degree angle and thus the light should
    // have no effect.  1 means the light is directly hitting the
    // surface and the light should be applied with full intensity.
    float incidenceFactor = posDot(normal, direction);
    // The intensity of the light.
    vec3 radiance = lightRadiance(light, direction);

    // Skip all the expensive math below if the light will have no
    // effect on the result.
    if(incidenceFactor == 0.0 || radiance == vec3(0.0)) {
      continue;
    }

    // The halfway vector is named as such because it is halfway
    // between the view direction and the light direction.  It is a
    // significant vector for microfacet lighting calculations.
    vec3 halfwayVector = normalize(viewDirection + direction);
    // Apply the Trowbridge-Reitz GGX normal distribution function.
    // This function approximates the percentage of microfacets that
    // are aligned with the halfway vector.  Smooth objects have a lot
    // of aligned microfacets over a small area, thus the viewer sees
    // a bright spot.  Rough objects have few aligned microfacets over
    // a larger area, thus it appears dull.
    float normalDistribution = distributionGGX(normal, halfwayVector, roughness);
    // Apply the geometry distribution function using Smith's method.
    // This function approximates how much light is reflected based on
    // the fragment's roughness.  Rougher objects reflect less light
    // because the microfacets of the surface obstruct or overshadow
    // the reflected light rays.
    float geoFactor = geometrySmith(normal, viewDirection, direction, roughness);
    // Compute the fresnel factor via Schlick's approximation to get
    // the specular reflection coeffecient, i.e. the percentage of
    // light that is reflected.  Since this is a microfacet lighting
    // model, we need to use the dot product of the halfway vector and
    // the view direction to get the cos(theta) value, according to:
    // https://en.wikipedia.org/wiki/Schlick%27s_approximation
    //
    // I also found this article useful:
    // http://psgraphics.blogspot.com/2020/03/fresnel-equations-schlick-approximation.html
    vec3 reflectionFactor = fresnelSchlick(posDot(halfwayVector, direction), baseColor);
    // Refracted light is the leftover light that hasn't been
    // reflected.  One additional complicating factor is that metallic
    // surfaces don't refract light, so we must scale the refraction
    // factor based on how metallic the fragment is.
    vec3 refractionFactor = (vec3(1.0) - reflectionFactor) * (1.0 - metallic);
    // The dot product of the surface normal and the view direction
    // gives a value in the range [0, 1] that tells us how directly
    // the viewer is looking at the fragment.  The effect of specular
    // lighting changes based on the angle from the viewer to the
    // surface.
    float viewFactor = posDot(normal, viewDirection);
    // Specular and diffuse lighting are processed together using the
    // Cook-Torrance bidirectional reflectance distribution function
    // (BRDF.)
    //
    // http://www.codinglabs.net/article_physically_based_rendering_cook_torrance.aspx
    vec3 cookTorranceNumerator = normalDistribution * reflectionFactor * geoFactor;
    float cookTorranceDenominator = 4.0 * viewFactor * incidenceFactor;
    // We need to be careful to avoid a division by zero error, such
    // as when the specular factor is 0, so we clamp the denominator
    // to some very small value to safeguard ourselves.
    vec3 specular = cookTorranceNumerator / max(cookTorranceDenominator, 0.0001);
    // Apply Lambertian reflectance to get the diffuse lighting
    // factor.  https://en.wikipedia.org/wiki/Lambertian_reflectance
    vec3 diffuse = refractionFactor * albedo / PI;
    // The final light value is the combination of diffuse and
    // specular light scaled by the intensity of the light source and
    // the angle of incidence (how directly the light hit the
    // fragment.)
    color += (diffuse + specular) * radiance * incidenceFactor;
  }

  // The emissive texture says which fragments emit light.  We simply
  // add this light value to the color accumulator.
  color += materialEmissive();
  // Apply image based ambient lighting.  The affect of the ambient
  // light is dampened by the ambient occlusion factor.
  //
  // TODO: Use fancy PBR equations instead of these basic ones.
  float fresnel = pow(1.0 - clamp(dot(viewDirection, normal), 0.0, 1.0), 5);
  vec3 ambientDiffuse = texture(skybox, normal).rgb;
  vec3 ambientSpecular = textureLod(skybox, reflection, roughness * 7.0).rgb;
  color += (ambientDiffuse * albedo + ambientSpecular * fresnel) * ambientOcclusion;
  // Apply Reinhard tone mapping to convert our high dynamic range
  // color value to low dynamic range.  All of the lighting
  // calculations stacked on top of each other is likely to create
  // color channel values that exceed the [0, 1] range of OpenGL's
  // linear color space, i.e. high dynamic range.  If we did nothing,
  // the resulting color values would be clamped to that range and the
  // final image would be mostly white and washed out.
  color = toneMap(color);
  // Apply gamma correction.  The power of 2.2 is the rough average
  // gamma value of most monitors.  So, scaling the color by the
  // inverse, the power of 1/2.2, the color that people see on the
  // monitor is closer to what we intend to present with our original
  // linear color value.
  color = gammaCorrect(color);
  // Add alpha channel back in and apply the appropriate blend mode.
  // Yay, we're done!
  vec4 finalColor = applyAlpha(vec4(color, rawAlbedo.a));

#ifdef GLSL330
  fragColor = finalColor;
#else
  gl_FragColor = finalColor;
#endif
}