// Copyright IHY.
#include "OptimizerSubsystem.h"

#include "OptimizerGeometry.h"
#include "OptimizerMatcher.h"
#include "OptimizerReconciler.h"

#include "Editor.h"
#include "Subsystems/EditorActorSubsystem.h"
#include "Engine/StaticMesh.h"
#include "Engine/StaticMeshActor.h"
#include "Engine/World.h"
#include "GameFramework/Actor.h"
#include "Components/StaticMeshComponent.h"
#include "Components/InstancedStaticMeshComponent.h"
#include "Components/HierarchicalInstancedStaticMeshComponent.h"
#include "ScopedTransaction.h"

DEFINE_LOG_CATEGORY_STATIC(LogOptimizer, Log, All);

namespace
{
	UStaticMesh* ChooseCanonical(
		const TArray<UStaticMesh*>& Members,
		const TMap<UStaticMesh*, FOptMeshGeom>& Geoms,
		const TMap<UStaticMesh*, int32>& Counts,
		EOptimizerCanonicalPolicy Policy)
	{
		UStaticMesh* Best = Members.Num() > 0 ? Members[0] : nullptr;
		if (!Best)
		{
			return nullptr;
		}
		switch (Policy)
		{
		case EOptimizerCanonicalPolicy::LowestVertexCount:
			for (UStaticMesh* M : Members)
			{
				const FOptMeshGeom* G = Geoms.Find(M);
				const FOptMeshGeom* BG = Geoms.Find(Best);
				if (G && BG && G->WeldedVertexCount < BG->WeldedVertexCount)
				{
					Best = M;
				}
			}
			break;
		case EOptimizerCanonicalPolicy::FirstAlphabetical:
			for (UStaticMesh* M : Members)
			{
				if (M->GetName() < Best->GetName())
				{
					Best = M;
				}
			}
			break;
		case EOptimizerCanonicalPolicy::MostInstances:
		default:
			for (UStaticMesh* M : Members)
			{
				const int32 C = Counts.FindRef(M);
				const int32 BC = Counts.FindRef(Best);
				if (C > BC)
				{
					Best = M;
				}
			}
			break;
		}
		return Best;
	}
}

