// Copyright IHY. Vertical Partition Streaming plugin.
#include "VerticalPartitionRuntimeSubsystem.h"
#include "VerticalPartitionDescriptor.h"
#include "VerticalPartitionSettings.h"
#include "VerticalPartitionLog.h"
#include "Engine/World.h"
#include "Engine/Level.h"
#include "Engine/LevelStreaming.h"
#include "Engine/LevelStreamingDynamic.h"
#include "Engine/StaticMesh.h"
#include "Engine/StaticMeshActor.h"
#include "Components/StaticMeshComponent.h"
#include "GameFramework/Actor.h"

UVerticalPartitionRuntimeSubsystem* UVerticalPartitionRuntimeSubsystem::Get(const UObject* WorldContextObject)
{
	if (!WorldContextObject)
	{
		return nullptr;
	}
	if (const UWorld* World = WorldContextObject->GetWorld())
	{
		return World->GetSubsystem<UVerticalPartitionRuntimeSubsystem>();
	}
	return nullptr;
}

void UVerticalPartitionRuntimeSubsystem::Deinitialize()
{
	UnregisterDescriptor();
	Super::Deinitialize();
}

void UVerticalPartitionRuntimeSubsystem::RegisterDescriptor(UVerticalPartitionDescriptor* InDescriptor, const FBox& InVolumeBounds)
{
	UnregisterDescriptor();

	Descriptor = InDescriptor;
	if (!Descriptor)
	{
		return;
	}

	EffectiveSettings = Descriptor->Settings;
	VolumeBounds = InVolumeBounds.IsValid ? InVolumeBounds : Descriptor->VolumeBounds;

	Cells.Reset();
	Cells.Reserve(Descriptor->Cells.Num());

	for (int32 i = 0; i < Descriptor->Cells.Num(); ++i)
	{
		const FVerticalCellDescriptor& D = Descriptor->Cells[i];

		FVerticalCellRuntime RT;
		RT.Id = D.CellId;
		RT.DescriptorIndex = i;
		RT.Bounds = D.Bounds;
		RT.EstimatedMemoryMB = D.EstimatedMemoryMB;
		RT.bUsesStreamingBackend = D.FullCellPackage.IsValid();

		if (!RT.bUsesStreamingBackend)
		{
			// NonDestructive: the actors are already in the persistent level. Resolve
			// them once and start from "Full" (they are currently visible).
			RT.ResolvedActors.Reserve(D.SourceActors.Num());
			for (const FSoftObjectPath& Path : D.SourceActors)
			{
				if (AActor* Actor = Cast<AActor>(Path.ResolveObject()))
				{
					RT.ResolvedActors.Add(Actor);
				}
			}
			RT.State = EVerticalCellState::Full;
		}
		else
		{
			RT.State = EVerticalCellState::Unloaded;
		}

		Cells.Add(RT.Id, MoveTemp(RT));
	}

	UE_LOG(LogVerticalPartition, Log, TEXT("[VP] Registered descriptor '%s' with %d cells (%s backend)."),
		*Descriptor->GetName(), Cells.Num(), Descriptor->HasFullPackages() ? TEXT("streaming") : TEXT("visibility"));
}

void UVerticalPartitionRuntimeSubsystem::UnregisterDescriptor()
{
	for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		FVerticalCellRuntime& C = Pair.Value;
		// Restore persistent actors so the level looks normal after PIE/unregister.
		if (!C.bUsesStreamingBackend)
		{
			SetActorGroupActive(C, true, true, true);
		}
		if (C.ProxyActor.IsValid())
		{
			C.ProxyActor->Destroy();
		}
		if (C.StreamingLevel.IsValid())
		{
			C.StreamingLevel->SetShouldBeVisible(false);
			C.StreamingLevel->SetShouldBeLoaded(false);
		}
	}
	Cells.Reset();
	Descriptor = nullptr;
	Stats = FVerticalRuntimeStats();
	bShowOnlyMode = false;
}

// ---------------------------------------------------------------------------
// Main update
// ---------------------------------------------------------------------------

