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

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

// 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 = -1;
option UvLocation: i32 = -1;

const HasNormal = (NormalLocation >= 0);
const HasVertexColor = (ColorLocation >= 0);
const HasColor = (HasVertexColor || Billboard);
const HasUV = (UvLocation >= 0);

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

	// PhongSettings
	AmbientColor: vec3<f32>,
	SpecularColor: vec3<f32>,
	Shininess: f32,
}

[layout(std140)]
struct InstanceData
{
	worldMatrix: mat4<f32>,
	invWorldMatrix: mat4<f32>
}

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

[layout(std140)]
struct Light
{
	type: i32,
	color: vec4<f32>,
	factor: vec2<f32>,
	parameter1: vec4<f32>,
	parameter2: vec4<f32>,
	parameter3: vec4<f32>,
	hasShadowMapping: bool
}

[layout(std140)]
struct LightData
{
	lights: [Light; MaxLightCount],
	lightCount: u32,
}

[layout(std140)]
struct ViewerData
{
	projectionMatrix: mat4<f32>,
	invProjectionMatrix: mat4<f32>,
	viewMatrix: mat4<f32>,
	invViewMatrix: mat4<f32>,
	viewProjMatrix: mat4<f32>,
	invViewProjMatrix: mat4<f32>,
	renderTargetSize: vec2<f32>,
	invRenderTargetSize: vec2<f32>,
	eyePosition: vec3<f32>
}

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

// Fragment stage
struct FragIn
{
	[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>,
}

struct FragOut
{
	[location(0)] RenderTarget0: vec4<f32>
}

[entry(frag)]
fn main(input: FragIn) -> 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);

		for i in 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); //< FIXME

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

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

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

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

				lightSpecular += specFactor * light.color.rgb;
			}
			else if (light.type == PointLight)
			{
	
			}
			else if (light.type == SpotLight)
			{
	
			}
		}

		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.w;
		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(Billboard), location(BillboardCenterLocation)]
	billboardCenter: vec3<f32>,

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

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

struct VertOut
{
	[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>,
	[builtin(position)] position: vec4<f32>,
}

[entry(vert), cond(Billboard)]
fn billboardMain(input: VertIn) -> VertOut
{
	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: VertOut;
	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) -> VertOut
{
	let worldPosition = instanceData.worldMatrix * vec4<f32>(input.pos, 1.0);

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

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

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

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

	return output;
}