// -----------------------------
// -- Based on Hazel PBR shader --
// -----------------------------
// Note: this shader is still very much in progress. There are likely many bugs and future additions that will go in.
//       Currently heavily updated.
//
// References upon which this is based:
// - Unreal Engine 4 PBR notes (https://blog.selfshadow.com/publications/s2013-shading-course/karis/s2013_pbs_epic_notes_v2.pdf)
// - Frostbite's SIGGRAPH 2014 paper (https://seblagarde.wordpress.com/2015/07/14/siggraph-2014-moving-frostbite-to-physically-based-rendering/)
// - Michał Siejak's PBR project (https://github.com/Nadrin)
// - My implementation from years ago in the Sparky engine (https://github.com/TheCherno/Sparky)
#type vertex
#version 430 core

layout(location = 0) in vec3 a_Position;
layout(location = 1) in vec3 a_Normal;
layout(location = 2) in vec3 a_Tangent;
layout(location = 3) in vec3 a_Binormal;
layout(location = 4) in vec2 a_TexCoord;

uniform mat4 u_ViewProjectionMatrix;
uniform mat4 u_ViewMatrix;
uniform mat4 u_Transform;

uniform mat4 u_LightSpaceMatrix;

out VertexOutput
{
	vec3 WorldPosition;
	vec3 Normal;
	vec2 TexCoord;
	mat3 WorldNormals;
	mat3 WorldTransform;
	vec3 Binormal;
	vec3 ViewPosition;
	vec4 FragPosLightSpace;
} vs_Output;

void main()
{
	vs_Output.WorldPosition = vec3(u_Transform * vec4(a_Position, 1.0));
	vs_Output.Normal = mat3(u_Transform) * a_Normal;
	vs_Output.TexCoord = vec2(a_TexCoord.x, 1.0 - a_TexCoord.y);
	vs_Output.WorldNormals = mat3(u_Transform) * mat3(a_Tangent, a_Binormal, a_Normal);
	vs_Output.WorldTransform = mat3(u_Transform);
	vs_Output.Binormal = a_Binormal;

	vs_Output.FragPosLightSpace = u_LightSpaceMatrix * u_Transform * vec4(a_Position, 1.0);

	vs_Output.ViewPosition = vec3(u_ViewMatrix * vec4(vs_Output.WorldPosition, 1.0));

	gl_Position = u_ViewProjectionMatrix * u_Transform * vec4(a_Position, 1.0);
}

#type fragment
#version 430 core

layout(location = 0) out vec4 outAlbedoMetal;
layout(location = 1) out vec4 outNormalRoughness;
layout(location = 2) out vec4 outEmissiveAO;
layout(location = 3) out vec4 outColor;
layout(location = 4) out vec4 outBloomColor;

const float PI = 3.141592;
const float Epsilon = 0.00001;

const int LightCount = 1;
const vec3 Fdielectric = vec3(0.04);

struct DirectionalLight {
	vec3 Direction;
	vec3 Radiance;
	float Intensity;
	bool CastShadows;
};

struct PointLight {
	vec3 Position;
	vec3 Radiance;
	float Intensity;
	float Range;
	bool CastShadows;
};

struct SpotLight {
	vec3 Position;
	vec3 Direction;
	vec3 Radiance;
	float Intensity;
	float Range;
	float InnerConeCos;
	float OuterConeCos;
	bool CastShadows;
};


in VertexOutput
{
	vec3 WorldPosition;
	vec3 Normal;
	vec2 TexCoord;
	mat3 WorldNormals;
	mat3 WorldTransform;
	vec3 Binormal;
	vec3 ViewPosition;
	vec4 FragPosLightSpace;
} vs_Input;


uniform bool u_GBufferMode;

uniform DirectionalLight u_DirectionalLights;

uniform int u_PointLightCount;
uniform PointLight u_PointLights;

uniform int u_SpotLightCount;
uniform SpotLight u_SpotLights;


uniform vec3 u_CameraPosition;

// PBR
uniform sampler2D u_AlbedoTexture;
uniform sampler2D u_NormalTexture;
uniform sampler2D u_MetalnessTexture;
uniform sampler2D u_RoughnessTexture;

// environment
uniform samplerCube u_EnvRadianceTex;
uniform samplerCube u_EnvIrradianceTex;

// BRDF LUT
uniform sampler2D u_BRDFLUTTexture;

uniform float u_IBLContribution;
uniform float u_BloomThreshold;
uniform float u_EnvMapRotation;

// baseColor
uniform vec3  u_AlbedoColor;
uniform float u_Metalness;
uniform float u_Roughness;

// textureToggle
uniform float u_AlbedoTexToggle;
uniform float u_NormalTexToggle;
uniform float u_MetalnessTexToggle;
uniform float u_RoughnessTexToggle;

