// Copyright Pacificao. All Rights Reserved.

#include "AgrarianFoliagePatch.h"

#include "AgrarianPerformanceStats.h"
#include "Components/HierarchicalInstancedStaticMeshComponent.h"
#include "Components/SceneComponent.h"
#include "ProfilingDebugging/CpuProfilerTrace.h"

namespace
{
void ConfigureFoliageComponent(UHierarchicalInstancedStaticMeshComponent* Component, const FName CollisionProfileName)
{
	if (!Component)
	{
		return;
	}

	Component->SetMobility(EComponentMobility::Static);
	Component->SetCollisionProfileName(CollisionProfileName);
	Component->SetGenerateOverlapEvents(false);
	Component->bCastDynamicShadow = true;
	Component->bCastStaticShadow = true;
	Component->InstanceStartCullDistance = 120000;
	Component->InstanceEndCullDistance = 180000;
}
}

AAgrarianFoliagePatch::AAgrarianFoliagePatch()
{
	PrimaryActorTick.bCanEverTick = false;

	SceneRoot = CreateDefaultSubobject<USceneComponent>(TEXT("SceneRoot"));
	RootComponent = SceneRoot;
	SceneRoot->SetMobility(EComponentMobility::Static);

	TreeInstances = CreateDefaultSubobject<UHierarchicalInstancedStaticMeshComponent>(TEXT("TreeInstances"));
	TreeInstances->SetupAttachment(SceneRoot);
	ConfigureFoliageComponent(TreeInstances, TEXT("BlockAll"));

	ShrubInstances = CreateDefaultSubobject<UHierarchicalInstancedStaticMeshComponent>(TEXT("ShrubInstances"));
	ShrubInstances->SetupAttachment(SceneRoot);
	ConfigureFoliageComponent(ShrubInstances, TEXT("NoCollision"));

	GrassInstances = CreateDefaultSubobject<UHierarchicalInstancedStaticMeshComponent>(TEXT("GrassInstances"));
	GrassInstances->SetupAttachment(SceneRoot);
	ConfigureFoliageComponent(GrassInstances, TEXT("NoCollision"));
}

void AAgrarianFoliagePatch::ClearFoliage()
{
	SCOPE_CYCLE_COUNTER(STAT_AgrarianFoliageInstanceMutation);
	TRACE_CPUPROFILER_EVENT_SCOPE(AgrarianFoliageClear);

	if (TreeInstances)
	{
		TreeInstances->ClearInstances();
	}

	if (ShrubInstances)
	{
		ShrubInstances->ClearInstances();
	}

	if (GrassInstances)
	{
		GrassInstances->ClearInstances();
	}
}

int32 AAgrarianFoliagePatch::AddTreeInstance(const FTransform& InstanceTransform)
{
	SCOPE_CYCLE_COUNTER(STAT_AgrarianFoliageInstanceMutation);
	TRACE_CPUPROFILER_EVENT_SCOPE(AgrarianFoliageAddTreeInstance);

	return TreeInstances ? TreeInstances->AddInstance(InstanceTransform, true) : INDEX_NONE;
}

int32 AAgrarianFoliagePatch::AddShrubInstance(const FTransform& InstanceTransform)
{
	SCOPE_CYCLE_COUNTER(STAT_AgrarianFoliageInstanceMutation);
	TRACE_CPUPROFILER_EVENT_SCOPE(AgrarianFoliageAddShrubInstance);

	return ShrubInstances ? ShrubInstances->AddInstance(InstanceTransform, true) : INDEX_NONE;
}

int32 AAgrarianFoliagePatch::AddGrassInstance(const FTransform& InstanceTransform)
{
	SCOPE_CYCLE_COUNTER(STAT_AgrarianFoliageInstanceMutation);
	TRACE_CPUPROFILER_EVENT_SCOPE(AgrarianFoliageAddGrassInstance);

	return GrassInstances ? GrassInstances->AddInstance(InstanceTransform, true) : INDEX_NONE;
}

int32 AAgrarianFoliagePatch::GetTreeInstanceCount() const
{
	return TreeInstances ? TreeInstances->GetInstanceCount() : 0;
}

int32 AAgrarianFoliagePatch::GetShrubInstanceCount() const
{
	return ShrubInstances ? ShrubInstances->GetInstanceCount() : 0;
}

int32 AAgrarianFoliagePatch::GetGrassInstanceCount() const
{
	return GrassInstances ? GrassInstances->GetInstanceCount() : 0;
}

void AAgrarianFoliagePatch::GetFuelCountsNearLocation(
	const FVector& WorldLocation,
	float Radius,
	int32& OutGrassCount,
	int32& OutShrubCount,
	int32& OutTreeCount) const
{
	SCOPE_CYCLE_COUNTER(STAT_AgrarianFoliageInstanceMutation);
	TRACE_CPUPROFILER_EVENT_SCOPE(AgrarianFoliageFuelCountsNearLocation);

	OutGrassCount = CountInstancesNearLocation(GrassInstances, WorldLocation, Radius);
	OutShrubCount = CountInstancesNearLocation(ShrubInstances, WorldLocation, Radius);
	OutTreeCount = CountInstancesNearLocation(TreeInstances, WorldLocation, Radius);
}

float AAgrarianFoliagePatch::GetDryVegetationFuelScoreNearLocation(const FVector& WorldLocation, float Radius) const
{
	int32 GrassCount = 0;
	int32 ShrubCount = 0;
	int32 TreeCount = 0;
	GetFuelCountsNearLocation(WorldLocation, Radius, GrassCount, ShrubCount, TreeCount);

	return static_cast<float>(GrassCount) + (static_cast<float>(ShrubCount) * 2.0f) + (static_cast<float>(TreeCount) * 4.0f);
}

int32 AAgrarianFoliagePatch::CountInstancesNearLocation(
	const UHierarchicalInstancedStaticMeshComponent* Component,
	const FVector& WorldLocation,
	float Radius) const
{
	if (!Component || Radius <= 0.0f)
	{
		return 0;
	}

	const float RadiusSquared = FMath::Square(Radius);
	int32 Count = 0;
	for (int32 Index = 0; Index < Component->GetInstanceCount(); ++Index)
	{
		FTransform InstanceTransform;
		if (Component->GetInstanceTransform(Index, InstanceTransform, true)
			&& FVector::DistSquared2D(InstanceTransform.GetLocation(), WorldLocation) <= RadiusSquared)
		{
			++Count;
		}
	}

	return Count;
}