void UVerticalPartitionRuntimeSubsystem::UpdateStreaming(const FVerticalStreamingContext& Ctx, float /*DeltaSeconds*/, bool bSnap)
{
	if (!Descriptor || Cells.Num() == 0)
	{
		return;
	}

	const double T0 = FPlatformTime::Seconds();
	Stats.PlayerCellZ = Ctx.PlayerCellZ;

	// 1) Desired state for every cell.
	for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		FVerticalCellRuntime& C = Pair.Value;
		PollStreaming(C);   // refresh State from in-flight async ops first

		if (C.bForced)
		{
			C.Desired = C.ForcedState;
		}
		else if (bShowOnlyMode)
		{
			C.Desired = (C.Id.Z == ShowOnlyZ) ? EVerticalCellDesiredState::Full : EVerticalCellDesiredState::Unloaded;
		}
		else
		{
			C.Desired = UVerticalPartitionStatics::GetDesiredStateForCell(C.Id, Ctx, EffectiveSettings, C.State);
		}

		C.bPreloadRequested = (C.Desired == EVerticalCellDesiredState::HLOD)
			&& UVerticalPartitionStatics::ShouldPreloadFull(C.Id, Ctx, EffectiveSettings);
	}

	// 2) Enforce budgets: keep the nearest MaxFullCells as Full, nearest
	//    MaxVisibleHLODCells as HLOD, demote the rest.
	{
		TArray<FVerticalCellRuntime*> Ordered;
		Ordered.Reserve(Cells.Num());
		for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
		{
			Ordered.Add(&Pair.Value);
		}
		const int32 PlayerZ = Ctx.PlayerCellZ;
		Ordered.Sort([PlayerZ](const FVerticalCellRuntime& A, const FVerticalCellRuntime& B)
		{
			return FMath::Abs(A.Id.Z - PlayerZ) < FMath::Abs(B.Id.Z - PlayerZ);
		});

		int32 FullCount = 0;
		int32 HlodCount = 0;
		const int32 MaxFull = FMath::Max(0, EffectiveSettings.MaxFullCells);
		const int32 MaxHlod = FMath::Max(0, EffectiveSettings.MaxVisibleHLODCells);
		for (FVerticalCellRuntime* C : Ordered)
		{
			if (C->bForced) { continue; }   // forced cells ignore budgets
			if (C->Desired == EVerticalCellDesiredState::Full)
			{
				if (FullCount < MaxFull) { ++FullCount; }
				else { C->Desired = EVerticalCellDesiredState::HLOD; }
			}
			if (C->Desired == EVerticalCellDesiredState::HLOD)
			{
				if (HlodCount < MaxHlod) { ++HlodCount; }
				else { C->Desired = EVerticalCellDesiredState::Unloaded; }
			}
		}

		// 3) Apply transitions, nearest first, within the async budget.
		int32 LoadBudget = bSnap ? MAX_int32 : FMath::Max(1, EffectiveSettings.MaxAsyncLoadsPerFrame);
		int32 UnloadBudget = bSnap ? MAX_int32 : FMath::Max(1, EffectiveSettings.MaxAsyncUnloadsPerFrame);
		for (FVerticalCellRuntime* C : Ordered)
		{
			ApplyDesiredState(*C, bSnap, LoadBudget, UnloadBudget);
		}
	}

	RebuildStatsFromCells();
	Stats.LastUpdateTimeMS = (float)((FPlatformTime::Seconds() - T0) * 1000.0);
	Stats.LastStreamingTimeMS = Stats.LastUpdateTimeMS;

	if (const UVerticalPartitionDeveloperSettings* Project = GetDefault<UVerticalPartitionDeveloperSettings>())
	{
		if (Stats.LastUpdateTimeMS > Project->HitchThresholdMS)
		{
			++Stats.HitchesOverThreshold;
			UE_LOG(LogVerticalPartition, Verbose, TEXT("[VP] Streaming update hitch: %.2f ms"), Stats.LastUpdateTimeMS);
		}
	}
}

