// Copyright 2025 https://yuewu.dev/en  All Rights Reserved.

#include "Collision/GCS_TraceSubsystem.h"

#include "GCS_LogChannels.h"
#include "KismetTraceUtils.h"
#include "Async/ParallelFor.h"
#include "Components/StaticMeshComponent.h"
#include "Collision/GCS_TraceSystemComponent.h"
#include "Kismet/KismetMathLibrary.h"

#define LOCTEXT_NAMESPACE "UGCS_CollisionTraceSubsystem"

// CVars
namespace GenericCombatSystemCVars
{
#if ENABLE_DRAW_DEBUG

	static bool bEnableTraceDebugging = false;
	FAutoConsoleVariableRef CvarForceEnableTraceDebugging(
		TEXT("gcs.debug.EnableTraceDebugging"),
		bEnableTraceDebugging,
		TEXT("Toggles whether enable draw debugs for all traces. (Enabled: true, Disabled: false)"));

	static float OverrideTraceDebuggingLifeTime = 0.f;
	FAutoConsoleVariableRef CvarOverrideTraceDebuggingLifeTime(
		TEXT("gcs.debug.OverrideTraceDebuggingLifeTime"),
		OverrideTraceDebuggingLifeTime,
		TEXT("Overrides the draws life time to ease the trace debugging"));
#endif
}

void UGCS_TraceSubsystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);
	this->TraceStates.Reserve(4068);
}

void UGCS_TraceSubsystem::Tick(float DeltaTime)
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::Tick"), STAT_UGCS_TraceSubsystem_Tick, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)

	TickIdx++;
	// Lock removals so we don't get any modifications to the array while we are iterating
	RemovalLock = true;

	PreTraceTick(DeltaTime);

	PrepareSubTicks(DeltaTime);
	PerformSubTicks(DeltaTime);

	PostTraceTick();

	// Reverse iterate, remove pending removals
	RemovalLock = false;
	for (int Idx = TraceStates.Num() - 1; Idx >= 0; Idx--)
	{
		const auto& TraceState = TraceStates[Idx];
		if (!TraceStates.IsValidIndex(Idx))
		{
			continue;
		}

		if (TraceState.IsPendingRemoval)
		{
			RemoveTraceStateAt(Idx, TraceState.Handle.Guid);
		}
	}
}

void UGCS_TraceSubsystem::RemoveTraceState(const int Idx, const FGuid Guid)
{
	if (RemovalLock)
	{
		if (TraceStates.IsValidIndex(Idx) && TraceStates[Idx].Handle.Guid == Guid)
		{
			TraceStates[Idx].IsPendingRemoval = true;
		}
	}
	else
	{
		RemoveTraceStateAt(Idx, Guid);
	}
}

int32 UGCS_TraceSubsystem::AddTraceState()
{
	FScopeLock ScopeLock(&CriticalSection);
	return TraceStates.AddDefaulted();
}

bool UGCS_TraceSubsystem::IsValidStateIdx(int32 StateIdx) const
{
	return TraceStates.IsValidIndex(StateIdx);
}

FGCS_TraceState& UGCS_TraceSubsystem::GetTraceStateAt(const int Index)
{
	return TraceStates[Index];
}

void UGCS_TraceSubsystem::RemoveTraceStateAt(const int Idx, const FGuid Guid)
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::RemoveTraceStateAt"), STAT_UGCS_TraceSubsystem_RemoveTraceStateAt, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)
	FScopeLock ScopeLock(&CriticalSection);
	if (TraceStates.IsValidIndex(Idx) && TraceStates[Idx].Handle.Guid == Guid)
	{
		// 移除并交换
		TraceStates.RemoveAtSwap(Idx);

		// If a state was moved to Idx, update its corresponding handle in the OwningSystem
		if (TraceStates.IsValidIndex(Idx))
		{
			FGCS_TraceState& MovedState = TraceStates[Idx];
			if (IsValid(MovedState.OwningSystem))
			{
				// Update the Component's HandleToStateIdx map
				MovedState.OwningSystem->HandleToStateIdx.Remove(MovedState.Handle); // Remove old mapping
				MovedState.OwningSystem->HandleToStateIdx.Add(MovedState.Handle, Idx); // Add new mapping
			}
			// mark pending removal if owning system become invalid.
			else
			{
				MovedState.IsPendingRemoval = true;
			}
		}
	}
}

