#include "MujocoMeshFactory.h"
#include "AssetRegistry/AssetRegistryModule.h"
#include "Engine/StaticMesh.h"
#include "StaticMeshAttributes.h"

namespace MujocoMeshFactory
{
	namespace
	{
		UStaticMesh* CreateMeshAsset(const FString& MeshPath, const FString& MeshName)
		{
			const FString PackageFullName = FPaths::Combine(MeshPath, MeshName);
			UPackage* Package = CreatePackage(*PackageFullName);
			UStaticMesh* Mesh = NewObject<UStaticMesh>(Package, FName(*MeshName), RF_Public | RF_Standalone);

			if (Mesh)
			{
				Mesh->InitResources();
				Mesh->SetLightingGuid();
			}
			return Mesh;
		}

		void CommitMesh(UStaticMesh* Mesh)
		{
			if (Mesh)
			{
				Mesh->MarkPackageDirty();
				FAssetRegistryModule::AssetCreated(Mesh);
				Mesh->PostEditChange();
			}
		}
	} // namespace

	UStaticMesh* CreateBoxMesh(const FVector& InSize, const FString& MeshPath, const FString& InGeomName)
	{
		UStaticMesh* Mesh = CreateMeshAsset(MeshPath, InGeomName);
		if (!Mesh)
		{
			return nullptr;
		}

		FMeshDescription MeshDesc;
		FStaticMeshAttributes Attributes(MeshDesc);
		Attributes.Register();
		FPolygonGroupID PolygonGroup = MeshDesc.CreatePolygonGroup();
		TArray<FVertexID> Vertices;
		const FVector Extents = InSize;
		TArray<FVector> Positions = { FVector(-Extents.X, -Extents.Y, -Extents.Z), FVector(Extents.X, -Extents.Y, -Extents.Z), FVector(Extents.X, Extents.Y, -Extents.Z), FVector(-Extents.X, Extents.Y, -Extents.Z), FVector(-Extents.X, -Extents.Y, Extents.Z), FVector(Extents.X, -Extents.Y, Extents.Z), FVector(Extents.X, Extents.Y, Extents.Z), FVector(-Extents.X, Extents.Y, Extents.Z) };

		for (const FVector& Pos : Positions)
		{
			FVertexID VertexID = MeshDesc.CreateVertex();
			Attributes.GetVertexPositions()[VertexID] = (FVector3f)Pos;
			Vertices.Add(VertexID);
		}

		TArray<TArray<int32>> Faces = {
			{ 0, 1, 2, 3 }, // Bottom Face
			{ 4, 7, 6, 5 }, // Top Face
			{ 0, 4, 5, 1 }, // Front Face
			{ 3, 2, 6, 7 }, // Back Face
			{ 1, 5, 6, 2 }, // Right Face
			{ 0, 3, 7, 4 }	// Left Face
		};

		for (const TArray<int32>& Face : Faces)
		{
			if (Face.Num() < 3)
			{
				continue;
			}
			TArray<FVertexInstanceID> VertexInstances;
			for (int32 VertexIndex : Face)
			{
				FVertexInstanceID VertexInstance = MeshDesc.CreateVertexInstance(Vertices[VertexIndex]);
				VertexInstances.Add(VertexInstance);
			}
			for (int32 i = 1; i < Face.Num() - 1; ++i)
			{
				MeshDesc.CreatePolygon(PolygonGroup, { VertexInstances[0], VertexInstances[i], VertexInstances[i + 1] });
			}
		}

		Mesh->BuildFromMeshDescriptions({ &MeshDesc });
		Mesh->PostEditChange();
		CommitMesh(Mesh);
		return Mesh;
	}

