#pragma once

#include "CoreMinimal.h"
#include "EngineUtils.h"

#include "Runtime/Engine/Public/SceneViewExtension.h"

#include "Runtime/Engine/Classes/Engine/Texture2D.h"

#include "MaterialShared.h"
#include "Materials/MaterialInstance.h"
#include "Delegates/Delegate.h"

#include "OptiXContext.h"
#include "OptiXObjectComponent.h"
#include "OptiXLaserComponent.h"
#include "OptiXCameraActor.h"



// Let's try some events!

DECLARE_EVENT(FOptiXContextManager, FLaserTraceFinishedEvent)
DECLARE_EVENT_OneParam(FOptiXContextManager, FWavelengthChangedEvent, const float)
DECLARE_MULTICAST_DELEGATE(FOnSceneChangedDelegate);


// DX

#if PLATFORM_WINDOWS
#include "AllowWindowsPlatformTypes.h"
#endif
#include <d3d11.h>
#if PLATFORM_WINDOWS
#include "HideWindowsPlatformTypes.h"
#endif

// print cuda error helper:

inline void PrintLastCudaError(FString Msg)
{
	cudaError_t Err = cudaGetLastError();
	if (cudaSuccess != Err) {
		UE_LOG(LogTemp, Fatal, TEXT("Cuda Error: %s. "), *Msg, static_cast<int>(Err), cudaGetErrorString(Err));
	}
}


class OPTIX_API FOptiXContextManager : public FSceneViewExtensionBase
{

public:

	// The auto register thing is used to make sure this constructor is only called via the NewExtension function
	FOptiXContextManager(const FAutoRegister& AutoRegister);

	~FOptiXContextManager() 
	{
		if (bIsInitialized)
		{
			cudaGraphicsUnregisterResource(CudaResourceDepthLeft);
			cudaGraphicsUnregisterResource(CudaResourceDepthRight);
			cudaGraphicsUnregisterResource(CudaResourceColorLeft);
			cudaGraphicsUnregisterResource(CudaResourceColorRight);
			PrintLastCudaError("cudaGraphicsUnregisterResource");
			cudaFree(CudaLinearMemoryDepth);
			cudaFree(CudaLinearMemoryColor);
			PrintLastCudaError("cudaFree");
		}

		AccelerationsToDeleteQueue.Empty();
		BuffersToDeleteQueue.Empty();
		GeometriesToDeleteQueue.Empty();
		GeometryGroupToDeleteQueue.Empty();
		GeometryInstancesToDeleteQueue.Empty();
		GroupsToDeleteQueue.Empty();
		MaterialsToDeleteQueue.Empty();
		ProgramToDeleteQueue.Empty();
		TextureSamplersToDeleteQueue.Empty();
		TransformsToDeleteQueue.Empty();
	}

	// ISceneViewExtension interface start, called by the render thread:
public:
	virtual void SetupViewFamily(FSceneViewFamily& InViewFamily) override;
	virtual void SetupView(FSceneViewFamily& InViewFamily, FSceneView& InView) override;
	virtual void BeginRenderViewFamily(FSceneViewFamily& InViewFamily) override;
	virtual void PreRenderView_RenderThread(FRHICommandListImmediate& RHICmdList, FSceneView& InView) override;
	virtual void PreRenderViewFamily_RenderThread(FRHICommandListImmediate& RHICmdList, FSceneViewFamily& InViewFamily) override;
	virtual void PostRenderViewFamily_RenderThread(FRHICommandListImmediate& RHICmdList, FSceneViewFamily& InViewFamily) override;
	virtual bool IsActiveThisFrame(class FViewport* InViewport) const override;
	virtual void PostRenderView_RenderThread(FRHICommandListImmediate& RHICmdList, FSceneView& InView) override;
	// ISceneViewExtension interface end


	// Initialization methods, called by the GAME thread

	void Init();

	void EndPlay()
	{
		CleanupOptiXOnEnd();
	}

	UOptiXContext* GetOptiXContext()
	{
		return OptiXContext.Get();
	}

	UMaterialInstanceDynamic* GetOptiXMID() // Used to set up the post process
	{
		return DynamicMaterial.Get();
	}

