// Fill out your copyright notice in the Description page of Project Settings.

#include "RootMotionRecorderComponent.h"
#include "RecorderSession.h"
#include "Serialization/JsonWriter.h"
#include "Serialization/JsonSerializer.h"
#include "Misc/Paths.h"
#include "HAL/PlatformFileManager.h"
#include "Misc/FileHelper.h"
#include "DrawDebugHelpers.h"
#include "TimerManager.h"

DEFINE_LOG_CATEGORY(LogRootRecorder);

// Sets default values for this component's properties
URootMotionRecorderComponent::URootMotionRecorderComponent()
{
	// Set this component to be initialized when the game starts, and to be ticked every frame.  You can turn these features
	// off to improve performance if you don't need them.
	PrimaryComponentTick.bCanEverTick = false;
	OutputDirectory.Path = FPaths::ConvertRelativePathToFull(FPaths::ProjectDir() + TEXT("RecordedPaths/"));
}

// Called when the game starts
void URootMotionRecorderComponent::BeginPlay()
{
	Super::BeginPlay();

	// Create the directory if it doesn't already exist
	if (!FPaths::DirectoryExists(OutputDirectory.Path))
	{
		bool bCreated = FPlatformFileManager::Get().GetPlatformFile().CreateDirectoryTree(*OutputDirectory.Path);
		if (!bCreated)
		{
			UE_LOG(LogRootRecorder, Error, TEXT("Failed to create output directory: %s"), *OutputDirectory.Path);
		}
		else
		{
			UE_LOG(LogRootRecorder, Warning, TEXT("Output directory did not exist. Created: %s"), *OutputDirectory.Path);
		}
	}
}

void URootMotionRecorderComponent::StartRecording()
{
	if (GetOwner() && GetWorld())
	{
		// Verify frame rate
		if (ExportFrameRate <= 0)
		{
			UE_LOG(LogRootRecorder, Warning, TEXT("Invalid framerate of %d, falling back to 30"), ExportFrameRate);
			ExportFrameRate = 30;
		}

		SampleIndex = 0;
		RecordStartTime = GetWorld()->GetTimeSeconds();

		StartPosition = GetOwner()->GetActorLocation();
		LastDebugPosition = StartPosition + DebugDrawOffset;

		// Reserve space for 3 seconds of data
		RecordedData.Empty();
		RecordedRotations.Empty();
		RecordedPositionCurve = FVectorCurve();
		TimeIndexCurve = FFloatCurve();
		RecordedRotations.Reserve(ExportFrameRate * 180);

		// Start recording
		UE_LOG(LogRootRecorder, Log, TEXT("Recording started"));
		bIsRecording = true;
		RecordTick();
	}
}

void URootMotionRecorderComponent::EndRecording()
{
	EndRecordingWithName(FString());
}

void URootMotionRecorderComponent::EndRecordingWithName(const FString &BaseName)
{
	if (!bIsRecording)
	{
		return;
	}
	bIsRecording = false;

	// Process recorded data
	ProcessRecordedData();

	// Stop recording
	UE_LOG(LogRootRecorder, Log, TEXT("Recording ended"));

	// Verify output directory
	if (!FPaths::DirectoryExists(OutputDirectory.Path))
	{
		UE_LOG(LogRootRecorder, Warning, TEXT("Output directory does not exist, creating it: %s"), *OutputDirectory.Path);
		FPlatformFileManager::Get().GetPlatformFile().CreateDirectoryTree(*OutputDirectory.Path);
	}

	// Group this recording into the current session's subfolder. The subfolder is created
	// lazily here, on the first save of the session, so empty sessions leave no trace.
	FString TargetDirectory = OutputDirectory.Path;
	if (UWorld *World = GetWorld())
	{
		if (URecorderSession *Session = World->GetSubsystem<URecorderSession>())
		{
			TargetDirectory = Session->GetOrCreateSessionDirectory(OutputDirectory.Path);
		}
	}

	// Save recorded data to file with timestamp as filename. Session subfolder grouping
	// comes from TargetDirectory; filename stem correlates with the input schema when
	// EndRecordingWithName supplied one.
	const FString Timestamp = FDateTime::Now().ToString();
	const FString FileStem = BaseName.IsEmpty()
								 ? (TEXT("RecordedPath_") + Timestamp)
								 : (BaseName + TEXT("_recorded_") + Timestamp);
	FString FilePath = TargetDirectory + FileStem + TEXT(".json");
	SaveFloatArrayToJson(RecordedData, *FilePath);
	UE_LOG(LogRootRecorder, Log, TEXT("Recorded data saved to %s"), *FilePath);
}

