[nzsl_version("1.0")]
module PhongMaterial;

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

import SkinLinearPosition, SkinLinearPositionNormal from Engine.SkinningLinear;

// Pass-specific options
option DepthPass: bool = false;

// Basic material options
option HasBaseColorTexture: 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 VertexColorLoc: i32 = -1;
option VertexNormalLoc: i32 = -1;
option VertexPositionLoc: i32;
option VertexTangentLoc: i32 = -1;
option VertexUvLoc: i32 = -1;

option VertexJointIndicesLoc: i32 = -1;
option VertexJointWeightsLoc: i32 = -1;

option MaxLightCount: u32 = u32(3); //< FIXME: Fix integral value types

const HasNormal = (VertexNormalLoc >= 0);
const HasVertexColor = (VertexColorLoc >= 0);
const HasColor = (HasVertexColor || Billboard);
const HasTangent = (VertexTangentLoc >= 0);
const HasUV = (VertexUvLoc >= 0);
const HasNormalMapping = HasNormalTexture && HasNormal && HasTangent && !DepthPass;
const HasSkinning = (VertexJointIndicesLoc >= 0 && VertexJointWeightsLoc >= 0);
const HasLighting = HasNormal && !DepthPass;

[layout(std140)]
struct MaterialSettings
{
	// Basic settings
	[tag("AlphaTestThreshold")]
	AlphaThreshold: f32,

	[tag("BaseColor")]
	BaseColor: vec4[f32],

	// Phong settings
	[tag("AmbientColor")]
	AmbientColor: vec4[f32], //< TODO: Switch to vec3[f32]

	[tag("SpecularColor")]
	SpecularColor: vec4[f32], //< TODO: Switch to vec3[f32

	[tag("Shininess")]
	Shininess: f32
}

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

[tag("Material")]
[auto_binding]
external
{
	[tag("Settings")] settings: uniform[MaterialSettings],
	[tag("BaseColorMap")] MaterialBaseColorMap: sampler2D[f32],
	[tag("AlphaMap")] MaterialAlphaMap: sampler2D[f32],
	[tag("EmissiveMap")] MaterialEmissiveMap: sampler2D[f32],
	[tag("HeightMap")] MaterialHeightMap: sampler2D[f32],
	[tag("NormalMap")] MaterialNormalMap: sampler2D[f32],
	[tag("SpecularMap")] MaterialSpecularMap: sampler2D[f32],
}

[tag("Engine")]
[auto_binding]
external
{
	[tag("TextureOverlay")] TextureOverlay: sampler2D[f32],
	[tag("InstanceData")] instanceData: uniform[InstanceData],
	[tag("ViewerData")] viewerData: uniform[ViewerData],
	[tag("SkeletalData")] skeletalData: uniform[SkeletalData],
	[tag("LightData")] lightData: uniform[LightData],
	[tag("ShadowMaps2D")] shadowMaps2D: array[depth_sampler2D[f32], MaxLightCount],
	[tag("ShadowMapsCube")] shadowMapsCube: array[sampler_cube[f32], MaxLightCount]
}

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],
	[location(5), cond(HasLighting)] lightProjPos: array[vec4[f32], MaxLightCount],
	[builtin(position)] position: vec4[f32],
}

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

fn LinearizeDepth(depth: f32, zNear: f32, zFar: f32) -> f32
{
    return zNear * zFar / (zFar + depth * (zNear - zFar));
}