	void SceneChangedCallback();

	// The OptiX context is not thread-safe, so any changes to variables/properties on the game thread 
	// while a trace is running will lead to errors. Game thread objects push their requested updates into several queues,
	// which then get updated just before the trace runs on the render thread.
	// TQueue is guaranteed to be thread-safe.
	// There is probably a better and faster way of doing this.

	void RegisterOptiXComponent(IOptiXComponentInterface* Component)
	{
		ComponentsToInitializeQueue.Enqueue(Component);
	}

	void QueueComponentUpdate(IOptiXComponentInterface* Component)
	{
		ComponentsToUpdateQueue.Enqueue(Component);
	}

	void RequestCubemapUpdate(UOptiXCubemapComponent* Component)
	{
		CubemapComponentsToUpdateQueue.Enqueue(Component);
	}

	void SetActiveLaserComponent(UOptiXLaserComponent* Component)
	{
		LaserComponent = Component;
	}

	void SetActiveCameraActor(AOptiXPlayerCameraManager* Cam)
	{
		CameraActor = Cam;
	}

	int32 RequestCubemapId();

	void DeleteCubemapId(int32 Id);

	void AddCubemapToBuffer(int32 CubemapId, int32 SamplerId);

	void BroadcastWavelengthChange(float WL)
	{
		WavelengthChangedEvent.Broadcast(WL);
	}


public:
	
	FThreadSafeBool bStartTracing = false;
	FThreadSafeBool bIsInitialized = false;
	FThreadSafeBool bLaserIsInitialized = false;
	FThreadSafeBool bSceneChanged = true;
	FThreadSafeBool bIsTracing = false;
	FThreadSafeBool bClearToLaunch = true;
	FThreadSafeBool bCleanup = false;
	FThreadSafeBool bValidCubemap = false;

	FLaserTraceFinishedEvent LaserTraceFinishedEvent;
	FWavelengthChangedEvent WavelengthChangedEvent;
	FOnSceneChangedDelegate OnSceneChangedDelegate;


public:

	// This is so ugly but the optix api structure doesn't really allow anything more abstract.
	TQueue<optix::Acceleration>			AccelerationsToDeleteQueue;
	TQueue<optix::Buffer>				BuffersToDeleteQueue;
	TQueue<optix::Geometry>				GeometriesToDeleteQueue;
	TQueue<optix::GeometryGroup>		GeometryGroupToDeleteQueue;
	TQueue<optix::GeometryInstance>		GeometryInstancesToDeleteQueue;
	TQueue<optix::Group>				GroupsToDeleteQueue;
	TQueue<optix::Material>				MaterialsToDeleteQueue;
	TQueue<optix::Program>				ProgramToDeleteQueue;
	TQueue<optix::TextureSampler>		TextureSamplersToDeleteQueue;
	TQueue<optix::Transform>			TransformsToDeleteQueue;

	TQueue<TPair<optix::Group, uint32>> GroupChildrenToRemoveQueue;
	TQueue<TPair<optix::GeometryGroup, uint32>> GeometryGroupChildrenToRemoveQueue;


private:
	void InitContext();
	void InitRendering();
	void InitBuffers();
	void InitPrograms();
	void InitLaser();

	void LaunchLaser();

	void InitCubemap();

	void InitOptiXComponents(FRHICommandListImmediate & RHICmdList)
	{
		// Possibly dangerous? Use limited for instead of while in case something goes wrong and we deadlock
		for (uint32 i = 0; i < 100 && !ComponentsToInitializeQueue.IsEmpty(); i++)
		{
			IOptiXComponentInterface* Component;
			if (ComponentsToInitializeQueue.Dequeue(Component))
			{
				Component->InitOptiXComponent(RHICmdList);				
			}
		}
	}

	void UpdateOptiXComponentVariables()
	{
		for (uint32 i = 0; i < 100 && !ComponentsToUpdateQueue.IsEmpty(); i++)
		{
			IOptiXComponentInterface* Component;
			if (ComponentsToUpdateQueue.Dequeue(Component))
			{
				if (Component != nullptr)
				{
					Component->UpdateOptiXComponentVariables();
					Component->SetUpdateQueued(false);//   bUpdateQueued.AtomicSet(false);
				}
			}
		}
	}