void URootMotionRecorderComponent::SaveFloatArrayToJson(const TArray<float> &FloatArray, const FString &FilePath)
{
	// Step 1: Create a JSON array
	TArray<TSharedPtr<FJsonValue>> JsonValues;
	for (float Value : FloatArray)
	{
		JsonValues.Add(MakeShared<FJsonValueNumber>(Value));
	}

	// Step 2: Create the JSON writer
	FString OutputString;
	TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&OutputString);
	FJsonSerializer::Serialize(JsonValues, Writer);

	// Step 3: Save to file
	UE_LOG(LogRootRecorder, Warning, TEXT("Saving path to: %s"), *FilePath);
	FFileHelper::SaveStringToFile(OutputString, *FilePath);
}

void URootMotionRecorderComponent::RecordTick()
{
	if (!bIsRecording)
	{
		return;
	}
	if (GetOwner() && GetWorld())
	{
		// Store data in arrays
		RecordCurrentTime = GetWorld()->GetTimeSeconds() - RecordStartTime;
		TimeIndexCurve.UpdateOrAddKey(SampleIndex, RecordCurrentTime);
		RecordedPositionCurve.UpdateOrAddKey(GetOwner()->GetActorLocation(), RecordCurrentTime);
		RecordedRotations.Add(GetOwner()->GetActorRotation().Quaternion());

		if (bShouldDebugDraw)
		{
			FVector DebugPosition = FVector(
				GetOwner()->GetActorLocation().X,
				GetOwner()->GetActorLocation().Y,
				bFlattenTrajectory ? StartPosition.Z : GetOwner()->GetActorLocation().Z);
			DebugPosition += DebugDrawOffset;
			DrawDebugLine(GetWorld(), LastDebugPosition, DebugPosition, FColor::Red, false, 15.0f, 0, 1.0f);
			LastDebugPosition = DebugPosition;
		}

		SampleIndex++;

		GetWorld()->GetTimerManager().SetTimerForNextTick(this, &URootMotionRecorderComponent::RecordTick);
	}
}

void URootMotionRecorderComponent::ProcessRecordedData()
{
	// Reserve appropriate space for the recorded data
	const int32 numSamples = FMath::FloorToInt32(RecordCurrentTime * ExportFrameRate);
	RecordedData.SetNum(numSamples * 7);

	FVector worldPos;
	FQuat worldRot;
	FVector pos;
	FQuat rot;

	float deltaTime = 1.0f / ExportFrameRate;

	LastDebugPosition = StartPosition + DebugDrawOffset;

	// Process the recorded data
	for (int32 i = 0; i < numSamples; ++i)
	{
		// Get the time for this sample
		const float time = i * deltaTime;

		worldPos = RecordedPositionCurve.Evaluate(time, 1.0f);
		pos = worldPos - StartPosition;

		// This will give us a fractional 'index' to tell us which values in the Rotations array we need to lerp between
		float timeSample = TimeIndexCurve.Evaluate(time);

		if (timeSample <= 0.0f)
		{
			worldRot = RecordedRotations[0];
		}
		else if (timeSample >= RecordedRotations.Num() - 1)
		{
			worldRot = RecordedRotations.Last();
		}
		else
		{
			// Get the two closest rotations
			int32 i0 = FMath::FloorToInt32(timeSample);
			int32 i1 = i0 + 1;

			// Get the two rotations
			FQuat rot0 = RecordedRotations[i0];
			FQuat rot1 = RecordedRotations[i1];

			// Lerp between the two rotations
			worldRot = FQuat::Slerp(rot0, rot1, timeSample - i0);
		}

		// Axis transformation for rotation: Q' = R * Q * Ri
		rot = BasisChangeQuat * worldRot * BasisChangeQuat.Inverse();

		// Store the data in the array in [Px, Py, Pz, Qw, Qx, Qy, Qz] with Y-axis up
		RecordedData[i * 7] = (pos.Y) * -1.0;
		if (bFlattenTrajectory)
		{
			RecordedData[(i * 7) + 1] = 0.0f;
		}
		else
		{
			RecordedData[(i * 7) + 1] = pos.Z;
		}
		RecordedData[(i * 7) + 2] = pos.X;
		RecordedData[(i * 7) + 3] = rot.W;
		RecordedData[(i * 7) + 4] = rot.X;
		RecordedData[(i * 7) + 5] = rot.Y;
		RecordedData[(i * 7) + 6] = rot.Z;

		if (bShouldDebugDraw)
		{
			FVector DebugPosition = FVector(
				worldPos.X,
				worldPos.Y,
				bFlattenTrajectory ? StartPosition.Z : worldPos.Z);
			DebugPosition += DebugDrawOffset;
			DrawDebugSphere(GetWorld(), DebugPosition, 2.0f, 8, FColor::Purple, false, 15.0f, 0, 1.0f);
			DrawDebugLine(GetWorld(), LastDebugPosition, DebugPosition, FColor::Purple, false, 15.0f, 0, 1.0f);
			if (bDebugDrawRotations)
			{
				DrawDebugDirectionalArrow(GetWorld(), DebugPosition, DebugPosition + (worldRot.GetForwardVector() * 25.0f), 15.0f, FColor::Green, false, 15.0f, 0, 1.0f);
			}
			LastDebugPosition = DebugPosition;
		}
	}
}