void UGCS_TraceSubsystem::PreTraceTick(const float DeltaTime)
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::PreTraceTick"), STAT_UGCS_TraceSubsystem_PreTraceTick, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)
	ParallelFor(TraceStates.Num(), [&](const int32 Idx)
	{
		if (!TraceStates.IsValidIndex(Idx))
		{
			return;
		}
		auto& TraceState = TraceStates[Idx];
		if (TraceState.ExecutionState == EGCS_TraceExecutionState::Stopped)
		{
			return;
		}

		check(TraceState.bShouldTickThisFrame == false)

		if (!IsValid(TraceState.SourceComponent))
		{
			TraceState.ExecutionState = EGCS_TraceExecutionState::Stopped;
			TraceState.bShouldTickThisFrame = false;
			return;
		}

		const FTransform& CurrentTransform = TraceState.GetCurrentTransform();

		if (TraceState.TransformsOverTime.IsEmpty())
		{
			TraceState.TransformsOverTime.Add(CurrentTransform);
		}

		if (TraceState.ExecutionState == EGCS_TraceExecutionState::PendingStop)
		{
			TraceState.bShouldTickThisFrame = true;
			TraceState.TransformsOverTime.Add(CurrentTransform);
			return;
		}

		switch (TraceState.TickPolicy)
		{
		case EGCS_TraceTickType::Default:
			TraceState.TransformsOverTime.Add(CurrentTransform);
			TraceState.bShouldTickThisFrame = true;
			return;
		case EGCS_TraceTickType::DistanceBased:
			{
				const FVector& PreviousLocation = TraceState.TransformsOverTime[0].GetLocation();
				const FVector& CurrentLocation = CurrentTransform.GetLocation();
				const FRotator& PreviousRotation = TraceState.TransformsOverTime[0].GetRotation().Rotator();
				const FRotator& CurrentRotation = CurrentTransform.GetRotation().Rotator();

				bool bDistanceThresholdMet = (PreviousLocation - CurrentLocation).Length() >= TraceState.TickInterval;
				bool bAngleThresholdMet = FMath::Abs(PreviousRotation.Yaw - CurrentRotation.Yaw) >= TraceState.AngleThreshold ||
					FMath::Abs(PreviousRotation.Pitch - CurrentRotation.Pitch) >= TraceState.AngleThreshold ||
					FMath::Abs(PreviousRotation.Roll - CurrentRotation.Roll) >= TraceState.AngleThreshold;

				if (bDistanceThresholdMet || bAngleThresholdMet)
				{
					TraceState.TransformsOverTime.Add(CurrentTransform);
					TraceState.bShouldTickThisFrame = true;
				}
				return;
			}

		case EGCS_TraceTickType::FixedFrameRate:
			TraceState.TimeSinceLastTick += DeltaTime;
			int32 SubTickCountPreview = FMath::FloorToInt(TraceState.TimeSinceLastTick / TraceState.TickInterval);
			if (SubTickCountPreview > 0)
			{
				TraceState.TransformsOverTime.Add(CurrentTransform);
				TraceState.bShouldTickThisFrame = true;
			}
		}
	});
}

void UGCS_TraceSubsystem::PostTraceTick()
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::PostTraceTick"), STAT_UGCS_TraceSubsystem_PostTraceTick, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)
	ParallelFor(TraceStates.Num(), [&](const int32 Idx)
	{
		auto& TraceState = TraceStates[Idx];
		if (!TraceState.bShouldTickThisFrame)
		{
			return;
		}

		// Reset time since last tick.
		if (TraceState.TickPolicy == EGCS_TraceTickType::FixedFrameRate)
		{
			int32 SubTickCount = TraceState.SubTicks.Num();
			TraceState.TimeSinceLastTick -= SubTickCount * TraceState.TickInterval;

			TraceState.TimeSinceLastTick = FMath::Max(TraceState.TimeSinceLastTick, 0.0f);
		}
		else
		{
			TraceState.TimeSinceLastTick = 0;
		}

		TraceState.bShouldTickThisFrame = false;

		if (TraceState.ExecutionState == EGCS_TraceExecutionState::PendingStop)
		{
			TraceState.ExecutionState = EGCS_TraceExecutionState::Stopped;
			TraceState.TransformsOverTime.Empty();
			TraceState.TimeSinceActive = 0;
		}
		else
		{
			if (!TraceState.TransformsOverTime.IsEmpty())
			{
				TraceState.TransformsOverTime.RemoveAtSwap(0);
			}
		}
	});
}