FOptimizerScanResult UOptimizerSubsystem::ScanLevel(const FOptimizerScanSettings& Settings)
{
	LastSettings = Settings;
	Groups.Reset();

	FOptimizerScanResult Result;
	UWorld* World = GEditor ? GEditor->GetEditorWorldContext().World() : nullptr;
	if (!World)
	{
		Result.Summary = TEXT("No editor world.");
		return Result;
	}
	Result.bWorldPartitionCoverageWarning = World->IsPartitionedWorld();

	// --- 1. Gather non-instanced static-mesh placements ---
	TArray<FPlacement> Placements;
	UEditorActorSubsystem* AS = GEditor->GetEditorSubsystem<UEditorActorSubsystem>();
	TArray<AActor*> Actors;
	if (AS)
	{
		Actors = (Settings.Scope == EOptimizerScanScope::SelectedActors)
			? AS->GetSelectedLevelActors()
			: AS->GetAllLevelActors();
	}
	Result.ActorsScanned = Actors.Num();

	for (AActor* Actor : Actors)
	{
		if (!Actor)
		{
			continue;
		}
		TArray<UStaticMeshComponent*> SMCs;
		Actor->GetComponents<UStaticMeshComponent>(SMCs);
		for (UStaticMeshComponent* SMC : SMCs)
		{
			// ISM/HISM sources are out of v1 scope (avoids per-instance removal complexity).
			if (SMC->IsA<UInstancedStaticMeshComponent>())
			{
				continue;
			}
			UStaticMesh* M = SMC->GetStaticMesh();
			if (!M)
			{
				continue;
			}
			FPlacement P;
			P.Component = SMC;
			P.Mesh = M;
			P.World = SMC->GetComponentTransform();
			Placements.Add(MoveTemp(P));
		}
	}
	Result.ComponentsScanned = Placements.Num();

	// --- 2. Fingerprint every unique mesh ---
	TSet<UStaticMesh*> UniqueSet;
	for (const FPlacement& P : Placements)
	{
		if (P.Mesh.IsValid())
		{
			UniqueSet.Add(P.Mesh.Get());
		}
	}
	TArray<UStaticMesh*> Unique = UniqueSet.Array();
	Result.UniqueMeshes = Unique.Num();

	TMap<UStaticMesh*, FOptMeshGeom> Geoms;
	Geoms.Reserve(Unique.Num());
	for (UStaticMesh* M : Unique)
	{
		FOptMeshGeom G;
		if (OptimizerGeometry::ExtractGeom(M, Settings.WeldEpsilon, /*bWantPositions*/false, G))
		{
			Geoms.Add(M, MoveTemp(G));
		}
	}

	auto EnsurePositions = [&](UStaticMesh* M)
	{
		FOptMeshGeom* G = Geoms.Find(M);
		if (G && G->Positions.Num() == 0)
		{
			OptimizerGeometry::ExtractGeom(M, Settings.WeldEpsilon, /*bWantPositions*/true, *G);
		}
	};

	TMap<UStaticMesh*, int32> MeshInstanceCount;
	for (const FPlacement& P : Placements)
	{
		if (P.Mesh.IsValid())
		{
			MeshInstanceCount.FindOrAdd(P.Mesh.Get())++;
		}
	}

	// --- 3. Bucket by (welded vertex count, triangle count) ---
	TMap<TPair<int32, int32>, TArray<UStaticMesh*>> Buckets;
	for (UStaticMesh* M : Unique)
	{
		const FOptMeshGeom* G = Geoms.Find(M);
		if (G && G->bValid)
		{
			Buckets.FindOrAdd(TPair<int32, int32>(G->WeldedVertexCount, G->TriangleCount)).Add(M);
		}
	}

	const double ShearTol = FMath::Max((double)Settings.AcceptTolerance, 0.01);

	// --- 4. Recover + group within each bucket ---
	for (auto& BucketPair : Buckets)
	{
		TArray<UStaticMesh*>& Bucket = BucketPair.Value;
		if (Bucket.Num() < 2)
		{
			continue;
		}
		for (UStaticMesh* M : Bucket)
		{
			EnsurePositions(M);
		}

		const int32 K = Bucket.Num();
		TArray<int32> Parent;
		Parent.SetNum(K);
		for (int32 i = 0; i < K; ++i)
		{
			Parent[i] = i;
		}
		auto Find = [&Parent](int32 x)
		{
			while (Parent[x] != x)
			{
				Parent[x] = Parent[Parent[x]];
				x = Parent[x];
			}
			return x;
		};

		for (int32 i = 0; i < K; ++i)
		{
			for (int32 j = i + 1; j < K; ++j)
			{
				const FOptMeshGeom& Gi = Geoms[Bucket[i]];
				const FOptMeshGeom& Gj = Geoms[Bucket[j]];
				if (!OptimizerMatcher::FingerprintCompatible(Gi, Gj, Settings))
				{
					continue;
				}
				FOptDelta D;
				if (OptimizerMatcher::RecoverDelta(Gi, Gj, Settings, D))
				{
					Parent[Find(i)] = Find(j);
				}
			}
		}

		TMap<int32, TArray<UStaticMesh*>> Components;
		for (int32 i = 0; i < K; ++i)
		{
			Components.FindOrAdd(Find(i)).Add(Bucket[i]);
		}

		for (auto& CompPair : Components)
		{
			TArray<UStaticMesh*>& Members = CompPair.Value;
			if (Members.Num() < 2)
			{
				continue;
			}

			UStaticMesh* Canon = ChooseCanonical(Members, Geoms, MeshInstanceCount, Settings.CanonicalPolicy);
			if (!Canon)
			{
				continue;
			}
			EnsurePositions(Canon);
			const FOptMeshGeom& CanonGeom = Geoms[Canon];

			// Recover canon->member delta for each member (direct re-verify, no transitivity trust).
			TMap<UStaticMesh*, FOptDelta> Deltas;
			FGroup Group;
			Group.Canonical = Canon;
			for (UStaticMesh* Mem : Members)
			{
				if (Geoms[Mem].MaterialHash != CanonGeom.MaterialHash)
				{
					Group.bMaterialMismatch = true;
				}
				if (Mem == Canon)
				{
					FOptDelta Id;
					Id.bValid = true;
					Id.bIdentity = true;
					Id.CanonToMember = FMatrix::Identity;
					Deltas.Add(Mem, Id);
					continue;
				}
				EnsurePositions(Mem);
				FOptDelta D;
				if (OptimizerMatcher::RecoverDelta(CanonGeom, Geoms[Mem], Settings, D))
				{
					Group.bHasMirrored |= D.bMirrored;
					Group.bHasScaled |= D.bScaled;
					Group.MaxDeviation = FMath::Max(Group.MaxDeviation, (float)D.MaxDev);
					Deltas.Add(Mem, D);
				}
			}
			if (Deltas.Num() < 2)
			{
				continue;
			}
			for (const auto& DPair : Deltas)
			{
				Group.Members.Add(DPair.Key);
			}

			// Plan each placement whose mesh is in this group.
			for (const FPlacement& P : Placements)
			{
				UStaticMesh* PM = P.Mesh.Get();
				const FOptDelta* D = PM ? Deltas.Find(PM) : nullptr;
				if (!D)
				{
					continue;
				}
				FPlacement Plan = P;
				if (D->bIdentity)
				{
					Plan.PlannedWorld = P.World;
					Plan.bNeedsTransformFix = false;
				}
				else
				{
					FTransform Wp;
					double CornerDev = 0.0;
					if (OptimizerReconciler::ComputeCorrectedWorld(D->CanonToMember, P.World, CanonGeom.LocalBounds, ShearTol, Wp, CornerDev))
					{
						Plan.PlannedWorld = Wp;
						Plan.bNeedsTransformFix = true;
					}
					else
					{
						Plan.bShearRejected = true;
						Plan.PlannedWorld = P.World;
						Group.ShearRejected++;
					}
				}
				Group.Placements.Add(MoveTemp(Plan));
			}

			if (Group.Placements.Num() > 0)
			{
				Groups.Add(MoveTemp(Group));
			}
		}
	}

	// --- 5. Build the UI view ---
	FOptimizerScanResult View = BuildResultView();
	View.ActorsScanned = Result.ActorsScanned;
	View.ComponentsScanned = Result.ComponentsScanned;
	View.UniqueMeshes = Result.UniqueMeshes;
	View.bWorldPartitionCoverageWarning = Result.bWorldPartitionCoverageWarning;

	int32 TotalDupMeshes = 0;
	for (const FOptimizerGroupView& GV : View.Groups)
	{
		TotalDupMeshes += FMath::Max(0, GV.MemberMeshNames.Num() - 1);
	}
	View.Summary = FString::Printf(
		TEXT("%d sibling group(s): %d duplicate mesh asset(s) across %d placement(s)%s%s"),
		View.Groups.Num(), TotalDupMeshes, View.InstancesCollapsible,
		View.bWorldPartitionCoverageWarning ? TEXT(" | World Partition: only loaded cells scanned") : TEXT(""),
		View.Groups.ContainsByPredicate([](const FOptimizerGroupView& G) { return G.ShearRejectedCount > 0; })
			? TEXT(" | some placements skipped (shear)") : TEXT(""));

	UE_LOG(LogOptimizer, Log, TEXT("Optimizer scan: %s"), *View.Summary);
	return View;
}

