[nzsl_version("1.0")]
module;

import Engine/InstanceData;
import Engine/ViewerData;

option HasDiffuseTexture: bool = false;
option HasAlphaTexture: bool = false;
option AlphaTest: 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 PosLocation: i32;
option UvLocation: i32 = -1;

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

[layout(std140)]
struct MaterialSettings
{
	AlphaThreshold: f32,
	DiffuseColor: vec4[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],
}

// Fragment stage
struct FragIn
{
	[location(0), cond(HasUV)] uv: vec2[f32],
	[location(1), cond(HasColor)] color: vec4[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;
	}

	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(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), cond(HasUV)] uv: vec2[f32],
	[location(1), cond(HasColor)] color: vec4[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 output: VertOut;
	output.position = viewerData.viewProjMatrix * instanceData.worldMatrix * vec4[f32](input.pos, 1.0);

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

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

	return output;
}