[nzsl_version("1.0")]
module PhongMaterial;

import InstanceData from Engine.InstanceData;
import LightData from Engine.LightData;
import ViewerData from Engine.ViewerData;

// Basic material options
option HasDiffuseTexture: bool = false;
option HasAlphaTexture: bool = false;
option AlphaTest: bool = false;

// Phong material options
option HasEmissiveTexture: bool = false;
option HasHeightTexture: bool = false;
option HasNormalTexture: bool = false;
option HasSpecularTexture: bool = false;

// Billboard related options
option Billboard: bool = false;
option BillboardCenterLocation: i32 = -1;
option BillboardColorLocation: i32 = -1;
option BillboardSizeRotLocation: i32 = -1;

// Vertex declaration related options
option ColorLocation: i32 = -1;
option NormalLocation: i32 = -1;
option PosLocation: i32;
option TangentLocation: i32 = -1;
option UvLocation: i32 = -1;

const HasNormal = (NormalLocation >= 0);
const HasVertexColor = (ColorLocation >= 0);
const HasColor = (HasVertexColor || Billboard);
const HasTangent = (TangentLocation >= 0);
const HasUV = (UvLocation >= 0);
const HasNormalMapping = HasNormalTexture && HasNormal && HasTangent;

[layout(std140)]
struct MaterialSettings
{
	// BasicSettings
	AlphaThreshold: f32,
	DiffuseColor: vec4[f32],

	// PhongSettings
	AmbientColor: vec3[f32],
	SpecularColor: vec3[f32],
	Shininess: f32,
}

// TODO: Add enums
const DirectionalLight = 0;
const PointLight = 1;
const SpotLight = 2;

external
{
	[binding(0)] settings: uniform[MaterialSettings],
	[binding(1)] MaterialDiffuseMap: sampler2D[f32],
	[binding(2)] MaterialAlphaMap: sampler2D[f32],
	[binding(3)] TextureOverlay: sampler2D[f32],
	[binding(4)] instanceData: uniform[InstanceData],
	[binding(5)] viewerData: uniform[ViewerData],
	[binding(6)] lightData: uniform[LightData],
	[binding(7)] MaterialEmissiveMap: sampler2D[f32],
	[binding(8)] MaterialHeightMap: sampler2D[f32],
	[binding(9)] MaterialNormalMap: sampler2D[f32],
	[binding(10)] MaterialSpecularMap: sampler2D[f32],
}

struct VertToFrag
{
	[location(0)] worldPos: vec3[f32],
	[location(1), cond(HasUV)] uv: vec2[f32],
	[location(2), cond(HasColor)] color: vec4[f32],
	[location(3), cond(HasNormal)] normal: vec3[f32],
	[location(4), cond(HasNormalMapping)] tangent: vec3[f32],
	[builtin(position)] position: vec4[f32],
}

// Fragment stage
struct FragOut
{
	[location(0)] RenderTarget0: vec4[f32]
}