FOptimizerScanResult UOptimizerSubsystem::BuildResultView() const
{
	FOptimizerScanResult View;
	int32 Collapsible = 0;
	for (int32 gi = 0; gi < Groups.Num(); ++gi)
	{
		const FGroup& G = Groups[gi];
		FOptimizerGroupView V;
		V.GroupId = gi;
		if (UStaticMesh* Canon = G.Canonical.Get())
		{
			V.CanonicalMeshName = Canon->GetName();
			V.CanonicalMeshPath = Canon->GetPathName();
		}
		for (const TWeakObjectPtr<UStaticMesh>& M : G.Members)
		{
			if (UStaticMesh* SM = M.Get())
			{
				V.MemberMeshNames.Add(SM->GetName());
			}
		}
		V.InstanceCount = G.Placements.Num();
		for (const FPlacement& P : G.Placements)
		{
			if (P.bNeedsTransformFix)
			{
				V.TransformFixCount++;
			}
		}
		V.MaxDeviation = G.MaxDeviation;
		V.bHasMirrored = G.bHasMirrored;
		V.bHasScaled = G.bHasScaled;
		V.bMaterialMismatch = G.bMaterialMismatch;
		V.ShearRejectedCount = G.ShearRejected;
		Collapsible += G.Placements.Num();
		View.Groups.Add(MoveTemp(V));
	}
	View.InstancesCollapsible = Collapsible;
	return View;
}