void UVerticalPartitionRuntimeSubsystem::ApplyDesiredState(FVerticalCellRuntime& C, bool bSnap, int32& LoadBudget, int32& UnloadBudget)
{
	switch (C.Desired)
	{
	case EVerticalCellDesiredState::Full:
		if (C.bUsesStreamingBackend)
		{
			if (C.State != EVerticalCellState::Full)
			{
				if (LoadBudget > 0)
				{
					LoadCellFullAsync(C, /*bVisible*/true);
					--LoadBudget;
				}
			}
			else
			{
				HideProxy(C);
			}
		}
		else
		{
			SetActorGroupActive(C, /*bVisible*/true, /*bCollision*/true, /*bTick*/true);
			HideProxy(C);
			C.State = EVerticalCellState::Full;
		}
		break;

	case EVerticalCellDesiredState::HLOD:
		ShowCellHLOD(C);
		if (C.bUsesStreamingBackend)
		{
			if (C.bPreloadRequested)
			{
				if (C.State != EVerticalCellState::HLOD && C.State != EVerticalCellState::PreloadingFull && LoadBudget > 0)
				{
					LoadCellFullAsync(C, /*bVisible*/false);   // resident, hidden
					--LoadBudget;
				}
			}
			else if (C.StreamingLevel.IsValid() && UnloadBudget > 0)
			{
				C.StreamingLevel->SetShouldBeVisible(false);
				C.StreamingLevel->SetShouldBeLoaded(false);   // free full data while showing proxy
				--UnloadBudget;
			}
			if (C.State == EVerticalCellState::Full || C.State == EVerticalCellState::Unloaded)
			{
				C.State = EVerticalCellState::HLOD;
			}
		}
		else
		{
			SetActorGroupActive(C, /*bVisible*/false, /*bCollision*/false, /*bTick*/false);
			C.State = EVerticalCellState::HLOD;
		}
		break;

	case EVerticalCellDesiredState::Unloaded:
	default:
		if (C.bUsesStreamingBackend)
		{
			if (C.State != EVerticalCellState::Unloaded && UnloadBudget > 0)
			{
				UnloadCell(C);
				--UnloadBudget;
			}
		}
		else
		{
			UnloadCell(C);
		}
		break;
	}
}

void UVerticalPartitionRuntimeSubsystem::LoadCellFullAsync(FVerticalCellRuntime& C, bool bVisible)
{
	ULevelStreaming* L = EnsureStreamingLevel(C);
	if (!L)
	{
		C.State = EVerticalCellState::Error;
		return;
	}
	L->SetShouldBeLoaded(true);
	L->SetShouldBeVisible(bVisible);
	C.State = bVisible ? EVerticalCellState::LoadingFull : EVerticalCellState::PreloadingFull;
}

void UVerticalPartitionRuntimeSubsystem::ShowCellHLOD(FVerticalCellRuntime& C)
{
	if (AStaticMeshActor* Proxy = EnsureProxy(C))
	{
		Proxy->SetActorHiddenInGame(false);
	}
}

void UVerticalPartitionRuntimeSubsystem::UnloadCell(FVerticalCellRuntime& C)
{
	if (C.bUsesStreamingBackend)
	{
		if (C.StreamingLevel.IsValid())
		{
			C.StreamingLevel->SetShouldBeVisible(false);
			C.StreamingLevel->SetShouldBeLoaded(false);
			C.State = EVerticalCellState::Unloading;
		}
		else
		{
			C.State = EVerticalCellState::Unloaded;
		}
	}
	else
	{
		SetActorGroupActive(C, /*bVisible*/false, /*bCollision*/false, /*bTick*/false);
		C.State = EVerticalCellState::Unloaded;
	}
	HideProxy(C);
}

