// Portage UE 5.5 - voir LPHairFromStaticMeshCommandlet.h

#include "LPHairFromStaticMeshCommandlet.h"

#include "AnimationRuntime.h"
#include "Animation/Skeleton.h"
#include "Engine/SkeletalMesh.h"
#include "Engine/StaticMesh.h"
#include "MeshDescription.h"
#include "Misc/PackageName.h"
#include "Rendering/SkeletalMeshLODModel.h"
#include "Rendering/SkeletalMeshModel.h"
#include "StaticMeshAttributes.h"
#include "SkeletalMeshAttributes.h"
#include "StaticToSkeletalMeshConverter.h"
#include "UObject/Package.h"
#include "UObject/SavePackage.h"

DEFINE_LOG_CATEGORY_STATIC(LogLPHairFromStatic, Log, All);

namespace
{
	template <typename T>
	T* LoadByPath(const FString& Path)
	{
		FString Full = Path;
		if (!Full.Contains(TEXT(".")))
		{
			Full += TEXT(".") + FPackageName::GetShortName(Path);
		}
		return LoadObject<T>(nullptr, *Full);
	}

	/** Sommets du corps rattaches (influence dominante) a l'os HeadIndex ou a un de ses descendants. */
	TArray<FVector> HeadPoints(const USkeletalMesh* Body, int32 HeadIndex)
	{
		TArray<FVector> Points;
		const FReferenceSkeleton& Ref = Body->GetRefSkeleton();
		auto IsUnderHead = [&](int32 Bone)
		{
			for (; Bone != INDEX_NONE; Bone = Ref.GetParentIndex(Bone))
			{
				if (Bone == HeadIndex)
				{
					return true;
				}
			}
			return false;
		};
		const FSkeletalMeshLODModel& LOD = Body->GetImportedModel()->LODModels[0];
		for (const FSkelMeshSection& Section : LOD.Sections)
		{
			for (const FSoftSkinVertex& V : Section.SoftVertices)
			{
				int32 Best = 0;
				for (int32 i = 1; i < MAX_TOTAL_INFLUENCES; ++i)
				{
					if (V.InfluenceWeights[i] > V.InfluenceWeights[Best])
					{
						Best = i;
					}
				}
				const int32 Bone = Section.BoneMap.IsValidIndex(V.InfluenceBones[Best]) ? Section.BoneMap[V.InfluenceBones[Best]] : INDEX_NONE;
				if (IsUnderHead(Bone))
				{
					Points.Add(FVector(V.Position));
				}
			}
		}
		return Points;
	}

	/** Sommets d'un maillage statique (donnees sources, LOD 0). */
	TArray<FVector> StaticPoints(UStaticMesh* Mesh)
	{
		TArray<FVector> Points;
		if (const FMeshDescription* Desc = Mesh->GetMeshDescription(0))
		{
			TVertexAttributesConstRef<FVector3f> Positions = FStaticMeshConstAttributes(*Desc).GetVertexPositions();
			for (const FVertexID Vertex : Desc->Vertices().GetElementIDs())
			{
				Points.Add(FVector(Positions[Vertex]));
			}
		}
		return Points;
	}

	/** Boite du haut du crane : sommets situes a moins de DepthFraction x profondeur de la tete sous le sommet. Exclut
	 *  oreilles, nez et cou, dont les proportions varient d'une tete de reference a l'autre. */
	FBox CraniumBox(const TArray<FVector>& Points, double DepthFraction)
	{
		FBox Full(Points);
		const double Floor = Full.Max.Z - DepthFraction * Full.GetSize().Y;
		FBox Box(ForceInit);
		for (const FVector& Point : Points)
		{
			if (Point.Z >= Floor)
			{
				Box += Point;
			}
		}
		return Box;
	}

	/** Sommet du corps pour le transfert de poids : position, normale, poids exprimes en os du squelette de la coiffure. */
	struct FBodySample
	{
		FVector Position;
		FVector Normal;
		TArray<TPair<int32, float>, TInlineAllocator<8>> Weights;
	};