int32 UOptimizerSubsystem::ApplyUnify()
{
	if (Groups.Num() == 0)
	{
		return 0;
	}
	const FScopedTransaction Transaction(NSLOCTEXT("Optimizer", "Unify", "Optimizer: unify sibling meshes"));
	int32 Changed = 0;
	for (FGroup& G : Groups)
	{
		UStaticMesh* Canon = G.Canonical.Get();
		if (!Canon)
		{
			continue;
		}
		for (FPlacement& P : G.Placements)
		{
			if (P.bShearRejected)
			{
				continue;
			}
			UStaticMeshComponent* C = P.Component.Get();
			if (!C)
			{
				continue;
			}
			if (P.Mesh.Get() == Canon && !P.bNeedsTransformFix)
			{
				continue; // already canonical, nothing to fix
			}
			C->Modify();
			if (AActor* Owner = C->GetOwner())
			{
				Owner->Modify();
			}
			C->SetStaticMesh(Canon);
			C->SetWorldTransform(P.PlannedWorld, false, nullptr, ETeleportType::TeleportPhysics);
			C->MarkRenderStateDirty();
			if (AActor* Owner = C->GetOwner())
			{
				Owner->MarkPackageDirty();
			}
			++Changed;
		}
	}
	if (GEditor)
	{
		GEditor->RedrawLevelEditingViewports();
	}
	UE_LOG(LogOptimizer, Log, TEXT("Optimizer unify: %d placement(s) reassigned."), Changed);
	return Changed;
}

int32 UOptimizerSubsystem::BuildHISM(bool bDestroyOriginals)
{
	if (Groups.Num() == 0 || !GEditor)
	{
		return 0;
	}
	UWorld* World = GEditor->GetEditorWorldContext().World();
	if (!World)
	{
		return 0;
	}
	UEditorActorSubsystem* AS = GEditor->GetEditorSubsystem<UEditorActorSubsystem>();
	const FScopedTransaction Transaction(NSLOCTEXT("Optimizer", "HISM", "Optimizer: build HISM"));
	int32 Built = 0;
	TSet<AActor*> Destroyed;

	for (FGroup& G : Groups)
	{
		UStaticMesh* Canon = G.Canonical.Get();
		if (!Canon)
		{
			continue;
		}
		TArray<FTransform> Instances;
		for (const FPlacement& P : G.Placements)
		{
			if (!P.bShearRejected)
			{
				Instances.Add(P.PlannedWorld);
			}
		}
		if (Instances.Num() == 0)
		{
			continue;
		}

		FActorSpawnParameters Sp;
		Sp.ObjectFlags |= RF_Transactional;
		AActor* Holder = World->SpawnActor<AActor>(AActor::StaticClass(), FTransform::Identity, Sp);
		if (!Holder)
		{
			continue;
		}
		UHierarchicalInstancedStaticMeshComponent* HISM =
			NewObject<UHierarchicalInstancedStaticMeshComponent>(Holder, NAME_None, RF_Transactional);
		HISM->SetStaticMesh(Canon);
		HISM->SetMobility(EComponentMobility::Static);
		Holder->SetRootComponent(HISM);
		Holder->AddInstanceComponent(HISM);
		HISM->RegisterComponent();
		for (const FTransform& T : Instances)
		{
			HISM->AddInstance(T, /*bWorldSpace*/true);
		}
		Holder->SetActorLabel(FString::Printf(TEXT("HISM_%s"), *Canon->GetName()));
		++Built;

		if (bDestroyOriginals)
		{
			for (FPlacement& P : G.Placements)
			{
				if (P.bShearRejected)
				{
					continue;
				}
				UStaticMeshComponent* C = P.Component.Get();
				if (!C)
				{
					continue;
				}
				AActor* Owner = C->GetOwner();
				if (Owner && Owner->IsA<AStaticMeshActor>())
				{
					if (!Destroyed.Contains(Owner))
					{
						Owner->Modify();
						if (AS)
						{
							AS->DestroyActor(Owner);
						}
						else
						{
							World->DestroyActor(Owner);
						}
						Destroyed.Add(Owner);
					}
				}
				else
				{
					C->Modify();
					C->DestroyComponent();
				}
			}
		}
	}

	// Plan is now stale (originals consumed); require a re-scan before further ops.
	if (bDestroyOriginals)
	{
		Groups.Reset();
	}
	if (GEditor)
	{
		GEditor->RedrawLevelEditingViewports();
	}
	UE_LOG(LogOptimizer, Log, TEXT("Optimizer HISM: %d actor(s) built."), Built);
	return Built;
}

void UOptimizerSubsystem::ClearPlan()
{
	Groups.Reset();
}

bool UOptimizerSubsystem::HasPlan() const
{
	return Groups.Num() > 0;
}