void UVerticalPartitionRuntimeSubsystem::PollStreaming(FVerticalCellRuntime& C)
{
	if (!C.bUsesStreamingBackend || !C.StreamingLevel.IsValid())
	{
		return;
	}
	ULevelStreaming* L = C.StreamingLevel.Get();
	const ELevelStreamingState St = L->GetLevelStreamingState();
	const bool bVisible = (St == ELevelStreamingState::LoadedVisible);
	const bool bLoaded = bVisible
		|| St == ELevelStreamingState::LoadedNotVisible
		|| St == ELevelStreamingState::MakingVisible
		|| St == ELevelStreamingState::MakingInvisible;

	if (St == ELevelStreamingState::FailedToLoad)
	{
		C.State = EVerticalCellState::Error;
		return;
	}

	if (C.Desired == EVerticalCellDesiredState::Full)
	{
		if (bVisible) { C.State = EVerticalCellState::Full; HideProxy(C); }
		else { C.State = EVerticalCellState::LoadingFull; }
	}
	else if (C.Desired == EVerticalCellDesiredState::HLOD)
	{
		C.State = (C.bPreloadRequested && !bLoaded) ? EVerticalCellState::PreloadingFull : EVerticalCellState::HLOD;
	}
	else
	{
		C.State = bLoaded ? EVerticalCellState::Unloading : EVerticalCellState::Unloaded;
	}
}

// ---------------------------------------------------------------------------
// Backend helpers
// ---------------------------------------------------------------------------

void UVerticalPartitionRuntimeSubsystem::SetActorGroupActive(FVerticalCellRuntime& C, bool bVisible, bool bCollision, bool bTick)
{
	for (const TWeakObjectPtr<AActor>& Weak : C.ResolvedActors)
	{
		if (AActor* Actor = Weak.Get())
		{
			Actor->SetActorHiddenInGame(!bVisible);
			Actor->SetActorEnableCollision(bCollision);
			Actor->SetActorTickEnabled(bTick);
		}
	}
}

AStaticMeshActor* UVerticalPartitionRuntimeSubsystem::EnsureProxy(FVerticalCellRuntime& C)
{
	if (C.ProxyActor.IsValid())
	{
		return C.ProxyActor.Get();
	}
	if (!Descriptor || !Descriptor->Cells.IsValidIndex(C.DescriptorIndex))
	{
		return nullptr;
	}
	const FVerticalCellDescriptor& D = Descriptor->Cells[C.DescriptorIndex];
	if (!D.HLODProxyAsset.IsValid())
	{
		if (!C.bProxyMissingWarned)
		{
			UE_LOG(LogVerticalPartition, Verbose, TEXT("[VP] Cell %s has no HLOD proxy; far cell will simply hide."), *C.Id.ToString());
			C.bProxyMissingWarned = true;
		}
		return nullptr;
	}

	UStaticMesh* Mesh = Cast<UStaticMesh>(D.HLODProxyAsset.TryLoad());
	UWorld* World = GetWorld();
	if (!Mesh || !World)
	{
		return nullptr;
	}

	FActorSpawnParameters Params;
	Params.ObjectFlags |= RF_Transient;
	Params.SpawnCollisionHandlingOverride = ESpawnActorCollisionHandlingMethod::AlwaysSpawn;
	// Proxy meshes are baked with bPivotPointAtZero (world-space verts), so the actor
	// goes at the world origin and the geometry lands exactly over the source actors.
	AStaticMeshActor* Actor = World->SpawnActor<AStaticMeshActor>(FVector::ZeroVector, FRotator::ZeroRotator, Params);
	if (!Actor)
	{
		return nullptr;
	}
	if (UStaticMeshComponent* Comp = Actor->GetStaticMeshComponent())
	{
		Comp->SetMobility(EComponentMobility::Movable);
		Comp->SetStaticMesh(Mesh);
		Comp->SetCollisionEnabled(ECollisionEnabled::NoCollision);
		Comp->SetCanEverAffectNavigation(false);
	}
	Actor->SetActorHiddenInGame(true);
	C.ProxyActor = Actor;
	return Actor;
}

void UVerticalPartitionRuntimeSubsystem::HideProxy(FVerticalCellRuntime& C)
{
	if (C.ProxyActor.IsValid())
	{
		C.ProxyActor->SetActorHiddenInGame(true);
	}
}

