#pragma once
#include "CoreMinimal.h"
#include "Robot/PilotComponent/RobotPilotComponent.h"
#include "Subsystems/WorldSubsystem.h"
#include "Stats/Stats.h"
#include "EpisodeSubSystem.generated.h"


class ULuckyDataTransferSubsystem;
class ALuckySensorPawnBase;
class ATextRenderActor;
class AMujocoStaticMeshActor;
class ARobotPawn;

USTRUCT()
struct FTrainingEpisode
{
	GENERATED_BODY()
	UPROPERTY()
	int32 EpisodeIndex = -1;

	UPROPERTY()
	TArray<FString> Tasks;

	UPROPERTY()
	int32 Length = -1;
};




UCLASS()
class LUCKYWORLDV2_API UEpisodeSubSystem : public UWorldSubsystem
{
	GENERATED_BODY()

public:
	// Setup
	UEpisodeSubSystem();
	virtual void Initialize(FSubsystemCollectionBase& Collection) override;
	virtual void Deinitialize() override;

	// ----------------
	// ----- TICK -----
	// ----------------
	// TODO I don't like this solution, it's hacky - Tick should be in a component, a primitive or scene
	// TODO + it's leaking, not properly teared down
	// It will allows us to remove all the episode logic from the SubSystem and having different types of episodes 
	void Tick(float DeltaTime);
	void StartTicking();
	void StopTicking();
	FTSTicker::FDelegateHandle TickHandle;
	bool bTickEnabled = true;


	// --------------------
	// ------- DEBUG ------
	// --------------------
	UPROPERTY()
	ATextRenderActor* DebugTextActor = nullptr;
	int32 SuccessEpisodes = 0;
	int32 FailEpisodes = 0;
	void UpdateDebugTextActor() const;
	
	
	// ---------------------
	// ------- START -------
	// ---------------------
	/**
	 * Called by the UI when pressing the "Capture" button
	 */
	UFUNCTION(BlueprintCallable)
	void StartTraining(int32 EpisodesCountIn, FString BaseImageDataPathIn, FString TaskDescriptionIn);
	void EndTraining();

	UPROPERTY()
	ULuckyDataTransferSubsystem* DataTransfer = nullptr;
	

private:
	// ---------------------
	// ------- FLOW --------
	// ---------------------
	void StartEpisode();
	void EndEpisode();
	FTransform EpisodeRewardZone = FTransform::Identity;
	float EpisodeRewardZoneRadius = 5.f; // TODO Not hardcode it - or only in the Robot? - Maybe we want different scenarios for the robot
	bool CheckEpisodeCompletion();
	
	// Where the robot has to place the object
	FTransform EpisodeObjectBaseTransform = FTransform::Identity;

	// The object that will serve for the episode
	TObjectPtr<AMujocoStaticMeshActor> EpisodeTargetObject;
	
	// ---------------------
	// ------- ROBOT -------
	// ---------------------
	
	// The state of capture - if true we should call the scene capture and data transfer
	bool bIsCapturing = false;
	bool bIsEpisodeRunning = false;
	int32 EpisodesToCapture = 0;
	int32 CapturedEpisodes = 0;
	
	void FindEpisodeObjectFromScene();
	void FindRobotPawnFromScene();
	
	UPROPERTY()
	TObjectPtr<ARobotPawn> CurrentRobot;

	// ---------------------
	// ------ SENSORS ------
	// ---------------------
	void InitCameras();
	TArray<TObjectPtr<ALuckySensorPawnBase>> Cameras;

	// --------------------
	// ------- DATA -------
	// --------------------
	FString BaseImageDataPath;
	FString TaskDescription;
	int32 EpisodeFrames = 0;
	// Noah here add anything you need
	void ConfigureDataCapture();

	// End of episode tasks
	void CreateEpisodeStatJsonLine(const FTrainingEpisodeData& TrainingEpisodeData);
	void CreateEpisodeParquetFile();
	void ConvertImagesToVideo();
	
	// End of training files
	TArray<FString> EpisodeStatLines;
	void CreateEpisodesStatsJsonFile();
	void CreateEpisodesJsonFile();
	void CreateInfoJsonFile();
	void CreateTasksJsonFile();
};