	void RemovePendingChildrenFromGroups()
	{
		for (uint32 i = 0; i < 100 && !GroupChildrenToRemoveQueue.IsEmpty(); i++)
		{
			TPair<optix::Group, uint32> Pair;
			if (GroupChildrenToRemoveQueue.Dequeue(Pair))
			{
				if (Pair.Key != NULL)
				{
					Pair.Key->removeChild(Pair.Value); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to remove optix child in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !GeometryGroupChildrenToRemoveQueue.IsEmpty(); i++)
		{
			TPair<optix::GeometryGroup, uint32> Pair;
			if (GeometryGroupChildrenToRemoveQueue.Dequeue(Pair))
			{
				if (Pair.Key != NULL)
				{
					Pair.Key->removeChild(Pair.Value); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to remove optix child in queue: Object was NULL"));
				}
			}
		}
	}

	void DestroyOptiXObjects()
	{
		for (uint32 i = 0; i < 100 && !AccelerationsToDeleteQueue.IsEmpty(); i++)
		{
			optix::Acceleration NativeObj;
			if (AccelerationsToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					if(NativeObj->getContext() != NULL)
						NativeObj->destroy(); // TODO do we need to do anything else here?
					else
					{
						UE_LOG(LogTemp, Warning, TEXT("Context already destroyed but somehow this buffer handle is still valid"));
					}
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}			
		}
		for (uint32 i = 0; i < 100 && !BuffersToDeleteQueue.IsEmpty(); i++)
		{
			optix::Buffer NativeObj;
			if (BuffersToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					try
					{
						NativeObj->destroy(); // TODO do we need to do anything else here?
					}
					catch (optix::Exception& E)
					{
						FString Message = FString(E.getErrorString().c_str());
						UE_LOG(LogTemp, Error, TEXT("Trying to remove buffer: OptiX Error: %s"), *Message);
						GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::Red, FString::Printf(TEXT("OptiX Error %s"), *Message));
					}
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !GeometriesToDeleteQueue.IsEmpty(); i++)
		{
			optix::Geometry NativeObj;
			if (GeometriesToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					NativeObj->destroy(); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !GeometryInstancesToDeleteQueue.IsEmpty(); i++)
		{
			optix::GeometryInstance NativeObj;
			if (GeometryInstancesToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					NativeObj->destroy(); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !GeometryGroupToDeleteQueue.IsEmpty(); i++)
		{
			optix::GeometryGroup NativeObj;
			if (GeometryGroupToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					NativeObj->destroy(); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !GroupsToDeleteQueue.IsEmpty(); i++)
		{
			optix::Group NativeObj;
			if (GroupsToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					NativeObj->destroy(); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !MaterialsToDeleteQueue.IsEmpty(); i++)
		{
			optix::Material NativeObj;
			if (MaterialsToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					NativeObj->destroy(); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !ProgramToDeleteQueue.IsEmpty(); i++)
		{
			optix::Program NativeObj;
			if (ProgramToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					NativeObj->destroy(); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !TextureSamplersToDeleteQueue.IsEmpty(); i++)
		{
			optix::TextureSampler NativeObj;
			if (TextureSamplersToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					try
					{
						NativeObj->destroy(); // TODO do we need to do anything else here?
					}
					catch (optix::Exception& E)
					{
						FString Message = FString(E.getErrorString().c_str());
						UE_LOG(LogTemp, Error, TEXT("Trying to remove texture sampler: OptiX Error: %s"), *Message);
						GEngine->AddOnScreenDebugMessage(-1, 5.f, FColor::Red, FString::Printf(TEXT("OptiX Error %s"), *Message));
					}
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
		for (uint32 i = 0; i < 100 && !TransformsToDeleteQueue.IsEmpty(); i++)
		{
			optix::Transform NativeObj;
			if (TransformsToDeleteQueue.Dequeue(NativeObj))
			{
				if (NativeObj != NULL)
				{
					NativeObj->destroy(); // TODO do we need to do anything else here?
				}
				else
				{
					UE_LOG(LogTemp, Error, TEXT("Error while trying to destroy optix object in queue: Object was NULL"));
				}
			}
		}
	}

	void UpdateRequestedCubemaps(FRHICommandListImmediate & RHICmdList)
	{
		// update only the first for now, shouldn't be more than 1 in queue anyway:

		if (!CubemapComponentsToUpdateQueue.IsEmpty())
		{
			UOptiXCubemapComponent* Comp;
			if (CubemapComponentsToUpdateQueue.Dequeue(Comp))
			{
				Comp->UpdateCubemap(RHICmdList);
			}
		}
	}

	void UpdateCubemapBuffer(FRHICommandListImmediate & RHICmdList);

	void CleanupOptiXOnEnd()
	{

		UE_LOG(LogTemp, Display, TEXT("Starting Cleanup in Context Manager"));


		//check(IsInRenderingThread());
		bStartTracing.AtomicSet(false);
		bIsInitialized.AtomicSet(false);
		bLaserIsInitialized.AtomicSet(false);
		bSceneChanged.AtomicSet(true);
		bIsTracing.AtomicSet(false);
		bClearToLaunch.AtomicSet(true);
		bCleanup.AtomicSet(false);
		bValidCubemap.AtomicSet(false);

		// Clear the queue
		DestroyOptiXObjects();

		if (OptiXContext.IsValid())
		{
			OptiXContext->RemoveFromRoot();
		}
		if (LaserOutputBuffer.IsValid())
		{
			LaserOutputBuffer->RemoveFromRoot();
		}
		if (OutputTexture.IsValid())
		{
			OutputTexture->RemoveFromRoot();
		}
		if (OutputTexture2.IsValid())
		{
			OutputTexture2->RemoveFromRoot();
		}
		if (DepthTexture.IsValid())
		{
			DepthTexture->RemoveFromRoot();
		}
		if (DepthTexture2.IsValid())
		{
			DepthTexture2->RemoveFromRoot();
		}
		if (CubemapSampler.IsValid())
		{
			CubemapSampler->RemoveFromRoot();
		}
		if (CubemapBuffer.IsValid())
		{
			CubemapBuffer->RemoveFromRoot();
		}

		OutputBuffer.Reset();
		OutputDepthBuffer.Reset();

		OutputTexture.Reset();
		DepthTexture.Reset();
		OutputTexture2.Reset();
		DepthTexture2.Reset();

		DynamicMaterial.Reset();
		RegularMaterial.Reset();
		VRMaterial.Reset();


		TopObject.Reset();
		TopAcceleration.Reset();
		OptiXContext.Reset();

		if (NativeContext != NULL)
		{
			NativeContext->destroy();
			NativeContext = NULL;
		}

		/*{
			
			OptiXContext->RemoveFromRoot();
			LaserOutputBuffer->RemoveFromRoot();
			OutputTexture->RemoveFromRoot();
			OutputTexture2->RemoveFromRoot();
			DepthTexture->RemoveFromRoot();
			DepthTexture2->RemoveFromRoot();

		

			OutputBuffer.Reset();
			OutputDepthBuffer.Reset();

			OutputTexture.Reset();
			DepthTexture.Reset();
			OutputTexture2.Reset();
			DepthTexture2.Reset();

			DynamicMaterial.Reset();
			RegularMaterial.Reset();
			VRMaterial.Reset();


			TopObject.Reset();
			TopAcceleration.Reset();
			OptiXContext.Reset();

			if (NativeContext != NULL)
			{
				NativeContext->destroy();
			}
		}*/
	}

	
private:
	// OptiX Part

	// Todo: refactor this to delegates maybe?
	TQueue<IOptiXComponentInterface*> ComponentsToInitializeQueue;
	TQueue<IOptiXComponentInterface*> ComponentsToUpdateQueue;
	TQueue<UOptiXCubemapComponent*> CubemapComponentsToUpdateQueue;

	TWeakObjectPtr<UOptiXLaserComponent> LaserComponent;
	TWeakObjectPtr<AOptiXPlayerCameraManager> CameraActor;

	// OptiX Objects to be kept in the context manager, TODO triple check that the GC doesn't nab them.
	// Those are always required, but some should also be changeable! TODO make custom setters for them later.
	TWeakObjectPtr<UOptiXContext> OptiXContext;
	TWeakObjectPtr<UOptiXProgram> RayGenerationProgram;
	TWeakObjectPtr<UOptiXProgram> MissProgram;
	TWeakObjectPtr<UOptiXProgram> ExceptionProgram;

	TWeakObjectPtr<UOptiXGroup> TopObject;
	TWeakObjectPtr<UOptiXAcceleration> TopAcceleration;

	TWeakObjectPtr<UOptiXBuffer> OutputBuffer;
	TWeakObjectPtr<UOptiXBuffer> OutputDepthBuffer;

	TWeakObjectPtr<UOptiXTextureSampler> CubemapSampler;
	TWeakObjectPtr<UOptiXBuffer> CubemapBuffer;
	TWeakObjectPtr<UOptiXBuffer> CubemapsInputBuffer;

	TQueue<int32> UnallocatedCubemapIds;

	TArray<TArray<FColor>> SurfaceDataCube;

	optix::Context NativeContext;


	// Laser Part
	TWeakObjectPtr<UOptiXBuffer> LaserOutputBuffer;
	TWeakObjectPtr<UOptiXProgram> LaserRayGenerationProgram;
	TWeakObjectPtr<UOptiXProgram> LaserMissProgram;
	TWeakObjectPtr<UOptiXProgram> LaserExceptionProgram;



	// Rendering Part

	FUpdateTextureRegion2D TextureRegion;

	//UTextureCube* TextureCube;

	int32 Width;
	int32 Height;

	FTexture2DRHIRef OutputTextureColorRightRef;
	FTexture2DRHIRef OutputTextureColorLeftRef;
	FTexture2DRHIRef OutputTextureDepthRightRef;
	FTexture2DRHIRef OutputTextureDepthLeftRef;

	TWeakObjectPtr<UTexture2D> OutputTexture;
	TWeakObjectPtr<UTexture2D> DepthTexture;
	TWeakObjectPtr<UTexture2D> OutputTexture2;
	TWeakObjectPtr<UTexture2D> DepthTexture2;

	TWeakObjectPtr<UMaterialInstanceDynamic> DynamicMaterial;
	TWeakObjectPtr<UMaterial> RegularMaterial;
	TWeakObjectPtr<UMaterial> VRMaterial;
	bool bWithHMD;


	// ---------------------------------------------------------------------------------------------------
	// Laser stuff
	// ---------------------------------------------------------------------------------------------------

public:

	TQueue<TPair<uint32, TArray<FVector>>> LaserIntersectionQueue;

	TArray<TArray<FVector>> PreviousLaserResults;


private:

	TArray<FVector4> IntersectionData;
	TArray<FVector4> OldIntersectionData;

	FThreadSafeBool bTracingLaser;

	FCriticalSection CriticalSection;

	int32 LaserMaxDepth;
	int32 LaserEntryPoint;
	int32 LaserBufferSize;


	// ---------------------------------------------------------------------------------------------------
	// DX <-> CUDA stuff
	// ---------------------------------------------------------------------------------------------------


public:
	// I'm not sure which ones I actually need
	IDXGIAdapter			*CudaCapableAdapter = NULL;		// Adapter to use
	ID3D11Device			*D3DDevice = NULL;				// Rendering device
	ID3D11DeviceContext		*D3DDeviceContext = NULL;

	// RTX mode
	int RTXOn = 1;


private:

	void InitCUDADX();


	//ID3D11Texture2D* D3D11Texture;
	cudaGraphicsResource* CudaResourceDepthLeft;
	cudaGraphicsResource* CudaResourceDepthRight;
	cudaGraphicsResource* CudaResourceColorLeft;
	cudaGraphicsResource* CudaResourceColorRight;
	void* CudaLinearMemoryDepth;
	void* CudaLinearMemoryColor;
	size_t Pitch; // fix me

	cudaGraphicsResource *Resources[4];
};