ULevelStreaming* UVerticalPartitionRuntimeSubsystem::EnsureStreamingLevel(FVerticalCellRuntime& C)
{
	if (C.StreamingLevel.IsValid())
	{
		return C.StreamingLevel.Get();
	}
	UWorld* World = GetWorld();
	if (!World || !Descriptor || !Descriptor->Cells.IsValidIndex(C.DescriptorIndex))
	{
		return nullptr;
	}
	const FVerticalCellDescriptor& D = Descriptor->Cells[C.DescriptorIndex];
	const FString PackageName = D.FullCellPackage.GetLongPackageName();
	if (PackageName.IsEmpty())
	{
		return nullptr;
	}
	const FName PackageFName(*PackageName);

	for (ULevelStreaming* Existing : World->GetStreamingLevels())
	{
		if (Existing && Existing->GetWorldAssetPackageFName() == PackageFName)
		{
			C.StreamingLevel = Existing;
			return Existing;
		}
	}

	ULevelStreamingDynamic* NewStream = NewObject<ULevelStreamingDynamic>(World, ULevelStreamingDynamic::StaticClass(), NAME_None, RF_Transient);
	NewStream->SetWorldAssetByPackageName(PackageFName);
	NewStream->SetShouldBeLoaded(false);
	NewStream->SetShouldBeVisible(false);
	NewStream->bShouldBlockOnLoad = false;
	World->AddStreamingLevel(NewStream);
	C.StreamingLevel = NewStream;
	return NewStream;
}

void UVerticalPartitionRuntimeSubsystem::RebuildStatsFromCells()
{
	Stats.FullCellsLoaded = 0;
	Stats.HLODCellsVisible = 0;
	Stats.CellsUnloaded = 0;
	Stats.CellsLoading = 0;
	Stats.ActiveAsyncLoads = 0;
	Stats.EstimatedMemoryMB = 0.f;

	for (const TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		const FVerticalCellRuntime& C = Pair.Value;
		switch (C.State)
		{
		case EVerticalCellState::Full:
			++Stats.FullCellsLoaded;
			Stats.EstimatedMemoryMB += C.EstimatedMemoryMB;
			break;
		case EVerticalCellState::HLOD:
			++Stats.HLODCellsVisible;
			break;
		case EVerticalCellState::LoadingFull:
		case EVerticalCellState::LoadingHLOD:
		case EVerticalCellState::PreloadingFull:
			++Stats.CellsLoading;
			++Stats.ActiveAsyncLoads;
			Stats.EstimatedMemoryMB += C.EstimatedMemoryMB; // resident while preloading
			break;
		case EVerticalCellState::Unloading:
			++Stats.CellsLoading;
			break;
		case EVerticalCellState::Unloaded:
		default:
			++Stats.CellsUnloaded;
			break;
		}
	}
}

FVerticalCellRuntime* UVerticalPartitionRuntimeSubsystem::FindCellByZ(int32 Z)
{
	for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		if (Pair.Value.Id.Z == Z)
		{
			return &Pair.Value;
		}
	}
	return nullptr;
}

// ---------------------------------------------------------------------------
// Live overrides + debug forcing
// ---------------------------------------------------------------------------

void UVerticalPartitionRuntimeSubsystem::SetFullRange(int32 Below, int32 Above)
{
	EffectiveSettings.FullLoadCellsBelow = FMath::Max(0, Below);
	EffectiveSettings.FullLoadCellsAbove = FMath::Max(0, Above);
	UE_LOG(LogVerticalPartition, Log, TEXT("[VP] Full range set to below=%d above=%d"), Below, Above);
}

void UVerticalPartitionRuntimeSubsystem::SetHLODRange(int32 Below, int32 Above)
{
	EffectiveSettings.HLODLoadCellsBelow = FMath::Max(0, Below);
	EffectiveSettings.HLODLoadCellsAbove = FMath::Max(0, Above);
	UE_LOG(LogVerticalPartition, Log, TEXT("[VP] HLOD range set to below=%d above=%d"), Below, Above);
}