// shadow
uniform sampler2D u_ShadowMap;
uniform float u_ShadowBias;
uniform float u_ShadowSoftness;
uniform float u_ShadowIntensity;
uniform int u_ShadowEnabled;

// Emissive
uniform sampler2D u_EmissiveTexture;
uniform float u_EmissiveTexToggle;
uniform vec3  u_EmissiveColor;
uniform float u_EmissiveIntensity;

struct PBRParameters
{
	vec3 Albedo;
	float Roughness;
	float Metalness;
	vec3 Normal;
	vec3 View;
	float NdotV;
};

PBRParameters m_Params;

// ---------- PBR param func ----------
float ndfGGX(float cosLh, float roughness)
{
	float alpha = roughness * roughness;
	float alphaSq = alpha * alpha;
	float denom = (cosLh * cosLh) * (alphaSq - 1.0) + 1.0;
	return alphaSq / (PI * denom * denom);
}

float GeometrySchlickGGX(float NdotV, float roughness)
{
	float r = (roughness + 1.0);
	float k = (r * r) / 8.0;
	float nom = NdotV;
	float denom = NdotV * (1.0 - k) + k;
	return nom / 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;
}

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

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

// ---------- direction light ----------
vec3 ComputeDirectionalLight(DirectionalLight light, vec3 F0, PBRParameters params)
{
	vec3 L = normalize(-light.Direction);
	vec3 Lradiance = light.Radiance * light.Intensity;

	vec3 Lh = normalize(L + params.View);
	float cosLi = max(0.0, dot(params.Normal, L));
	float cosLh = max(0.0, dot(params.Normal, Lh));

	vec3 F = fresnelSchlick(F0, max(0.0, dot(Lh, params.View)));
	float D = ndfGGX(cosLh, params.Roughness);
	float G = GeometrySmith(params.Normal, params.View, L, params.Roughness);

	vec3 kd = (1.0 - F) * (1.0 - params.Metalness);
	vec3 diffuseBRDF = kd * params.Albedo;
	vec3 specularBRDF = (F * D * G) / max(Epsilon, 4.0 * cosLi * params.NdotV);

	return (diffuseBRDF + specularBRDF) * Lradiance * cosLi;
}

vec3 ComputePointLight(PointLight light, vec3 F0, PBRParameters params, vec3 worldPos)
{
	vec3 lightVec = light.Position - worldPos;
	float dist = length(lightVec);
	if (dist > light.Range) return vec3(0.0);

	vec3 L = lightVec / dist;
	vec3 Lradiance = light.Radiance * light.Intensity;

	// 距离衰减：通常使用平方衰减，但为避免分母为零，加一个小值
	float attenuation = 1.0 / (dist * dist + 0.0001);

	// 可选：范围平滑衰减
	float rangeFactor = clamp(1.0 - (dist / light.Range), 0.0, 1.0);
	rangeFactor = rangeFactor * rangeFactor; // 平滑
	attenuation *= rangeFactor;

	vec3 Lh = normalize(L + params.View);
	float cosLi = max(0.0, dot(params.Normal, L));
	float cosLh = max(0.0, dot(params.Normal, Lh));

	vec3 F = fresnelSchlick(F0, max(0.0, dot(Lh, params.View)));
	float D = ndfGGX(cosLh, params.Roughness);
	float G = GeometrySmith(params.Normal, params.View, L, params.Roughness);

	vec3 kd = (1.0 - F) * (1.0 - params.Metalness);
	vec3 diffuseBRDF = kd * params.Albedo;
	vec3 specularBRDF = (F * D * G) / max(Epsilon, 4.0 * cosLi * params.NdotV);

	return (diffuseBRDF + specularBRDF) * Lradiance * cosLi * attenuation;
}

vec3 ComputeSpotLight(SpotLight light, vec3 F0, PBRParameters params, vec3 worldPos)
{
	vec3 lightVec = light.Position - worldPos;
	float dist = length(lightVec);
	if (dist > light.Range) return vec3(0.0);

	vec3 L = lightVec / dist;
	vec3 Lradiance = light.Radiance * light.Intensity;

	// 距离衰减
	float attenuation = 1.0 / (dist * dist + 0.0001);
	float rangeFactor = clamp(1.0 - (dist / light.Range), 0.0, 1.0);
	rangeFactor = rangeFactor * rangeFactor;
	attenuation *= rangeFactor;

	// 角度衰减（聚光锥）
	float cosAngle = dot(-L, normalize(light.Direction)); // 光方向指向外，所以用 -L
	if (cosAngle < light.OuterConeCos) return vec3(0.0);
	float angleFalloff = (cosAngle - light.OuterConeCos) / (light.InnerConeCos - light.OuterConeCos);
	angleFalloff = clamp(angleFalloff, 0.0, 1.0);
	attenuation *= angleFalloff;

	vec3 Lh = normalize(L + params.View);
	float cosLi = max(0.0, dot(params.Normal, L));
	float cosLh = max(0.0, dot(params.Normal, Lh));

	vec3 F = fresnelSchlick(F0, max(0.0, dot(Lh, params.View)));
	float D = ndfGGX(cosLh, params.Roughness);
	float G = GeometrySmith(params.Normal, params.View, L, params.Roughness);

	vec3 kd = (1.0 - F) * (1.0 - params.Metalness);
	vec3 diffuseBRDF = kd * params.Albedo;
	vec3 specularBRDF = (F * D * G) / max(Epsilon, 4.0 * cosLi * params.NdotV);

	return (diffuseBRDF + specularBRDF) * Lradiance * cosLi * attenuation;
}

