// -*- mode: c -*-

// Heavily based upon: https://learnopengl.com/PBR/Lighting

struct Material {
  vec3 baseColorFactor;
  bool baseColorTextureEnabled;
  int baseColorTexcoord;
  float metallicFactor;
  float roughnessFactor;
  bool metallicRoughnessTextureEnabled;
  int metallicRoughnessTexcoord;
  vec3 normalFactor;
  bool normalTextureEnabled;
  int normalTexcoord;
  vec3 occlusionFactor;
  bool occlusionTextureEnabled;
  int occlusionTexcoord;
  vec3 emissiveFactor;
  bool emissiveTextureEnabled;
  int emissiveTexcoord;
  int alphaMode;
  float alphaCutoff;
};

#ifdef GLSL120
attribute vec3 fragWorldPos;
attribute vec3 fragCamPos;
attribute vec3 fragNormal;
attribute vec2 fragTexcoord0;
attribute vec2 fragTexcoord1
attribute vec4 fragColor0;
#else
in vec3 fragWorldPos;
in vec3 fragCamPos;
in vec3 fragNormal;
in vec2 fragTexcoord0;
in vec2 fragTexcoord1;
in vec4 fragColor0;
#endif

#ifdef GLSL330
out vec4 fragColor;
#endif

uniform Material material;
uniform bool vertexColored;
uniform sampler2D baseColorTexture;
uniform sampler2D metallicRoughnessTexture;
uniform sampler2D normalTexture;
uniform sampler2D occlusionTexture;
uniform sampler2D emissiveTexture;

const float PI = 3.14159265359;

vec4 sampleTexture(sampler2D tex, bool enabled, int texcoord, vec3 factor, vec4 defaultColor) {
  if(enabled && texcoord == 0) {
    return texture(tex, fragTexcoord0) * vec4(factor, 1.0);
  } else if(enabled && texcoord == 1) {
    return texture(tex, fragTexcoord1) * vec4(factor, 1.0);
  } else {
    return defaultColor;
  }
}

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

float DistributionGGX(vec3 N, vec3 H, float roughness)
{
    float a      = roughness*roughness;
    float a2     = a*a;
    float NdotH  = max(dot(N, H), 0.0);
    float NdotH2 = NdotH*NdotH;

    float num   = a2;
    float denom = (NdotH2 * (a2 - 1.0) + 1.0);
    denom = PI * denom * denom;

    return num / denom;
}

float GeometrySchlickGGX(float NdotV, float roughness)
{
    float r = (roughness + 1.0);
    float k = (r*r) / 8.0;

    float num   = NdotV;
    float denom = NdotV * (1.0 - k) + k;

    return num / denom;
}

float GeometrySmith(vec3 N, vec3 V, vec3 L, float roughness)
{
    float NdotV = max(dot(N, V), 0.0);
    float NdotL = max(dot(N, L), 0.0);
    float ggx2  = GeometrySchlickGGX(NdotV, roughness);
    float ggx1  = GeometrySchlickGGX(NdotL, roughness);

    return ggx1 * ggx2;
}

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) {
  if(i == 0) {
    return fragTexcoord0;
  } else {
    return fragTexcoord1;
  }
}

float materialMetallic() {
  float m = material.metallicFactor;

  if(material.metallicRoughnessTextureEnabled) {
    m *= texture2D(metallicRoughnessTexture,
                   texcoord(material.metallicRoughnessTexcoord)).b;
  }

  return m;
}

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

  if(material.metallicRoughnessTextureEnabled) {
    r *= texture2D(metallicRoughnessTexture,
                   texcoord(material.metallicRoughnessTexcoord)).g;
  }

  return r;
}

vec4 materialAlbedo() {
  vec4 color = vec4(0.0, 0.0, 1.0, 1.0);

  if(material.baseColorTextureEnabled) {
    color = texture2D(baseColorTexture,
                      texcoord(material.baseColorTexcoord));
  }

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

  if(vertexColored) {
    color *= fragColor0;
  }

  return color;
}

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

  if(material.occlusionTextureEnabled) {
    color = texture2D(occlusionTexture,
                      texcoord(material.occlusionTexcoord)).rgb;
  }

  return color * material.occlusionFactor;
}

vec3 materialNormal() {
  vec3 normal;

  if(material.normalTextureEnabled) {
    normal = texture2D(normalTexture,
                       texcoord(material.normalTexcoord)).rgb * 2.0 - 1.0;
  } else {
    normal = fragNormal;
  }

  return normalize(normal);
}

struct PointLight {
  vec3 position;
  vec3 color;
};

PointLight lights[1];
vec3 ambientLightColor = vec3(0.05);

void main(void) {
  // TODO: Support user supplied lights.
  lights[0].position = vec3(0.0, 0.2, 2.0);
  lights[0].color = vec3(1.0, 1.0, 1.0);

  float metallic = materialMetallic();
  float roughness = materialRoughness();
  vec3 N = materialNormal();
  vec3 V = normalize(-fragCamPos - fragWorldPos);
  vec3 albedo = materialAlbedo().rgb;
  vec3 ao = materialOcclusion();
  vec3 F0 = mix(vec3(0.4), albedo, metallic);

  // reflectance equation
  vec3 Lo = vec3(0.0);
  for(int i = 0; i < 1; ++i)
  {
    PointLight light = lights[i];
    // calculate per-light radiance
    vec3 L = normalize(light.position - fragWorldPos);
    vec3 H = normalize(V + L);
    float distance = length(light.position - fragWorldPos);
    float attenuation = 1.0 / (distance * distance);
    vec3 radiance = light.color * attenuation;

    // cook-torrance brdf
    float NDF = DistributionGGX(N, H, roughness);
    float G = GeometrySmith(N, V, L, roughness);
    vec3 F = fresnelSchlick(max(dot(H, V), 0.0), F0);

    vec3 kS = F;
    vec3 kD = vec3(1.0) - kS;
    kD *= 1.0 - metallic;

    vec3 numerator = NDF * G * F;
    float denominator = 4.0 * max(dot(N, V), 0.0) * max(dot(N, L), 0.0);
    vec3 specular = numerator / max(denominator, 0.001);

    // add to outgoing radiance Lo
    float NdotL = max(dot(N, L), 0.0);
    Lo += (kD * albedo / PI + specular) * radiance * NdotL;
  }

  // TODO: Process emissive color properly.
  Lo += sampleTexture(emissiveTexture,
                      material.emissiveTextureEnabled,
                      material.emissiveTexcoord,
                      material.emissiveFactor,
                      vec4(0.0, 0.0, 0.0, 0.0)).xyz;

  // Apply ambient lighting.
  vec3 ambient = ambientLightColor * albedo * ao;
  vec3 color = ambient + Lo;

  // Apply HDR.
  color = color / (color + vec3(1.0));
  color = pow(color, vec3(1.0 / 2.2));

  // TODO: Preserve alpha channel throughout lighting.
  vec4 finalColor = applyAlpha(vec4(color, 1.0));

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