void UVerticalPartitionRuntimeSubsystem::SetCellHeight(float NewHeight)
{
	EffectiveSettings.CellHeight = FMath::Max(1.f, NewHeight);
	UE_LOG(LogVerticalPartition, Warning, TEXT("[VP] CellHeight changed at runtime to %.0f. Note: cell boundaries are baked in the descriptor; this only affects velocity prediction. Rebuild to re-segment."), NewHeight);
}

void UVerticalPartitionRuntimeSubsystem::ForceCellFull(int32 Z)
{
	for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		if (Pair.Value.Id.Z == Z) { Pair.Value.bForced = true; Pair.Value.ForcedState = EVerticalCellDesiredState::Full; }
	}
}

void UVerticalPartitionRuntimeSubsystem::ForceCellHLOD(int32 Z)
{
	for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		if (Pair.Value.Id.Z == Z) { Pair.Value.bForced = true; Pair.Value.ForcedState = EVerticalCellDesiredState::HLOD; }
	}
}

void UVerticalPartitionRuntimeSubsystem::UnforceCell(int32 Z)
{
	for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		if (Pair.Value.Id.Z == Z) { Pair.Value.bForced = false; }
	}
}

void UVerticalPartitionRuntimeSubsystem::UnforceAll()
{
	for (TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		Pair.Value.bForced = false;
	}
	bShowOnlyMode = false;
}

void UVerticalPartitionRuntimeSubsystem::ShowOnlyCell(int32 Z)
{
	bShowOnlyMode = true;
	ShowOnlyZ = Z;
}

// ---------------------------------------------------------------------------
// Reporting
// ---------------------------------------------------------------------------

void UVerticalPartitionRuntimeSubsystem::DumpRuntimeStats() const
{
	UE_LOG(LogVerticalPartition, Log, TEXT("[VP] ==== Runtime Stats ===="));
	UE_LOG(LogVerticalPartition, Log, TEXT("[VP] PlayerCellZ=%d Full=%d HLOD=%d Unloaded=%d Loading=%d Async=%d Mem~%.1fMB Update=%.2fms Hitches=%d"),
		Stats.PlayerCellZ, Stats.FullCellsLoaded, Stats.HLODCellsVisible, Stats.CellsUnloaded,
		Stats.CellsLoading, Stats.ActiveAsyncLoads, Stats.EstimatedMemoryMB, Stats.LastUpdateTimeMS, Stats.HitchesOverThreshold);
}

void UVerticalPartitionRuntimeSubsystem::DumpCells() const
{
	UE_LOG(LogVerticalPartition, Log, TEXT("[VP] ==== Cells (%d) ===="), Cells.Num());
	for (const TPair<FVerticalCellId, FVerticalCellRuntime>& Pair : Cells)
	{
		const FVerticalCellRuntime& C = Pair.Value;
		UE_LOG(LogVerticalPartition, Log, TEXT("[VP]  %s state=%d desired=%d actors=%d backend=%s forced=%d"),
			*C.Id.ToString(), (int32)C.State, (int32)C.Desired, C.ResolvedActors.Num(),
			C.bUsesStreamingBackend ? TEXT("stream") : TEXT("vis"), C.bForced ? 1 : 0);
	}
}

void UVerticalPartitionRuntimeSubsystem::DumpWarnings() const
{
	if (!Descriptor)
	{
		UE_LOG(LogVerticalPartition, Log, TEXT("[VP] No descriptor registered."));
		return;
	}
	const FVerticalPartitionBuildReport& R = Descriptor->LastReport;
	UE_LOG(LogVerticalPartition, Log, TEXT("[VP] ==== Build Warnings (%d) / Errors (%d) ===="), R.Warnings.Num(), R.Errors.Num());
	for (const FVerticalPartitionWarning& W : R.Warnings)
	{
		UE_LOG(LogVerticalPartition, Warning, TEXT("[VP]  WARN %s :: %s"),
			*UVerticalPartitionStatics::WarningTypeToString(W.Type), *W.Message);
	}
	for (const FVerticalPartitionError& E : R.Errors)
	{
		UE_LOG(LogVerticalPartition, Error, TEXT("[VP]  ERR  %s :: %s"),
			*UVerticalPartitionStatics::WarningTypeToString(E.Type), *E.Message);
	}
}