[entry(frag), cond(!DepthPass || AlphaTest)]
fn main(input: VertToFrag) -> FragOut
{
	let color = settings.BaseColor;

	const if (HasUV)
		color.a *= TextureOverlay.Sample(input.uv).r;

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

	const if (HasBaseColorTexture)
		color *= MaterialBaseColorMap.Sample(input.uv);

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

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

	const if (HasLighting)
	{
		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.rgb;

				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 lightRadius = light.parameter2.x;
				let lightInvRadius = light.parameter2.y;

				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);
				let lambert = max(dot(normal, -lightToPosNorm), 0.0);

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

				let shadowFactor = 1.0;
				if (light.invShadowMapSize.x > 0.0)
				{
					shadowFactor = 0.0;

					let sampleDir = vec3[f32](lightToPosNorm.x, lightToPosNorm.y, -lightToPosNorm.z);

					const sampleCount = 4;
					const offset = 0.005;

					const invSampleCount = 1.0 / f32(sampleCount);
					const start = vec3[f32](offset * 0.5, offset * 0.5, offset * 0.5);
					const shadowContribution = 1.0 / f32(sampleCount * sampleCount * sampleCount);

					[unroll]
					for x in 0 -> sampleCount
					{
						[unroll]
						for y in 0 -> sampleCount
						{
							[unroll]
							for z in 0 -> sampleCount
							{
								let dirOffset = vec3[f32](f32(x), f32(y), f32(z)) * invSampleCount * offset - start;
								let sampleDir = sampleDir + dirOffset;

								let depth = shadowMapsCube[i].Sample(sampleDir).r;
								depth = LinearizeDepth(depth, 0.01, lightRadius);

								if (depth > dist)
									shadowFactor += shadowContribution;
							}
						}
					}
				}

				lightAmbient += attenuationFactor * light.color.rgb * lightAmbientFactor * settings.AmbientColor.rgb;
				lightDiffuse += shadowFactor * attenuationFactor * lambert * light.color.rgb * lightDiffuseFactor;
				lightSpecular += shadowFactor * 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);			

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

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

				let shadowFactor = 1.0;
				if (light.invShadowMapSize.x > 0.0)
				{
					let shadowCoords = input.lightProjPos[i].xyz / input.lightProjPos[i].w;
					shadowFactor = 0.0;
					[unroll]
					for x in -1 -> 2
					{
						[unroll]
						for y in -1 -> 2
						{
							let coords = shadowCoords.xy + vec2[f32](f32(x), f32(y)) * light.invShadowMapSize;
							shadowFactor += shadowMaps2D[i].SampleDepthComp(coords, shadowCoords.z).r;
						}
					}
					shadowFactor /= 9.0;
				}

				lightAmbient += attenuationFactor * light.color.rgb * lightAmbientFactor * settings.AmbientColor.rgb;
				lightDiffuse += shadowFactor * attenuationFactor * lambert * light.color.rgb * lightDiffuseFactor;
				lightSpecular += shadowFactor * attenuationFactor * specFactor * light.color.rgb;
			}
		}

		lightSpecular *= settings.SpecularColor.rgb;

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

		let lightColor = lightAmbient + lightDiffuse + lightSpecular;

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

// Dummy fragment shader (TODO: Add a way to delete stage?)
[entry(frag), cond(DepthPass && !AlphaTest)]
fn main() {}


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

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

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

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

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

	[cond(HasSkinning), location(VertexJointIndicesLoc)]
	jointIndices: vec4[i32],

	[cond(HasSkinning), location(VertexJointWeightsLoc)]
	jointWeights: vec4[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 pos: vec3[f32];
	const if (HasNormal) let normal: vec3[f32];

	const if (HasSkinning)
	{
		let jointMatrices = array[mat4[f32]](
			skeletalData.jointMatrices[input.jointIndices[0]],
			skeletalData.jointMatrices[input.jointIndices[1]],
			skeletalData.jointMatrices[input.jointIndices[2]],
			skeletalData.jointMatrices[input.jointIndices[3]]
		);

		const if (HasNormal)
		{
			let skinningOutput = SkinLinearPositionNormal(jointMatrices, input.jointWeights, input.pos, input.normal);
			pos = skinningOutput.position;
			normal = skinningOutput.normal;
		}
		else
		{
			let skinningOutput = SkinLinearPosition(jointMatrices, input.jointWeights, input.pos);
			pos = skinningOutput.position;
		}
	}
	else
	{
		pos = input.pos;
		const if (HasNormal)
			normal = input.normal;
	}

	let worldPosition = instanceData.worldMatrix * vec4[f32](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 * normal;

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

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

	const if (HasLighting)
	{
		for i in u32(0) -> lightData.lightCount
			output.lightProjPos[i] = lightData.lights[i].viewProjMatrix * worldPosition;
	}

	return output;
}