	UStaticMesh* CreateSphereMesh(const FVector& InSize, const FString& MeshPath, const FString& InGeomName)
	{
		UStaticMesh* Mesh = CreateMeshAsset(MeshPath, InGeomName);
		if (!Mesh)
		{
			return nullptr;
		}

		FMeshDescription MeshDesc;
		FStaticMeshAttributes Attributes(MeshDesc);
		Attributes.Register();
		FPolygonGroupID PolygonGroup = MeshDesc.CreatePolygonGroup();

		TArray<FVertexID> Vertices;
		const int32 LatitudeSegments = 16;
		const int32 LongitudeSegments = 32;
		const float Radius = FMath::Max(InSize.X, FMath::Max(InSize.Y, InSize.Z));

		for (int32 Lat = 0; Lat <= LatitudeSegments; ++Lat)
		{
			float Theta = Lat * PI / LatitudeSegments;
			float SinTheta = FMath::Sin(Theta);
			float CosTheta = FMath::Cos(Theta);

			for (int32 Lon = 0; Lon <= LongitudeSegments; ++Lon)
			{
				float Phi = Lon * 2.0f * PI / LongitudeSegments;
				float SinPhi = FMath::Sin(Phi);
				float CosPhi = FMath::Cos(Phi);

				FVector Position = FVector(CosPhi * SinTheta, SinPhi * SinTheta, CosTheta) * Radius;
				FVertexID VertexID = MeshDesc.CreateVertex();
				Attributes.GetVertexPositions()[VertexID] = (FVector3f)Position;
				Vertices.Add(VertexID);
			}
		}

		for (int32 Lat = 0; Lat < LatitudeSegments; ++Lat)
		{
			for (int32 Lon = 0; Lon < LongitudeSegments; ++Lon)
			{
				int32 Index0 = Lat * (LongitudeSegments + 1) + Lon;
				int32 Index1 = Index0 + LongitudeSegments + 1;
				int32 Index2 = Index0 + 1;
				int32 Index3 = Index1 + 1;

				FVertexInstanceID V0 = MeshDesc.CreateVertexInstance(Vertices[Index0]);
				FVertexInstanceID V1 = MeshDesc.CreateVertexInstance(Vertices[Index1]);
				FVertexInstanceID V2 = MeshDesc.CreateVertexInstance(Vertices[Index2]);
				FVertexInstanceID V3 = MeshDesc.CreateVertexInstance(Vertices[Index3]);

				MeshDesc.CreatePolygon(PolygonGroup, { V0, V2, V1 });
				MeshDesc.CreatePolygon(PolygonGroup, { V2, V3, V1 });
			}
		}

		Mesh->BuildFromMeshDescriptions({ &MeshDesc });
		Mesh->PostEditChange();
		CommitMesh(Mesh);
		return Mesh;
	}

	UStaticMesh* CreateCylinderMesh(const FVector& InSize, const FString& MeshPath, const FString& InGeomName)
	{
		UStaticMesh* Mesh = CreateMeshAsset(MeshPath, InGeomName);
		if (!Mesh)
		{
			return nullptr;
		}

		FMeshDescription MeshDesc;
		FStaticMeshAttributes Attributes(MeshDesc);
		Attributes.Register();
		FPolygonGroupID PolygonGroup = MeshDesc.CreatePolygonGroup();

		TArray<FVertexID> Vertices;
		const int32 RadialSegments = 32;
		const float Radius = InSize.X;
		const float HalfHeight = InSize.Z * 0.5f;

		// Create top and bottom ring vertices
		for (int32 j = 0; j <= RadialSegments; ++j)
		{
			float Angle = j * 2.0f * PI / RadialSegments;
			float X = Radius * FMath::Cos(Angle);
			float Y = Radius * FMath::Sin(Angle);

			FVector BottomVertex(X, Y, -HalfHeight);
			FVector TopVertex(X, Y, HalfHeight);

			FVertexID BottomVertexID = MeshDesc.CreateVertex();
			Attributes.GetVertexPositions()[BottomVertexID] = (FVector3f)BottomVertex;
			Vertices.Add(BottomVertexID);

			FVertexID TopVertexID = MeshDesc.CreateVertex();
			Attributes.GetVertexPositions()[TopVertexID] = (FVector3f)TopVertex;
			Vertices.Add(TopVertexID);
		}

		// Create top and bottom center vertices
		FVertexID BottomCenter = MeshDesc.CreateVertex();
		Attributes.GetVertexPositions()[BottomCenter] = (FVector3f)FVector(0, 0, -HalfHeight);
		FVertexID TopCenter = MeshDesc.CreateVertex();
		Attributes.GetVertexPositions()[TopCenter] = (FVector3f)FVector(0, 0, HalfHeight);

		// Generate body faces with corrected counter-clockwise winding
		for (int32 j = 0; j < RadialSegments; ++j)
		{
			int32 Index0 = j * 2;
			int32 Index1 = ((j + 1) % RadialSegments) * 2;
			int32 Index2 = Index0 + 1;
			int32 Index3 = Index1 + 1;

			FVertexInstanceID V0 = MeshDesc.CreateVertexInstance(Vertices[Index0]);
			FVertexInstanceID V1 = MeshDesc.CreateVertexInstance(Vertices[Index1]);
			FVertexInstanceID V2 = MeshDesc.CreateVertexInstance(Vertices[Index2]);
			FVertexInstanceID V3 = MeshDesc.CreateVertexInstance(Vertices[Index3]);

			MeshDesc.CreatePolygon(PolygonGroup, { V0, V2, V3 }); // Corrected winding order
			MeshDesc.CreatePolygon(PolygonGroup, { V0, V3, V1 });
		}

		// Create bottom cap faces (Counter-clockwise winding)
		for (int32 j = 0; j < RadialSegments; ++j)
		{
			int32 Index0 = j * 2;
			int32 Index1 = ((j + 1) % RadialSegments) * 2;

			FVertexInstanceID VB0 = MeshDesc.CreateVertexInstance(Vertices[Index0]);
			FVertexInstanceID VB1 = MeshDesc.CreateVertexInstance(Vertices[Index1]);
			FVertexInstanceID VBC = MeshDesc.CreateVertexInstance(BottomCenter);

			MeshDesc.CreatePolygon(PolygonGroup, { VB0, VB1, VBC }); // Counter-clockwise
		}

		// Create top cap faces (Counter-clockwise winding)
		for (int32 j = 0; j < RadialSegments; ++j)
		{
			int32 Index0 = j * 2 + 1;
			int32 Index1 = ((j + 1) % RadialSegments) * 2 + 1;

			FVertexInstanceID VT0 = MeshDesc.CreateVertexInstance(Vertices[Index1]);
			FVertexInstanceID VT1 = MeshDesc.CreateVertexInstance(Vertices[Index0]);
			FVertexInstanceID VTC = MeshDesc.CreateVertexInstance(TopCenter);

			MeshDesc.CreatePolygon(PolygonGroup, { VT0, VT1, VTC }); // Counter-clockwise
		}

		Mesh->BuildFromMeshDescriptions({ &MeshDesc });
		Mesh->PostEditChange();
		CommitMesh(Mesh);
		return Mesh;
	}