vec3 Lighting(vec3 F0)
{
	vec3 result = vec3(0.0);
	for(int i = 0; i < LightCount; i++)
	{
		vec3 Li = u_DirectionalLights.Direction;
		vec3 Lradiance = u_DirectionalLights.Radiance * u_DirectionalLights.Intensity;
		vec3 Lh = normalize(Li + m_Params.View);

		float cosLi = max(0.0, dot(m_Params.Normal, Li));
		float cosLh = max(0.0, dot(m_Params.Normal, Lh));

		vec3 F = fresnelSchlick(F0, max(0.0, dot(Lh, m_Params.View)));
		float D = ndfGGX(cosLh, m_Params.Roughness);
		float G = GeometrySmith(m_Params.Normal, m_Params.View, Li, m_Params.Roughness);

		vec3 kd = (1.0 - F) * (1.0 - m_Params.Metalness);
		vec3 diffuseBRDF = kd * m_Params.Albedo;

		vec3 specularBRDF = (F * D * G) / max(Epsilon, 4.0 * cosLi * m_Params.NdotV);

		result += (diffuseBRDF + specularBRDF) * Lradiance * cosLi;
	}
	return result;
}

// ---------- IBL ----------
vec3 RotateVectorAboutY(float angle, vec3 vec)
{
	angle = radians(angle);
	mat3 rotationMatrix = mat3(
		vec3(cos(angle), 0.0, sin(angle)),
		vec3(0.0, 1.0, 0.0),
		vec3(-sin(angle), 0.0, cos(angle))
	);
	return rotationMatrix * vec;
}

vec3 IBL(vec3 F0, vec3 Lr)
{
	vec3 irradiance = texture(u_EnvIrradianceTex, m_Params.Normal).rgb;
	vec3 F = fresnelSchlickRoughness(F0, m_Params.NdotV, m_Params.Roughness);
	vec3 kd = (1.0 - F) * (1.0 - m_Params.Metalness);
	vec3 diffuseIBL = m_Params.Albedo * irradiance;

	int u_EnvRadianceTexLevels = textureQueryLevels(u_EnvRadianceTex);
	vec3 specularIrradiance = textureLod(
		u_EnvRadianceTex,
		RotateVectorAboutY(u_EnvMapRotation, Lr),
		m_Params.Roughness * u_EnvRadianceTexLevels
	).rgb;

	vec2 specularBRDF = texture(u_BRDFLUTTexture, vec2(m_Params.NdotV, 1.0 - m_Params.Roughness)).rg;
	vec3 specularIBL = specularIrradiance * (F * specularBRDF.x + specularBRDF.y);

	return kd * diffuseIBL + specularIBL;
}

// shadow

float calculateShadow(vec4 fragPosLightSpace, vec3 normal, vec3 lightDir)
{
	// Perspective divide
	vec3 projCoords = fragPosLightSpace.xyz / fragPosLightSpace.w;

	// Transform to [0,1] range
	projCoords = projCoords * 0.5 + 0.5;

	// If outside shadow map bounds, assume no shadow
	if(projCoords.z > 1.0 || projCoords.x < 0.0 || projCoords.x > 1.0 || projCoords.y < 0.0 || projCoords.y > 1.0)
	return 0.0;

	// Get closest depth value from light's perspective
	float closestDepth = texture(u_ShadowMap, projCoords.xy).r;
	float currentDepth = projCoords.z;

	// Calculate bias based on surface angle
	float bias = max(u_ShadowBias * (1.0 - dot(normal, lightDir)), u_ShadowBias * 0.1);

	// PCF (Percentage Closer Filtering) for soft shadows
	float shadow = 0.0;
	vec2 texelSize = 1.0 / textureSize(u_ShadowMap, 0);
	int pcfRange = int(u_ShadowSoftness);
	int sampleCount = 0;

	for(int x = -pcfRange; x <= pcfRange; ++x)
	{
		for(int y = -pcfRange; y <= pcfRange; ++y)
		{
			float pcfDepth = texture(u_ShadowMap, projCoords.xy + vec2(x, y) * texelSize).r;
			shadow += currentDepth - bias > pcfDepth ? 1.0 : 0.0;
			sampleCount++;
		}
	}
	shadow /= float(sampleCount);

	return shadow;
}


