// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "../Common.ush"
#include "../ViewData.ush"
#include "NaniteSceneCommon.ush"

// Represents vertex data for a Nanite mesh in local space
struct FNaniteLocalVertex
{
	// Index of the vertex in the cluster
	uint VertIndex;
	// Decoded vertex position (NOTE: On assembly parts, this is transformed into local, and is not "part"-local)
	float3 Position;
	// Prev vertex position (NOTE: This is currently only different than Position with skinned assembly parts)
	float3 PrevPosition;
	// Decoded vertex attribute data
	FNaniteRawAttributeData RawAttributeData;
	// Whether or not this vertex should be skinned
	bool bEnableSkinning;
};

// Transforms a vertex that is part of an assembly into local space
void TransformNaniteAssemblyVertex(float4x4 AssemblyTransform, float4x4 PrevAssemblyTransform, inout FNaniteLocalVertex Vert)
{
	const float3x3 NormalTransform = float3x3(
		normalize(AssemblyTransform[0].xyz),
		normalize(AssemblyTransform[1].xyz),
		normalize(AssemblyTransform[2].xyz)
	);

	// NOTE: We assume non-negative scale. Otherwise, we'd have to flip the binormal sign and reverse the winding order
	Vert.Position = mul(float4(Vert.Position, 1.0f), AssemblyTransform).xyz;
	Vert.PrevPosition = mul(float4(Vert.PrevPosition, 1.0f), PrevAssemblyTransform).xyz;
	Vert.RawAttributeData.TangentZ = mul(Vert.RawAttributeData.TangentZ, NormalTransform);
	Vert.RawAttributeData.TangentXAndSign.xyz = mul(Vert.RawAttributeData.TangentXAndSign.xyz, NormalTransform);
}

template<uint N>
void FetchLocalNaniteVerts(
	FPrimitiveSceneData PrimitiveData,
	FInstanceSceneData InstanceData,
	FInstanceViewData InstanceViewData,
	FCluster Cluster,
	FVisibleCluster VisibleCluster,
	uint3 VertIndexes,
	uint CompileTimeMaxTexCoords,
	inout FNaniteLocalVertex OutVerts[N])
{
	float3 Positions[N];
	UNROLL
	for (uint i = 0; i < N; i++)
	{
		Positions[i] = DecodePosition(VertIndexes[i], Cluster);
	}
	
	FNaniteRawAttributeData RawAttributeData[N];
	GetRawAttributeData<N>(RawAttributeData, Cluster, Positions, VertIndexes, CompileTimeMaxTexCoords);

	UNROLL
	for (uint i = 0; i < N; ++i)
	{
		OutVerts[i].VertIndex			= VertIndexes[i];
		OutVerts[i].Position			= Positions[i];
		OutVerts[i].PrevPosition		= Positions[i];
		OutVerts[i].RawAttributeData	= RawAttributeData[i];

		// HACK: disable skinning on assembly parts until we get their own local skeleton/bones hooked up
		// TODO: Nanite-Assemblies: Remove when part bones are accessible
		OutVerts[i].bEnableSkinning = !IsAssemblyPartCluster(VisibleCluster);
	}

	if (IsAssemblyPartCluster(VisibleCluster))
	{
#if NANITE_USE_RAYTRACING_UNIFORM_BUFFER
		const float4x4 AssemblyTransform = LoadNaniteHierarchyAssemblyTransform(PrimitiveData.NaniteAssemblyTransformOffset, VisibleCluster.AssemblyTransformIndex);
		float4x4 PrevAssemblyTransform = AssemblyTransform;
#else
		const float4x4 AssemblyTransform = LoadNaniteAssemblyTransform(VisibleCluster.AssemblyTransformIndex);
		float4x4 PrevAssemblyTransform = AssemblyTransform;
		BRANCH
		if (Cluster.bSkinning && InstanceViewData.bIsDeforming)
		{
			PrevAssemblyTransform = LoadNaniteAssemblyTransform(VisibleCluster.AssemblyTransformIndex + 1);
		}
#endif

		for (uint i = 0; i < N; ++i)
		{
			TransformNaniteAssemblyVertex(AssemblyTransform, PrevAssemblyTransform, OutVerts[i]);
		}
	}
}

FNaniteLocalVertex FetchLocalNaniteVertex(
	FPrimitiveSceneData PrimitiveData,
	FInstanceSceneData InstanceData,
	FInstanceViewData InstanceViewData,
	FCluster Cluster,
	FVisibleCluster VisibleCluster,
	uint VertIndex,
	uint CompileTimeMaxTexCoords)
{
	FNaniteLocalVertex Verts[1];
	FetchLocalNaniteVerts(PrimitiveData, InstanceData, InstanceViewData, Cluster, VisibleCluster, VertIndex, CompileTimeMaxTexCoords, Verts);

	return Verts[0];
}


template<uint N>
void FetchLocalNaniteVertexPositions(
	FInstanceSceneData InstanceData,
	FCluster Cluster,
	FVisibleCluster VisibleCluster,
	uint3 VertIndexes,
	inout float3 OutPositions[N])
{
	UNROLL
	for (uint i = 0; i < N; i++)
	{
		OutPositions[i] = DecodePosition(VertIndexes[i], Cluster);
	}
	
	if (IsAssemblyPartCluster(VisibleCluster))
	{
		const float4x4 AssemblyTransform = LoadNaniteAssemblyTransform(VisibleCluster.AssemblyTransformIndex);

		UNROLL
		for (uint i = 0; i < N; ++i)
		{
			OutPositions[i] = mul(float4(OutPositions[i], 1), AssemblyTransform).xyz;
		}
	}
}

float3 FetchLocalNaniteVertexPosition(
	FInstanceSceneData InstanceData,
	FCluster Cluster,
	FVisibleCluster VisibleCluster,
	uint VertIndex)
{
	float3 Positions[1];
	FetchLocalNaniteVertexPositions(InstanceData, Cluster, VisibleCluster, VertIndex, Positions);

	return Positions[0];
}