void UGCS_TraceSubsystem::PrepareSubTicks(const float DeltaTime)
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::PrepareSubTicks"), STAT_UGCS_TraceSubsystem_PrepareSubTicks, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)

	ParallelFor(TraceStates.Num(), [&](const int32 Idx)
	{
		if (!TraceStates.IsValidIndex(Idx))
		{
			return;
		}
		auto& TraceState = TraceStates[Idx];

		TraceState.SubTicks.Reset();
		if (TraceState.bShouldTickThisFrame)
		{
			int32 SubTickCount = 1;
			if (TraceState.TickPolicy == EGCS_TraceTickType::DistanceBased)
			{
				SubTickCount = FMath::CeilToInt((TraceState.TransformsOverTime[0].GetLocation() - TraceState.TransformsOverTime[1].GetLocation()).Length() / TraceState.TickInterval);
			}
			else if (TraceState.TickPolicy == EGCS_TraceTickType::FixedFrameRate)
			{
				SubTickCount = FMath::FloorToInt(TraceState.TimeSinceLastTick / TraceState.TickInterval);
			}
			SubTickCount = FMath::Min(10, SubTickCount);

			if (TraceState.TransformsOverTime.Num() > 1)
			{
				const float SubTickRatio = 1.0 / SubTickCount;
				const FTransform CurrentTransform = TraceState.TransformsOverTime.Last();
				const FTransform PreviousTransform = TraceState.TransformsOverTime[0];
				TraceState.CollisionShapeOverTime = TraceState.Shape.Get<FGCS_CollisionShape>().GetDynamicCollisionShape(TraceState.SourceComponent, TraceState.TimeSinceActive);
				for (int32 i = 0; i < SubTickCount; i++)
				{
					FGCS_TraceSubTick SubTick{};
					SubTick.StartTransform = UKismetMathLibrary::TLerp(PreviousTransform, CurrentTransform, SubTickRatio * i, ELerpInterpolationMode::DualQuatInterp);
					SubTick.EndTransform = UKismetMathLibrary::TLerp(PreviousTransform, CurrentTransform, SubTickRatio * (i + 1), ELerpInterpolationMode::DualQuatInterp);
					SubTick.AverageTransform = UKismetMathLibrary::TLerp(SubTick.StartTransform, SubTick.EndTransform, 0.5, ELerpInterpolationMode::DualQuatInterp);
					TraceState.SubTicks.Add(SubTick);
				}
			}
			else
			{
				GCS_LOG(Warning, "Trace [%s] has less than 2 transforms in TransformsOverTime array!", *TraceState.Handle.TraceTag.ToString())
			}
		}

		if (TraceState.ExecutionState == EGCS_TraceExecutionState::InProgress)
		{
			TraceState.TimeSinceActive += DeltaTime;
		}
	});
}