	/** Os des membres : une meche ne doit pas suivre un bras ni une jambe ; ces os sont ramenes a leur ancetre. */
	bool IsLimbBone(const FString& Name)
	{
		static const TCHAR* Parts[] = { TEXT("Shldr"), TEXT("Forearm"), TEXT("Hand"), TEXT("Carpal"), TEXT("Thumb"), TEXT("Index"),
			TEXT("Mid"), TEXT("Ring"), TEXT("Pinky"), TEXT("Thigh"), TEXT("Shin"), TEXT("Foot"), TEXT("Toe"), TEXT("Metatarsal") };
		for (const TCHAR* Part : Parts)
		{
			if (Name.Contains(Part))
			{
				return true;
			}
		}
		return false;
	}

	/** Sommets du corps au-dessus de MinZ, poids convertis vers les os du squelette cible (ancetre le plus proche present). */
	TArray<FBodySample> BodySamples(const USkeletalMesh* Body, const FReferenceSkeleton& Target, double MinZ)
	{
		const FReferenceSkeleton& Ref = Body->GetRefSkeleton();
		TArray<int32> BoneMap;
		BoneMap.SetNum(Ref.GetRawBoneNum());
		for (int32 Bone = 0; Bone < Ref.GetRawBoneNum(); ++Bone)
		{
			int32 Mapped = 0;
			for (int32 B = Bone; B != INDEX_NONE; B = Ref.GetParentIndex(B))
			{
				const FString Name = Ref.GetBoneName(B).ToString();
				const int32 T = IsLimbBone(Name) ? INDEX_NONE : Target.FindBoneIndex(*Name);
				if (T != INDEX_NONE)
				{
					Mapped = T;
					break;
				}
			}
			BoneMap[Bone] = Mapped;
		}
		TArray<FBodySample> Samples;
		const FSkeletalMeshLODModel& LOD = Body->GetImportedModel()->LODModels[0];
		for (const FSkelMeshSection& Section : LOD.Sections)
		{
			for (const FSoftSkinVertex& V : Section.SoftVertices)
			{
				if (V.Position.Z < MinZ)
				{
					continue;
				}
				FBodySample& Sample = Samples.AddDefaulted_GetRef();
				Sample.Position = FVector(V.Position);
				Sample.Normal = FVector(FVector3f(V.TangentZ)).GetSafeNormal();
				for (int32 i = 0; i < MAX_TOTAL_INFLUENCES; ++i)
				{
					if (V.InfluenceWeights[i] > 0 && Section.BoneMap.IsValidIndex(V.InfluenceBones[i]))
					{
						Sample.Weights.Add({ BoneMap[Section.BoneMap[V.InfluenceBones[i]]], V.InfluenceWeights[i] / 65535.f });
					}
				}
			}
		}
		return Samples;
	}

	/** Point de reference d'une tete : centre en X et Y, sommet du crane en Z. */
	FVector TopCenter(const FBox& Box)
	{
		return FVector(Box.GetCenter().X, Box.GetCenter().Y, Box.Max.Z);
	}

	bool SaveAsset(UObject* Asset)
	{
		UPackage* Package = Asset->GetOutermost();
		Package->MarkPackageDirty();
		const FString Filename = FPackageName::LongPackageNameToFilename(Package->GetName(), FPackageName::GetAssetPackageExtension());
		FSavePackageArgs SaveArgs;
		SaveArgs.TopLevelFlags = RF_Public | RF_Standalone;
		SaveArgs.Error = GWarn;
		const bool bOk = UPackage::SavePackage(Package, Asset, *Filename, SaveArgs);
		UE_LOG(LogLPHairFromStatic, Warning, TEXT("  %s %s"), bOk ? TEXT("sauvegarde :") : TEXT("ECHEC de sauvegarde :"), *Filename);
		return bOk;
	}
}

ULPHairFromStaticMeshCommandlet::ULPHairFromStaticMeshCommandlet()
{
	IsClient = false;
	IsEditor = true;
	IsServer = false;
	LogToConsole = true;
}

