#version 330 core
precision mediump float;

// osmium.head look: the standard Verso::Model3d PBR shading (so it stays compatible with what
// Model3d::render() feeds it), then pushed HARD to a two-tone "mostly light or dark" look and
// dusted with animated grain. The key light itself is snapped to a new direction on every beat by
// the CPU side, so the lit/dark split flips with the music.

in vec2 ex_Uv;
in vec3 ex_FragPos;
in vec3 ex_Normal;
in vec4 ex_Tangent;

out vec4 fragColor;

struct Material {
	sampler2D baseColorTexture;
	sampler2D normalTexture;
	sampler2D specularTexture;
	vec4 baseColorFactor;
	vec3 specularFactor;
	float roughness;
	float metallic;
	int workflow;
};

struct DirectionalLight {
	vec3 direction;

	vec3 ambient;
	vec3 diffuse;
	vec3 specular;
};

struct PointLight {
	vec3 position;

	float constant;
	float linear;
	float quadratic;

	vec3 ambient;
	vec3 diffuse;
	vec3 specular;
};

struct SpotLight {
	vec3 position;
	vec3 direction;

	float cutOff;
	float outerCutOff;

	float constant;
	float linear;
	float quadratic;

	vec3 ambient;
	vec3 diffuse;
	vec3 specular;
};

#define MAX_DIRECTIONAL_LIGHTS 2
#define MAX_POINT_LIGHTS 10
#define MAX_SPOT_LIGHTS 10

uniform vec3 viewPos;
uniform Material material;

uniform DirectionalLight directionalLights[MAX_DIRECTIONAL_LIGHTS];
uniform int directionalLightsCount;

uniform PointLight pointLights[MAX_POINT_LIGHTS];
uniform int pointLightsCount;

uniform SpotLight spotLights[MAX_SPOT_LIGHTS];
uniform int spotLightsCount;

uniform int uTonemap; // 0 = none, 1 = reinhard, 2 = aces
uniform float alpha;

// osmium.head extras
uniform float uNoise; // grain amplitude added to the colour
uniform float uHarsh; // [0,1] mix toward the two-tone (1 = fully light-or-dark)
uniform float uTime;  // track seconds (animates the grain)

const float PI = 3.14159265359;


float distributionGGX(float NdotH, float roughness)
{
	float a = roughness * roughness;
	float a2 = a * a;
	float denom = NdotH * NdotH * (a2 - 1.0) + 1.0;
	return a2 / max(PI * denom * denom, 1e-6);
}


float geometrySchlickGGX(float NdotX, float roughness)
{
	float r = roughness + 1.0;
	float k = (r * r) / 8.0;
	return NdotX / (NdotX * (1.0 - k) + k);
}


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


vec3 shadeLight(vec3 N, vec3 V, vec3 L, vec3 lightDiffuse, vec3 lightSpecular,
                vec3 albedo, vec3 F0, float roughness, float metallic)
{
	float NdotL = dot(N, L);
	if (NdotL <= 0.0) {
		return vec3(0.0);
	}
	NdotL = min(NdotL, 1.0);
	float NdotV = max(dot(N, V), 1e-4);
	vec3 H = normalize(V + L);

	float D = distributionGGX(max(dot(N, H), 0.0), roughness);
	float G = geometrySchlickGGX(NdotV, roughness) * geometrySchlickGGX(NdotL, roughness);
	vec3 F = fresnelSchlick(max(dot(H, V), 0.0), F0);

	vec3 specular = (D * G * F) * PI / max(4.0 * NdotV * NdotL, 1e-4);

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

	return (kD * albedo * lightDiffuse + specular * lightSpecular) * NdotL;
}


vec3 acesFilm(vec3 x)
{
	const float a = 2.51;
	const float b = 0.03;
	const float c = 2.43;
	const float d = 0.59;
	const float e = 0.14;
	return clamp((x * (a * x + b)) / (x * (c * x + d) + e), 0.0, 1.0);
}


float hash(vec2 p)
{
	p = fract(p * vec2(123.34, 456.21));
	p += dot(p, p + 45.32);
	return fract(p.x * p.y);
}