	UStaticMesh* CreateCapsuleMesh(const FVector& InSize, const FString& MeshPath, const FString& InGeomName)
	{
		UStaticMesh* Mesh = CreateMeshAsset(MeshPath, InGeomName);
		if (!Mesh)
		{
			return nullptr;
		}

		FMeshDescription MeshDesc;
		FStaticMeshAttributes Attributes(MeshDesc);
		Attributes.Register();
		FPolygonGroupID PolygonGroup = MeshDesc.CreatePolygonGroup();

		TArray<FVertexID> Vertices;
		const int32 NumSubdivisionsHeight = 12;
		const int32 NumSegments = 12;
		const float Radius = InSize.X;
		const float Height = InSize.Z + Radius;

		auto CalculateRing = [&](int32 Segments, float r, float y, float dy) {
			float SegIncr = 1.0f / (Segments - 1);
			for (int32 s = 0; s < Segments; s++)
			{
				float Angle = (PI * 2) * s * SegIncr;
				float X = -FMath::Cos(Angle) * r;
				float Z = FMath::Sin(Angle) * r;

				FVector Position = FVector(Radius * X, Radius * Z, Radius * y + Height * dy);
				FVertexID VertexID = MeshDesc.CreateVertex();
				Attributes.GetVertexPositions()[VertexID] = (FVector3f)Position;
				Vertices.Add(VertexID);
			}
		};

		int32 RingsBody = NumSubdivisionsHeight + 1;
		int32 RingsTotal = NumSubdivisionsHeight + RingsBody;
		float BodyIncr = 1.0f / (RingsBody - 1);
		float RingIncr = 1.0f / (NumSubdivisionsHeight - 1);

		for (int32 r = 0; r < NumSubdivisionsHeight / 2; r++)
		{
			CalculateRing(NumSegments, FMath::Sin(PI * r * RingIncr), FMath::Sin(PI * (r * RingIncr - 0.5f)), -0.5f);
		}

		for (int32 r = 0; r < RingsBody; r++)
		{
			CalculateRing(NumSegments, 1.0f, 0.0f, r * BodyIncr - 0.5f);
		}

		for (int32 r = NumSubdivisionsHeight / 2; r < NumSubdivisionsHeight; r++)
		{
			CalculateRing(NumSegments, FMath::Sin(PI * r * RingIncr), FMath::Sin(PI * (r * RingIncr - 0.5f)), +0.5f);
		}

		// Recalculate normals counter-clockwise
		for (int32 r = 0; r < RingsTotal - 1; r++)
		{
			for (int32 s = 0; s < NumSegments - 1; s++)
			{
				FVertexInstanceID V0 = MeshDesc.CreateVertexInstance(Vertices[r * NumSegments + (s + 1)]);
				FVertexInstanceID V1 = MeshDesc.CreateVertexInstance(Vertices[r * NumSegments + s]);
				FVertexInstanceID V2 = MeshDesc.CreateVertexInstance(Vertices[(r + 1) * NumSegments + (s + 1)]);

				MeshDesc.CreatePolygon(PolygonGroup, { V0, V2, V1 });

				FVertexInstanceID V3 = MeshDesc.CreateVertexInstance(Vertices[(r + 1) * NumSegments + s]);
				FVertexInstanceID V4 = MeshDesc.CreateVertexInstance(Vertices[(r + 1) * NumSegments + (s + 1)]);
				FVertexInstanceID V5 = MeshDesc.CreateVertexInstance(Vertices[r * NumSegments + s]);

				MeshDesc.CreatePolygon(PolygonGroup, { V3, V5, V4 });
			}
		}

		Mesh->BuildFromMeshDescriptions({ &MeshDesc });
		Mesh->PostEditChange();
		CommitMesh(Mesh);
		return Mesh;
	}