float ComputeShadow(vec4 fragPosLightSpace, float NdotL)
{
	if (u_ShadowEnabled == 0) return 1.0;

	vec3 projCoords = fragPosLightSpace.xyz / fragPosLightSpace.w;
	projCoords = projCoords * 0.5 + 0.5;

	if (projCoords.x < 0.0 || projCoords.x > 1.0 ||
		projCoords.y < 0.0 || projCoords.y > 1.0 ||
		projCoords.z > 1.0) return 1.0;

	float closestDepth = texture(u_ShadowMap, projCoords.xy).r;
	float currentDepth = projCoords.z;

	float bias = max(u_ShadowBias * (1.0 - NdotL), u_ShadowBias * 0.5);

	float shadow = (currentDepth - bias) > closestDepth ? 1.0 : 0.0;
	return mix(1.0, 1.0 - u_ShadowIntensity, shadow);
}


void main()
{
	// === 1. 采样基础属性（所有模式都需要） ===
	vec4 albedoWithAlpha = texture(u_AlbedoTexture, vs_Input.TexCoord);
	vec3 albedo = u_AlbedoTexToggle > 0.5 ? albedoWithAlpha.rgb : u_AlbedoColor;
	float alpha = u_AlbedoTexToggle > 0.5 ? albedoWithAlpha.a : 1.0;

	float metallic = u_MetalnessTexToggle > 0.5 ? texture(u_MetalnessTexture, vs_Input.TexCoord).r : u_Metalness;
	float roughness = u_RoughnessTexToggle > 0.5 ? texture(u_RoughnessTexture, vs_Input.TexCoord).r : u_Roughness;
	roughness = max(roughness, 0.05);

	// === 2. 法线计算（世界空间） ===
	vec3 normal = normalize(vs_Input.Normal);
	if (u_NormalTexToggle > 0.5)
	{
		vec3 tangentNormal = texture(u_NormalTexture, vs_Input.TexCoord).rgb * 2.0 - 1.0;
		normal = normalize(vs_Input.WorldNormals * tangentNormal);
	}

	// === 3. 自发光计算 ===
	vec3 emissive = u_EmissiveColor;
	if (u_EmissiveTexToggle > 0.5)
	emissive = texture(u_EmissiveTexture, vs_Input.TexCoord).rgb;
	emissive *= u_EmissiveIntensity;

	// === 4. GBuffer 模式：直接输出到多个目标 ===
	if (u_GBufferMode)
	{
		outAlbedoMetal = vec4(albedo, metallic);
		outNormalRoughness = vec4(normal * 0.5 + 0.5, roughness);
		outEmissiveAO = vec4(emissive, 1.0); // AO 暂设为 1.0
		return; // 提前结束
	}

	// === 5. 非 GBuffer 模式：继续 PBR 光照计算 ===
	// 填充 PBRParameters
	m_Params.Albedo = albedo;
	m_Params.Metalness = metallic;
	m_Params.Roughness = roughness;
	m_Params.Normal = normal;
	m_Params.View = normalize(u_CameraPosition - vs_Input.WorldPosition);
	m_Params.NdotV = max(dot(m_Params.Normal, m_Params.View), 0.0);

	vec3 Lr = 2.0 * m_Params.NdotV * m_Params.Normal - m_Params.View;

	vec3 F0 = mix(Fdielectric, m_Params.Albedo, m_Params.Metalness);

	// Shadow
	float shadowFactor = 1.0;
	if (u_ShadowEnabled > 0.5) {
		float shadow = calculateShadow(vs_Input.FragPosLightSpace, m_Params.Normal, u_DirectionalLights.Direction);
		shadowFactor = 1.0 - shadow;
	}

	// directional light with with shadow
	vec3 lightContribution = u_DirectionalLights.Intensity > 0.0 ? Lighting(F0) * shadowFactor : vec3(0.0);

	if(u_PointLightCount > 0)
		lightContribution += ComputePointLight(u_PointLights, F0, m_Params, vs_Input.WorldPosition);

	if(u_SpotLightCount > 0)
		lightContribution += ComputeSpotLight(u_SpotLights, F0, m_Params, vs_Input.WorldPosition);

	vec3 iblContribution = IBL(F0, Lr) * u_IBLContribution;

	vec3 finalRGB = lightContribution + iblContribution + emissive;
	vec4 finalColor = vec4(finalRGB, alpha);

	outColor = finalColor;

	// Bloom
	float brightness = dot(finalColor.rgb, vec3(0.2126, 0.7152, 0.0722));
	outBloomColor = brightness > u_BloomThreshold ? finalColor : vec4(0.0, 0.0, 0.0, 1.0);
}