void main()
{
	vec4 baseTex = texture(material.baseColorTexture, ex_Uv);
	vec3 albedo = baseTex.rgb * material.baseColorFactor.rgb;
	float outAlpha = baseTex.a * material.baseColorFactor.a * alpha;

	vec3 N = normalize(ex_Normal);
	vec3 T = ex_Tangent.xyz - N * dot(N, ex_Tangent.xyz);
	float tLen = length(T);
	if (tLen > 1e-5) {
		T /= tLen;
		vec3 B = cross(N, T) * ex_Tangent.w;
		vec3 nTex = texture(material.normalTexture, ex_Uv).rgb * 2.0 - 1.0;
		N = normalize(mat3(T, B, N) * nTex);
	}

	vec3 V = normalize(viewPos - ex_FragPos);

	vec3 F0;
	float roughness;
	float metallic;
	if (material.workflow == 1) {
		vec3 specTex = texture(material.specularTexture, ex_Uv).rgb;
		F0 = specTex * material.specularFactor;
		roughness = material.roughness;
		metallic = 0.0;
	} else {
		vec3 mrTex = texture(material.specularTexture, ex_Uv).rgb;
		roughness = material.roughness * mrTex.g;
		metallic = clamp(material.metallic * mrTex.b, 0.0, 1.0);
		F0 = mix(vec3(0.04), albedo, metallic);
	}
	roughness = clamp(roughness, 0.045, 1.0);

	vec3 result = vec3(0.0);

	for (int i = 0; i < directionalLightsCount; i++) {
		vec3 L = normalize(-directionalLights[i].direction);
		result += directionalLights[i].ambient * albedo;
		result += shadeLight(N, V, L, directionalLights[i].diffuse, directionalLights[i].specular,
		                     albedo, F0, roughness, metallic);
	}

	for (int i = 0; i < pointLightsCount; i++) {
		vec3 toLight = pointLights[i].position - ex_FragPos;
		float distance = length(toLight);
		vec3 L = toLight / max(distance, 1e-6);
		float attenuation = 1.0 / max(pointLights[i].constant + pointLights[i].linear * distance +
		                              pointLights[i].quadratic * (distance * distance), 1e-6);
		result += pointLights[i].ambient * albedo * attenuation;
		result += shadeLight(N, V, L, pointLights[i].diffuse, pointLights[i].specular,
		                     albedo, F0, roughness, metallic) * attenuation;
	}

	for (int i = 0; i < spotLightsCount; i++) {
		vec3 toLight = spotLights[i].position - ex_FragPos;
		float distance = length(toLight);
		vec3 L = toLight / max(distance, 1e-6);
		float attenuation = 1.0 / max(spotLights[i].constant + spotLights[i].linear * distance +
		                              spotLights[i].quadratic * (distance * distance), 1e-6);
		float theta = dot(L, normalize(-spotLights[i].direction));
		float epsilon = max(spotLights[i].cutOff - spotLights[i].outerCutOff, 1e-6);
		float intensity = clamp((theta - spotLights[i].outerCutOff) / epsilon, 0.0, 1.0);
		result += spotLights[i].ambient * albedo * attenuation;
		result += shadeLight(N, V, L, spotLights[i].diffuse, spotLights[i].specular,
		                     albedo, F0, roughness, metallic) * attenuation * intensity;
	}

	// Tonemap + gamma to display space.
	vec3 color = result;
	if (uTonemap == 1) {
		color = color / (1.0 + color);
	} else if (uTonemap == 2) {
		color = acesFilm(color);
	}
	color = pow(max(color, vec3(0.0)), vec3(1.0 / 2.2));

	// Harsh two-tone: collapse toward almost-black or almost-white by luminance.
	float lum = dot(color, vec3(0.299, 0.587, 0.114));
	float lit = smoothstep(0.34, 0.50, lum);
	vec3 twoTone = mix(vec3(0.02), vec3(0.98), lit);
	color = mix(color, twoTone, clamp(uHarsh, 0.0, 1.0));

	// Animated grain (screen-space), amount = uNoise.
	float g = hash(gl_FragCoord.xy + vec2(uTime * 137.0, uTime * 91.0)) * 2.0 - 1.0;
	color += g * uNoise;

	color = clamp(color, 0.0, 1.0);
	fragColor = vec4(color, outAlpha);
}