int32 ULPHairFromStaticMeshCommandlet::Main(const FString& Params)
{
	TArray<FString> Tokens, Switches;
	TMap<FString, FString> P;
	ParseCommandLine(*Params, Tokens, Switches, P);
	const FString StaticPath = P.FindRef(TEXT("Static"));
	const FString FitHeadPath = P.FindRef(TEXT("FitHead"));
	const FString OutPath = P.FindRef(TEXT("Out"));
	const FString BodyPath = P.Contains(TEXT("Body")) ? P[TEXT("Body")] : TEXT("/Game/Heavy/Daz/Bridget/bridget_gp");
	const FString RefHairPath = P.Contains(TEXT("RefHair")) ? P[TEXT("RefHair")] : TEXT("/Game/Heavy/New4/Hairs/FloraHair_New");
	const FName BoneName(P.Contains(TEXT("Bone")) ? *P[TEXT("Bone")] : TEXT("head"));
	const double ExtraScale = P.Contains(TEXT("Scale")) ? FCString::Atod(*P[TEXT("Scale")]) : 1.0;
	const double OffsetZ = P.Contains(TEXT("OffsetZ")) ? FCString::Atod(*P[TEXT("OffsetZ")]) : 0.0;
	const double OffsetY = P.Contains(TEXT("OffsetY")) ? FCString::Atod(*P[TEXT("OffsetY")]) : 0.0;
	const double Margin = P.Contains(TEXT("Margin")) ? FCString::Atod(*P[TEXT("Margin")]) : 1.05;
	const double CraniumDepth = P.Contains(TEXT("Cranium")) ? FCString::Atod(*P[TEXT("Cranium")]) : 0.3;
	const bool bDryRun = Switches.Contains(TEXT("DryRun"));
	// Liaison : "transfer" (defaut) = poids copies des sommets du corps les plus proches, sommets sous la peau repousses ;
	// "rigid" = tout sur l'os -Bone (premiere version ; les meches longues traversent epaules et poitrine quand la tete
	// tourne).
	const bool bRigid = P.FindRef(TEXT("Skin")).Equals(TEXT("rigid"), ESearchCase::IgnoreCase);
	const double Gap = P.Contains(TEXT("Gap")) ? FCString::Atod(*P[TEXT("Gap")]) : 0.3;
	const int32 Neighbors = P.Contains(TEXT("Neighbors")) ? FCString::Atoi(*P[TEXT("Neighbors")]) : 6;
	if (StaticPath.IsEmpty() || FitHeadPath.IsEmpty() || OutPath.IsEmpty())
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("Usage : -Static=<StaticMesh> -FitHead=<StaticMesh> -Out=</Game/...> [-Body=] [-RefHair=] [-Bone=head] [-Scale=] [-OffsetZ=] [-OffsetY=] [-DryRun]"));
		return 1;
	}

	UStaticMesh* Static = LoadByPath<UStaticMesh>(StaticPath);
	UStaticMesh* FitHead = LoadByPath<UStaticMesh>(FitHeadPath);
	USkeletalMesh* Body = LoadByPath<USkeletalMesh>(BodyPath);
	USkeletalMesh* RefHair = LoadByPath<USkeletalMesh>(RefHairPath);
	if (!Static || !FitHead || !Body || !RefHair || !RefHair->GetSkeleton())
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("Asset introuvable : Static=%d FitHead=%d Body=%d RefHair=%d"), !!Static, !!FitHead, !!Body, RefHair && RefHair->GetSkeleton());
		return 1;
	}
	if (FPackageName::DoesPackageExist(OutPath))
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("%s existe deja : rien n'est ecrase"), *OutPath);
		return 1;
	}

	// 1. Mesures.
	const int32 BodyHead = Body->GetRefSkeleton().FindBoneIndex(BoneName);
	const FReferenceSkeleton& HairRef = RefHair->GetSkeleton()->GetReferenceSkeleton();
	const int32 HairHead = HairRef.FindBoneIndex(BoneName);
	if (BodyHead == INDEX_NONE || HairHead == INDEX_NONE)
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("Os '%s' absent du corps (%d) ou du squelette de reference (%d)"), *BoneName.ToString(), BodyHead, HairHead);
		return 1;
	}
	const TArray<FVector> BodyPointsList = HeadPoints(Body, BodyHead);
	const TArray<FVector> FitPointsList = StaticPoints(FitHead);
	const int32 HeadVerts = BodyPointsList.Num();
	if (HeadVerts == 0 || FitPointsList.Num() == 0)
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("Mesure impossible : %d sommets de tete sur le corps, %d sur la tete de reference"), HeadVerts, FitPointsList.Num());
		return 1;
	}
	// Ajustement sur le haut du crane (largeur, profondeur, centre) et le sommet : la boite complete inclut oreilles,
	// nez et cou, et donnait une coiffure trop etroite enfoncee dans le crane (premier essai O7e).
	const FBox BodyFull(BodyPointsList);
	const FBox FitFull(FitPointsList);
	const FBox BodyBox = CraniumBox(BodyPointsList, CraniumDepth);
	const FBox FitBox = CraniumBox(FitPointsList, CraniumDepth);
	// Echelle par axe (largeur X, profondeur Y), hauteur : moyenne des deux ; -Uniform : echelle unique (largeur).
	// Marge : la calotte reste au-dessus du cuir chevelu.
	const double ScaleX = BodyBox.GetSize().X / FitBox.GetSize().X;
	const double ScaleY = Switches.Contains(TEXT("Uniform")) ? ScaleX : BodyBox.GetSize().Y / FitBox.GetSize().Y;
	const FVector Scale = FVector(ScaleX, ScaleY, (ScaleX + ScaleY) * 0.5) * Margin * ExtraScale;
	const FVector From = TopCenter(FitBox);
	const FVector To = TopCenter(BodyBox) + FVector(0.0, OffsetY, OffsetZ);
	const FTransform BodyHeadCS = FAnimationRuntime::GetComponentSpaceTransformRefPose(Body->GetRefSkeleton(), BodyHead);
	const FTransform HairHeadCS = FAnimationRuntime::GetComponentSpaceTransformRefPose(HairRef, HairHead);
	const FTransform Correction = BodyHeadCS.Inverse() * HairHeadCS;

	UE_LOG(LogLPHairFromStatic, Warning, TEXT("Tete du corps (%d sommets) : taille %s ; haut du crane %s, taille %s"), HeadVerts, *BodyFull.GetSize().ToString(), *BodyBox.ToString(), *BodyBox.GetSize().ToString());
	UE_LOG(LogLPHairFromStatic, Warning, TEXT("Tete de reference : taille %s ; haut du crane %s, taille %s"), *FitFull.GetSize().ToString(), *FitBox.ToString(), *FitBox.GetSize().ToString());
	UE_LOG(LogLPHairFromStatic, Warning, TEXT("Ajustement : echelle %s, %s -> %s"), *Scale.ToString(), *From.ToString(), *To.ToString());
	UE_LOG(LogLPHairFromStatic, Warning, TEXT("Os '%s' : corps %s, squelette cible %s, correction %s"), *BoneName.ToString(),
		*BodyHeadCS.GetLocation().ToString(), *HairHeadCS.GetLocation().ToString(), Correction.Equals(FTransform::Identity, 0.01) ? TEXT("aucune") : *Correction.ToString());

	// 2. Copie de travail de la coiffure, sommets places sur la tete du corps.
	UStaticMesh* Work = DuplicateObject<UStaticMesh>(Static, GetTransientPackage(), NAME_None);
	FMeshDescription* Desc = Work ? Work->GetMeshDescription(0) : nullptr;
	if (!Desc)
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("Donnees sources de %s indisponibles"), *StaticPath);
		return 1;
	}
	FBox Placed(ForceInit);
	TVertexAttributesRef<FVector3f> Positions = FStaticMeshAttributes(*Desc).GetVertexPositions();
	for (const FVertexID Vertex : Desc->Vertices().GetElementIDs())
	{
		const FVector OnBody = (FVector(Positions[Vertex]) - From) * Scale + To;
		const FVector Final = Correction.TransformPosition(OnBody);
		Positions[Vertex] = FVector3f(Final);
		Placed += Final;
	}
	UE_LOG(LogLPHairFromStatic, Warning, TEXT("Coiffure placee : %s (avant : %s)"), *Placed.ToString(), *Static->GetBoundingBox().ToString());

	// 2b. Transfert de poids depuis le corps et sortie de la peau.
	TMap<FVertexID, TArray<UE::AnimationCore::FBoneWeight>> VertexWeights;
	if (!bRigid)
	{
		const TArray<FBodySample> Samples = BodySamples(Body, HairRef, Placed.Min.Z - 15.0);
		int32 Pushed = 0;
		TMap<int32, double> BoneShare;
		for (const FVertexID Vertex : Desc->Vertices().GetElementIDs())
		{
			FVector Position(Positions[Vertex]);
			// k plus proches sommets du corps (recherche exhaustive : quelques milliers x quelques dizaines de milliers).
			TArray<TPair<double, int32>, TInlineAllocator<16>> Nearest;
			for (int32 i = 0; i < Samples.Num(); ++i)
			{
				const double D2 = FVector::DistSquared(Position, Samples[i].Position);
				if (Nearest.Num() < Neighbors || D2 < Nearest.Last().Key)
				{
					int32 At = Nearest.Num();
					while (At > 0 && Nearest[At - 1].Key > D2) { --At; }
					Nearest.Insert(TPair<double, int32>(D2, i), At);
					if (Nearest.Num() > Neighbors) { Nearest.Pop(EAllowShrinking::No); }
				}
			}
			if (Nearest.Num() == 0)
			{
				continue;
			}
			// Sortie de la peau : un sommet a moins de Gap devant la surface (ou derriere) est repousse le long de la normale.
			const FBodySample& Closest = Samples[Nearest[0].Value];
			const double Depth = FVector::DotProduct(Position - Closest.Position, Closest.Normal);
			if (Depth < Gap && FMath::Sqrt(Nearest[0].Key) < 4.0)
			{
				Position += Closest.Normal * (Gap - Depth);
				Positions[Vertex] = FVector3f(Position);
				++Pushed;
			}
			// Poids : moyenne des voisins ponderee par l'inverse de la distance, 4 influences au plus.
			TMap<int32, float> Accum;
			for (const TPair<double, int32>& N : Nearest)
			{
				const float W = 1.f / (float(FMath::Sqrt(N.Key)) + 0.5f);
				for (const TPair<int32, float>& BW : Samples[N.Value].Weights)
				{
					Accum.FindOrAdd(BW.Key) += BW.Value * W;
				}
			}
			Accum.ValueSort([](float A, float B) { return A > B; });
			// Normalisation avant creation des poids : FBoneWeight borne sa valeur a 1.
			TArray<TPair<int32, float>, TInlineAllocator<4>> Top;
			float Total = 0.f;
			for (const TPair<int32, float>& A : Accum)
			{
				if (Top.Num() == 4) { break; }
				Top.Add(A);
				Total += A.Value;
			}
			TArray<UE::AnimationCore::FBoneWeight> Weights;
			for (const TPair<int32, float>& T : Top)
			{
				Weights.Add(UE::AnimationCore::FBoneWeight(FBoneIndexType(T.Key), T.Value / Total));
				BoneShare.FindOrAdd(T.Key) += T.Value / Total;
			}
			VertexWeights.Add(Vertex, MoveTemp(Weights));
		}
		BoneShare.ValueSort([](double A, double B) { return A > B; });
		FString Shares;
		for (const TPair<int32, double>& S : BoneShare)
		{
			Shares += FString::Printf(TEXT("%s %.0f%%, "), *HairRef.GetBoneName(S.Key).ToString(), 100.0 * S.Value / FMath::Max(1, VertexWeights.Num()));
		}
		UE_LOG(LogLPHairFromStatic, Warning, TEXT("Transfert de poids : %d sommets du corps, %d sommets de coiffure, %d repousse(s) hors de la peau (%.1f mm) ; os : %s"),
			Samples.Num(), VertexWeights.Num(), Pushed, Gap * 10.0, *Shares.LeftChop(2));
	}
	if (bDryRun)
	{
		UE_LOG(LogLPHairFromStatic, Warning, TEXT("-DryRun : rien n'est cree"));
		return 0;
	}
	if (bRigid)
	{
		Work->CommitMeshDescription(0);
		Work->Build(true);
	}

	// 3. Squelette (copie de celui de la coiffure de reference) et SkeletalMesh.
	const FString OutName = FPackageName::GetShortName(OutPath);
	UPackage* SkeletonPackage = CreatePackage(*(OutPath + TEXT("_Skeleton")));
	USkeleton* Skeleton = DuplicateObject<USkeleton>(RefHair->GetSkeleton(), SkeletonPackage, *(OutName + TEXT("_Skeleton")));
	Skeleton->SetFlags(RF_Public | RF_Standalone);
	Skeleton->SetPreviewMesh(nullptr);

	UPackage* MeshPackage = CreatePackage(*OutPath);
	USkeletalMesh* Hair = NewObject<USkeletalMesh>(MeshPackage, *OutName, RF_Public | RF_Standalone);
	bool bConverted = false;
	if (bRigid)
	{
		bConverted = FStaticToSkeletalMeshConverter::InitializeSkeletalMeshFromStaticMesh(Hair, Work, Skeleton->GetReferenceSkeleton(), BoneName);
	}
	else
	{
		// Description du maillage (positions placees et repoussees) + poids de peau ; materiaux de la coiffure statique.
		FMeshDescription SkinnedDesc(*Desc);
		FSkeletalMeshAttributes Attributes(SkinnedDesc);
		Attributes.Register(true);
		FSkinWeightsVertexAttributesRef SkinWeights = Attributes.GetVertexSkinWeights();
		const int32 HeadIndex = HairRef.FindBoneIndex(BoneName);
		for (const FVertexID Vertex : SkinnedDesc.Vertices().GetElementIDs())
		{
			const TArray<UE::AnimationCore::FBoneWeight>* Weights = VertexWeights.Find(Vertex);
			if (Weights && Weights->Num() > 0)
			{
				SkinWeights.Set(Vertex, *Weights);
			}
			else
			{
				const UE::AnimationCore::FBoneWeight HeadOnly(FBoneIndexType(HeadIndex), 1.0f);
				SkinWeights.Set(Vertex, MakeArrayView(&HeadOnly, 1));
			}
		}
		TArray<FSkeletalMaterial> Materials;
		for (const FStaticMaterial& M : Static->GetStaticMaterials())
		{
			Materials.Add(FSkeletalMaterial(M.MaterialInterface, M.MaterialSlotName, M.ImportedMaterialSlotName));
		}
		const FMeshDescription* Descs[] = { &SkinnedDesc };
		bConverted = FStaticToSkeletalMeshConverter::InitializeSkeletalMeshFromMeshDescriptions(Hair, Descs, Materials, Skeleton->GetReferenceSkeleton(), false, true);
	}
	if (!bConverted)
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("Conversion en SkeletalMesh refusee par le moteur"));
		return 1;
	}
	Hair->SetSkeleton(Skeleton);
	Skeleton->MergeAllBonesToBoneTree(Hair);
	Skeleton->SetPreviewMesh(Hair);

	// 4. Controle puis sauvegarde.
	const int32 Bones = Hair->GetRefSkeleton().GetNum();
	const FBox Bounds = Hair->GetImportedBounds().GetBox();
	UE_LOG(LogLPHairFromStatic, Warning, TEXT("%s : %d os, %d materiau(x), bornes %s"), *OutName, Bones, Hair->GetMaterials().Num(), *Bounds.ToString());
	if (Bones != HairRef.GetNum() || Hair->GetRefSkeleton().FindBoneIndex(BoneName) == INDEX_NONE)
	{
		UE_LOG(LogLPHairFromStatic, Error, TEXT("Squelette inattendu : rien n'est sauvegarde"));
		return 1;
	}
	return SaveAsset(Skeleton) && SaveAsset(Hair) ? 0 : 1;
}