[entry(frag)]
fn main(input: VertToFrag) -> FragOut
{
	let diffuseColor = settings.DiffuseColor;

	const if (HasUV)
		diffuseColor *= TextureOverlay.Sample(input.uv);

	const if (HasColor)
		diffuseColor *= input.color;

	const if (HasDiffuseTexture)
		diffuseColor *= MaterialDiffuseMap.Sample(input.uv);

	const if (HasAlphaTexture)
		diffuseColor.w *= MaterialAlphaMap.Sample(input.uv).x;

	const if (AlphaTest)
	{
		if (diffuseColor.w < settings.AlphaThreshold)
			discard;
	}

	const if (HasNormal)
	{
		let lightAmbient = vec3[f32](0.0, 0.0, 0.0);
		let lightDiffuse = vec3[f32](0.0, 0.0, 0.0);
		let lightSpecular = vec3[f32](0.0, 0.0, 0.0);

		let eyeVec = normalize(viewerData.eyePosition - input.worldPos);

		let normal: vec3[f32];
		const if (HasNormalMapping)
		{
			let N = normalize(input.normal);
			let T = normalize(input.tangent);
			let B = cross(N, T);
			let tbnMatrix = mat3[f32](T, B, N);

			normal = normalize(tbnMatrix * (MaterialNormalMap.Sample(input.uv).xyz * 2.0 - vec3[f32](1.0, 1.0, 1.0)));
		}
		else
			normal = normalize(input.normal);

		for i in u32(0) -> lightData.lightCount
		{
			let light = lightData.lights[i];

			let lightAmbientFactor = light.factor.x;
			let lightDiffuseFactor = light.factor.y;

			// TODO: Add switch instruction
			if (light.type == DirectionalLight)
			{
				let lightDir = light.parameter1.xyz;

				lightAmbient += light.color.rgb * lightAmbientFactor * settings.AmbientColor;

				let lambert = max(dot(normal, -lightDir), 0.0);

				lightDiffuse += lambert * light.color.rgb * lightDiffuseFactor;

				let reflection = reflect(lightDir, normal);
				let specFactor = max(dot(reflection, eyeVec), 0.0);
				specFactor = pow(specFactor, settings.Shininess);

				lightSpecular += specFactor * light.color.rgb;
			}
			else if (light.type == PointLight)
			{
				let lightPos = light.parameter1.xyz;
				let lightInvRadius = light.parameter1.w;

				let lightToPos = input.worldPos - lightPos;
				let dist = length(lightToPos);
				let lightToPosNorm = lightToPos / max(dist, 0.0001);

				let attenuationFactor = max(1.0 - dist * lightInvRadius, 0.0);

				lightAmbient += attenuationFactor * light.color.rgb * lightAmbientFactor * settings.AmbientColor;

				let lambert = max(dot(normal, -lightToPosNorm), 0.0);

				lightDiffuse += attenuationFactor * lambert * light.color.rgb * lightDiffuseFactor;

				let reflection = reflect(lightToPosNorm, normal);
				let specFactor = max(dot(reflection, eyeVec), 0.0);
				specFactor = pow(specFactor, settings.Shininess);

				lightSpecular += attenuationFactor * specFactor * light.color.rgb;
			}
			else if (light.type == SpotLight)
			{
				let lightPos = light.parameter1.xyz;
				let lightDir = light.parameter2.xyz;
				let lightInvRadius = light.parameter1.w;
				let lightInnerAngle = light.parameter3.x;
				let lightOuterAngle = light.parameter3.y;

				let lightToPos = input.worldPos - lightPos;
				let dist = length(lightToPos);
				let lightToPosNorm = lightToPos / max(dist, 0.0001);

				let curAngle = dot(lightDir, lightToPosNorm);
				let innerMinusOuterAngle = lightInnerAngle - lightOuterAngle;

				let attenuationFactor = max(1.0 - dist * lightInvRadius, 0.0);
				attenuationFactor *= max((curAngle - lightOuterAngle) / innerMinusOuterAngle, 0.0);			

				lightAmbient += attenuationFactor * light.color.rgb * lightAmbientFactor * settings.AmbientColor;

				let lambert = max(dot(normal, -lightToPosNorm), 0.0);

				lightDiffuse += attenuationFactor * lambert * light.color.rgb * lightDiffuseFactor;

				let reflection = reflect(lightToPosNorm, normal);
				let specFactor = max(dot(reflection, eyeVec), 0.0);
				specFactor = pow(specFactor, settings.Shininess);

				lightSpecular += attenuationFactor * specFactor * light.color.rgb;
			}
		}

		lightSpecular *= settings.SpecularColor;

		const if (HasSpecularTexture)
			lightSpecular *= MaterialSpecularMap.Sample(input.uv).rgb;

		let lightColor = lightAmbient + lightDiffuse + lightSpecular;

		let output: FragOut;
		output.RenderTarget0 = vec4[f32](lightColor, 1.0) * diffuseColor;
		return output;
	}
	else
	{
		let output: FragOut;
		output.RenderTarget0 = diffuseColor;
		return output;
	}
}

// Vertex stage
struct VertIn
{
	[location(PosLocation)] 
	pos: vec3[f32],

	[cond(HasVertexColor), location(ColorLocation)] 
	color: vec4[f32],

	[cond(HasUV), location(UvLocation)] 
	uv: vec2[f32],

	[cond(HasNormal), location(NormalLocation)]
	normal: vec3[f32],

	[cond(HasTangent), location(TangentLocation)]
	tangent: vec3[f32],

	[cond(Billboard), location(BillboardCenterLocation)]
	billboardCenter: vec3[f32],

	[cond(Billboard), location(BillboardSizeRotLocation)]
	billboardSizeRot: vec4[f32], //< width,height,sin,cos

	[cond(Billboard), location(BillboardColorLocation)]
	billboardColor: vec4[f32]
}

[entry(vert), cond(Billboard)]
fn billboardMain(input: VertIn) -> VertToFrag
{
	let size = input.billboardSizeRot.xy;
	let sinCos = input.billboardSizeRot.zw;

	let rotatedPosition = vec2[f32](
		input.pos.x * sinCos.y - input.pos.y * sinCos.x,
		input.pos.y * sinCos.y + input.pos.x * sinCos.x
	);
	rotatedPosition *= size;

	let cameraRight = vec3[f32](viewerData.viewMatrix[0][0], viewerData.viewMatrix[1][0], viewerData.viewMatrix[2][0]);
	let cameraUp = vec3[f32](viewerData.viewMatrix[0][1], viewerData.viewMatrix[1][1], viewerData.viewMatrix[2][1]);

	let vertexPos = input.billboardCenter;
	vertexPos += cameraRight * rotatedPosition.x;
	vertexPos += cameraUp * rotatedPosition.y;

	let output: VertToFrag;
	output.position = viewerData.viewProjMatrix * instanceData.worldMatrix * vec4[f32](vertexPos, 1.0);
	
	const if (HasColor)
		output.color = input.billboardColor;
	
	const if (HasUV)
		output.uv = input.pos.xy + vec2[f32](0.5, 0.5);

	return output;
}

[entry(vert), cond(!Billboard)]
fn main(input: VertIn) -> VertToFrag
{
	let worldPosition = instanceData.worldMatrix * vec4[f32](input.pos, 1.0);

	let output: VertToFrag;
	output.worldPos = worldPosition.xyz;
	output.position = viewerData.viewProjMatrix * worldPosition;

	let rotationMatrix = transpose(inverse(mat3[f32](instanceData.worldMatrix)));

	const if (HasColor)
		output.color = input.color;

	const if (HasNormal)
		output.normal = rotationMatrix * input.normal;

	const if (HasUV)
		output.uv = input.uv;

	const if (HasNormalMapping)
		output.tangent = rotationMatrix * input.tangent;

	return output;
}