	UStaticMesh* CreateEllipsoidMesh(const FVector& InSize, const FString& MeshPath, const FString& InGeomName)
	{
		UStaticMesh* Mesh = CreateMeshAsset(MeshPath, InGeomName);
		if (!Mesh)
		{
			return nullptr;
		}

		FMeshDescription MeshDesc;
		FStaticMeshAttributes Attributes(MeshDesc);
		Attributes.Register();
		FPolygonGroupID PolygonGroup = MeshDesc.CreatePolygonGroup();

		TArray<FVertexID> Vertices;
		const int32 LatitudeSegments = 16;
		const int32 LongitudeSegments = 32;
		const float RadiusX = InSize.X;
		const float RadiusY = InSize.Y;
		const float RadiusZ = InSize.Z;

		// Generate vertices
		for (int32 Lat = 0; Lat <= LatitudeSegments; ++Lat)
		{
			float Theta = Lat * PI / LatitudeSegments;
			float SinTheta = FMath::Sin(Theta);
			float CosTheta = FMath::Cos(Theta);

			for (int32 Lon = 0; Lon <= LongitudeSegments; ++Lon)
			{
				float Phi = Lon * 2.0f * PI / LongitudeSegments;
				float SinPhi = FMath::Sin(Phi);
				float CosPhi = FMath::Cos(Phi);

				FVector Normal = FVector(CosPhi * SinTheta, SinPhi * SinTheta, CosTheta).GetSafeNormal();
				FVector Position = FVector(CosPhi * SinTheta * RadiusX, SinPhi * SinTheta * RadiusY, CosTheta * RadiusZ);

				FVertexID VertexID = MeshDesc.CreateVertex();
				Attributes.GetVertexPositions()[VertexID] = (FVector3f)Position;
				Vertices.Add(VertexID);
			}
		}

		// Generate faces with correct counter-clockwise order
		for (int32 Lat = 0; Lat < LatitudeSegments; ++Lat)
		{
			for (int32 Lon = 0; Lon < LongitudeSegments; ++Lon)
			{
				int32 Index0 = Lat * (LongitudeSegments + 1) + Lon;
				int32 Index1 = Index0 + 1;
				int32 Index2 = Index0 + LongitudeSegments + 1;
				int32 Index3 = Index2 + 1;

				if (Lat < LatitudeSegments - 1) // Skip top pole
				{
					FVertexInstanceID V0 = MeshDesc.CreateVertexInstance(Vertices[Index0]);
					FVertexInstanceID V1 = MeshDesc.CreateVertexInstance(Vertices[Index2]);
					FVertexInstanceID V2 = MeshDesc.CreateVertexInstance(Vertices[Index3]);
					MeshDesc.CreatePolygon(PolygonGroup, { V0, V2, V1 });
				}

				if (Lat > 0) // Skip bottom pole
				{
					FVertexInstanceID V3 = MeshDesc.CreateVertexInstance(Vertices[Index3]);
					FVertexInstanceID V4 = MeshDesc.CreateVertexInstance(Vertices[Index1]);
					FVertexInstanceID V5 = MeshDesc.CreateVertexInstance(Vertices[Index0]);
					MeshDesc.CreatePolygon(PolygonGroup, { V3, V5, V4 });
				}
			}
		}

		Mesh->BuildFromMeshDescriptions({ &MeshDesc });
		Mesh->PostEditChange();
		CommitMesh(Mesh);
		return Mesh;
	}

} // namespace BasicShapeMeshFactory