// -*- 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;
  bool occlusionTextureEnabled;
  int occlusionTexcoord;
  vec3 emissiveFactor;
  bool emissiveTextureEnabled;
  int emissiveTexcoord;
  int alphaMode;
  float alphaCutoff;
};

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

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

#ifdef GLSL330
out vec4 fragColor;
#endif

uniform Material material;
uniform Light lights[4];
uniform vec4 ambientLightColor;
uniform bool vertexColored;
uniform vec3 cameraPosition;
uniform sampler2D baseColorTexture;
uniform sampler2D metallicRoughnessTexture;
uniform sampler2D normalTexture;
uniform sampler2D occlusionTexture;
uniform sampler2D emissiveTexture;

const float PI = 3.14159265359;

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;
  }
}

vec4 sRGBtoLinear(vec4 srgb) {
  return vec4(pow(srgb.r, 2.2),
              pow(srgb.g, 2.2),
              pow(srgb.b, 2.2),
              srgb.a);
}

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

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

  return m;
}

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

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

  return r;
}

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

  if(material.baseColorTextureEnabled) {
    vec4 texColor = texture(baseColorTexture,
                            texcoord(material.baseColorTexcoord));
    color = sRGBtoLinear(texColor);
  }

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

  if(vertexColored) {
    color *= fragColor0;
  }

  return color;
}

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

  if(material.emissiveTextureEnabled) {
    vec4 texColor = texture(emissiveTexture,
                            texcoord(material.emissiveTexcoord));
    color = sRGBtoLinear(texColor);
  }

  return color * vec4(material.emissiveFactor, 1.0);
}

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

vec3 materialNormal() {
  vec3 normal;

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

  return normalize(normal);
}

void main(void) {
  float metallic = materialMetallic();
  float roughness = materialRoughness();
  vec3 N = materialNormal();
  vec3 V = normalize(cameraPosition - fragWorldPos);
  vec4 rawAlbedo = materialAlbedo();
  vec3 albedo = rawAlbedo.rgb;
  vec3 ao = materialOcclusion();
  vec3 F0 = mix(vec3(0.4), albedo, metallic);

  // reflectance equation
  vec3 Lo = vec3(0.0);
  for(int i = 0; i < 4; ++i)
  {
    Light light = lights[i];

    if(!light.enabled) {
      continue;
    }

    vec3 L;
    vec3 radiance;

    // calculate per-light radiance
    if(light.type == 0) { // point light
      L = normalize(light.position - fragWorldPos);
      float distance = length(light.position - fragWorldPos);
      float attenuation = 1.0 / (distance * distance);
      radiance = light.color.rgb * attenuation;
    } else if(light.type == 1) { // directional light
      L = normalize(-light.direction);
      radiance = light.color.rgb;
    } else if(light.type == 2) { // spotlight
      L = normalize(light.position - fragWorldPos);
      float theta = dot(L, normalize(-light.direction));

      if(theta > light.cutOff) {
        float distance = length(light.position - fragWorldPos);
        float attenuation = 1.0 / (distance * distance);
        radiance = light.color.rgb * attenuation;
      } else {
        continue;
      }
    }

    vec3 H = normalize(V + L);

    // 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;
  }

  // Add emissive color.
  Lo += materialEmissive().rgb;

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

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

  vec4 finalColor = applyAlpha(vec4(color, rawAlbedo.a));

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