void UGCS_TraceSubsystem::PerformSubTicks(const float DeltaTime)
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::PerformSubTicks"), STAT_UGCS_TraceSubsystem_PerformSubTicks, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)

	for (int32 Idx = 0; Idx < TraceStates.Num(); Idx++)
	{
		if (!TraceStates.IsValidIndex(Idx))
		{
			continue;
		}
		auto& TraceState = TraceStates[Idx];

		if (TraceState.bShouldTickThisFrame)
		{
			for (int32 SubTickIdx = 0; SubTickIdx < TraceState.SubTicks.Num(); SubTickIdx++)
			{
				const FGCS_TraceSubTick& SubTick = TraceState.SubTicks[SubTickIdx];

				// Respect cancellations by user-defined code immediately.
				if (TraceState.ExecutionState == EGCS_TraceExecutionState::Stopped)
				{
					break;
				}
				TraceState.TotalTickNumDuringExecution += 1;

				FTraceDelegate Delegate = FTraceDelegate::CreateUObject(this, &UGCS_TraceSubsystem::HandleTraceResults, Idx, TickIdx, TraceState.TimeSinceActive);

				PerformAsyncTrace(SubTick.StartTransform, SubTick.EndTransform, SubTick.AverageTransform, TraceState.World, TraceState.SweepSetting,
				                  TraceState.CollisionShapeOverTime,
				                  TraceState.CollisionParams,
				                  TraceState.ResponseParams,
				                  TraceState.ObjectQueryParams, &Delegate);
			}
		}
	}
}

void UGCS_TraceSubsystem::HandleTraceResults(const FTraceHandle& InTraceHandle, FTraceDatum& InTraceDatum, int32 InTraceStateIdx, uint32 InTickIdx, float InShapeTime)
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::HandleTraceResults"), STAT_UGCS_TraceSubsystem_HandleTraceResults, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)

	if (!TraceStates.IsValidIndex(InTraceStateIdx))
	{
		return;
	}

	auto& TraceState = TraceStates[InTraceStateIdx];

	if (!IsValid(TraceState.OwningSystem) || TraceState.ExecutionState == EGCS_TraceExecutionState::Stopped)
	{
		TraceState.TimeSinceActive = 0;
		return;
	}

	// GCS_LOG(Warning, "Trace [%s] has ticked %d within %f s", *TraceState.Handle.TraceTag.ToString(), TraceState.TotalTickNumDuringExecution, TraceState.TimeSinceActive)

#if ENABLE_DRAW_DEBUG
	if (GenericCombatSystemCVars::bEnableTraceDebugging)
	{
		const EDrawDebugTrace::Type DrawDebugType = GenericCombatSystemCVars::OverrideTraceDebuggingLifeTime > 0 ? EDrawDebugTrace::ForDuration : EDrawDebugTrace::ForOneFrame;
		const float DrawDebugTime = GenericCombatSystemCVars::OverrideTraceDebuggingLifeTime > 0
			                            ? GenericCombatSystemCVars::OverrideTraceDebuggingLifeTime
			                            : (DrawDebugType == EDrawDebugTrace::ForOneFrame ? 0.0f : 0.5f);

		const bool bHasAuthority = TraceState.SourceComponent->GetOwner()->HasAuthority();


		DrawDebug(InTraceDatum.Start,
		          InTraceDatum.End,
		          InTraceDatum.Rot,
		          InTraceDatum.OutHits, TraceState.Shape.Get<FGCS_CollisionShape>().GetDynamicCollisionShape(TraceState.SourceComponent, InShapeTime), TraceState.World,
		          DrawDebugType,
		          DrawDebugTime,
		          bHasAuthority ? FLinearColor::Red : FLinearColor::White,
		          bHasAuthority ? FLinearColor::Green : FLinearColor::Black);
		// GCS_OWNED_CLOG(TraceState.SourceComponent, Display, "Did collision trace: start(%s) end(%s)", *InTraceDatum.Start.ToCompactString(), *InTraceDatum.End.ToCompactString())
	}
#endif // ENABLE_DRAW_DEBUG

	TArray<FHitResult> FilteredHits;
	for (const FHitResult& NewHit : InTraceDatum.OutHits)
	{
		if (!TraceState.HitActors.Contains(NewHit.GetActor()))
		{
			FilteredHits.Add(NewHit);
			TraceState.HitActors.Add(NewHit.GetActor());
		}
	}
	if (FilteredHits.Num() > 0)
	{
		TraceState.OwningSystem->OnTraceHitDetected(TraceState.Handle, FilteredHits, TraceState.TimeSinceLastTick, InTickIdx);
	}
}

void UGCS_TraceSubsystem::PerformAsyncTrace(const FTransform& StartTransform, const FTransform& EndTransform, const FTransform& AverageTransform, UWorld* World,
                                            const FGCS_TraceSweepSetting& TraceSettings, const FCollisionShape& CollisionShape, const FCollisionQueryParams& CollisionParams,
                                            const FCollisionResponseParams& CollisionResponseParams,
                                            const FCollisionObjectQueryParams& ObjectQueryParams, const FTraceDelegate* InDelegate)
{
	switch (TraceSettings.SweepType)
	{
	case EGCS_TraceSweepType::ByChannel:
		World->AsyncSweepByChannel(
			EAsyncTraceType::Multi,
			StartTransform.GetLocation(),
			EndTransform.GetLocation(),
			AverageTransform.GetRotation(),
			TraceSettings.TraceChannel,
			CollisionShape,
			CollisionParams,
			CollisionResponseParams, InDelegate);
		break;
	case EGCS_TraceSweepType::ByObject:
		World->AsyncSweepByObjectType(
			EAsyncTraceType::Multi,
			StartTransform.GetLocation(),
			EndTransform.GetLocation(),
			AverageTransform.GetRotation(),
			ObjectQueryParams,
			CollisionShape,
			CollisionParams, InDelegate);
		break;
	case EGCS_TraceSweepType::ByProfile:
		World->AsyncSweepByProfile(
			EAsyncTraceType::Multi,
			StartTransform.GetLocation(),
			EndTransform.GetLocation(),
			AverageTransform.GetRotation(),
			TraceSettings.ProfileName,
			CollisionShape,
			CollisionParams, InDelegate);
		break;
	}
}

#if ENABLE_DRAW_DEBUG
void UGCS_TraceSubsystem::DrawDebug(const FVector& StartLocation, const FVector& EndLocation, const FQuat& Orientation, TArray<FHitResult> Hits, const FCollisionShape& CollisionShape,
                                    const UWorld* World,
                                    const EDrawDebugTrace::Type DrawDebugType, float DrawDebugTime, const FLinearColor& DrawDebugColor,
                                    const FLinearColor& DrawDebugHitColor)
{
	DECLARE_SCOPE_CYCLE_COUNTER(TEXT("UGCS_TraceSubsystem::DrawDebug"), STAT_UGCS_TraceSubsystem_DrawDebug, STATGROUP_GCS)
	TRACE_CPUPROFILER_EVENT_SCOPE_STR(__FUNCTION__)

	// We have to manually find if there is a blocking hit.
	bool bHasBlockingHit = false;
	for (const FHitResult& HitResult : Hits)
	{
		if (HitResult.bBlockingHit)
		{
			bHasBlockingHit = true;
			break;
		}
	}

	switch (CollisionShape.ShapeType)
	{
	case ECollisionShape::Sphere:
		DrawDebugSphereTraceMulti(World, StartLocation, EndLocation, CollisionShape.GetSphereRadius(), DrawDebugType, bHasBlockingHit, Hits, DrawDebugColor, DrawDebugHitColor, DrawDebugTime);
		break;
	case ECollisionShape::Capsule:
		DrawDebugCapsuleTraceMulti(World, StartLocation, EndLocation, CollisionShape.GetCapsuleRadius(), CollisionShape.GetCapsuleHalfHeight(), Orientation.Rotator(), DrawDebugType,
		                           bHasBlockingHit, Hits, DrawDebugColor, DrawDebugHitColor, DrawDebugTime);
		break;
	case ECollisionShape::Box:
		DrawDebugBoxTraceMulti(World, StartLocation, EndLocation, CollisionShape.GetBox(), Orientation.Rotator(), DrawDebugType, bHasBlockingHit, Hits, DrawDebugColor, DrawDebugHitColor,
		                       DrawDebugTime);
		break;
	default:
	case ECollisionShape::Line:
		DrawDebugLineTraceMulti(World, StartLocation, EndLocation, DrawDebugType, bHasBlockingHit, Hits, DrawDebugColor, DrawDebugHitColor, DrawDebugTime);
		break;
	}
}
#endif

#undef LOCTEXT_NAMESPACE