// Copyright 2023 PICO Inc. All Rights Reserved.

#pragma once

#include "CoreMinimal.h"
#include "IOpenXRExtensionPlugin.h"
#include "OpenXRHMD.h"
#include "PICO_HMDFunctionLibrary.h"
#include <ThirdParty/openxr/bytedance_composition_layer_eac.h>

class IXRTrackingSystem;

class FHMDPICO : public IOpenXRExtensionPlugin
{
public:
	FHMDPICO();
	virtual ~FHMDPICO() {}

	void Register();
	void Unregister();

	/** IOpenXRExtensionPlugin implementation */
	virtual FString GetDisplayName() override
	{
		return FString(TEXT("HMDPICO"));
	}

	virtual bool GetCustomLoader(PFN_xrGetInstanceProcAddr* OutGetProcAddr) override;
	virtual bool GetOptionalExtensions(TArray<const ANSICHAR*>& OutExtensions) override;
	virtual void PostCreateInstance(XrInstance InInstance) override;
	virtual void PostGetSystem(XrInstance InInstance, XrSystemId InSystem) override;
	virtual bool IsStandaloneStereoOnlyDevice() override;
	virtual const void* OnCreateSession(XrInstance InInstance, XrSystemId InSystem, const void* InNext) override;
	virtual void PostCreateSession(XrSession InSession) override;
	virtual void OnDestroySession(XrSession InSession) override;
	virtual bool UseCustomReferenceSpaceType(XrReferenceSpaceType& OutReferenceSpaceType);
	virtual bool GetSpectatorScreenController(FHeadMountedDisplayBase* InHMDBase, TUniquePtr<FDefaultSpectatorScreenController>& OutSpectatorScreenController) override;
	virtual void OnEvent(XrSession InSession, const XrEventDataBaseHeader* InHeader) override;
	virtual const void* OnEndProjectionLayer(XrSession InSession, int32 InLayerIndex, const void* InNext, XrCompositionLayerFlags& OutFlags) override;
	virtual void UpdateDeviceLocations(XrSession InSession, XrTime DisplayTime, XrSpace TrackingSpace) override;
	virtual void OnBeginRendering_RenderThread(XrSession InSession) override;
	virtual void UpdateCompositionLayers(XrSession InSession, TArray<const XrCompositionLayerBaseHeader*>& Headers);
	virtual void UpdateCompositionLayers(XrSession InSession, TArray<XrCompositionLayerBaseHeader*>& Headers);

	bool GetSupportedDisplayRefreshRates(TArray<float>& Rates);
	bool GetCurrentDisplayRefreshRate(float& Rate);
	bool SetDisplayRefreshRate(float Rate);

	void EnableContentProtect(bool Enable);

	bool SetPerformanceLevel(int domain, int level);
	bool GetDevicePoseForTime(const EControllerHand Hand, bool UseDefaultTime, FTimespan Timespan, bool& OutTimeWasUsed, FRotator& Orientation, FVector& Position, bool& bProvidedLinearVelocity, FVector& LinearVelocity, bool& bProvidedAngularVelocity, FVector& AngularVelocityRadPerSec, bool& bProvidedLinearAcceleration, FVector& LinearAcceleration, float InWorldToMetersScale);
	EHMDWornState::Type GetHMDWornState(bool& ResultValid);
	bool GetFieldOfView(float& OutHFOVInDegrees, float& OutVFOVInDegrees);
	bool GetInterpupillaryDistance(float& IPD);
	void SetBaseRotationAndBaseOffset(FRotator Rotation, FVector BaseOffset, EOrientPositionSelector::Type Options);
	void GetBaseRotationAndBaseOffset(FRotator& OutRotation, FVector& OutBaseOffset);
    FTimespan GetDisplayTime();

	void CopyTexture_RenderThread(FRHICommandListImmediate& RHICmdList, FRHITexture* DstTexture, FRHITexture* SrcTexture, FIntRect DstRect = FIntRect(), FIntRect SrcRect = FIntRect(), bool bAlphaPremultiply = false, bool bNoAlpha = false, bool bClearGreen = false, bool bInvertY = false, bool bInvertAlpha = false) const;

	bool IsEyeTrackerSupported(bool& Supported);
private:
	void*								 LoaderHandle;
	bool								 bPICORuntime;
	FRWLock								 SessionHandleMutex;
	XrInstance							 Instance;
	XrSystemId							 System;
	XrSession							 Session;
	XrPath								 CommonInteractionProfile;

	bool								 bSupportLocalFloorLevelEXT;

	bool								 bSupportDisplayRefreshRate;
	float								 CurrentDisplayRefreshRate;
	TArray<float>						 SupportedDisplayRefreshRates;
	PFN_xrEnumerateDisplayRefreshRatesFB xrEnumerateDisplayRefreshRatesFB = nullptr;
	PFN_xrGetDisplayRefreshRateFB		 xrGetDisplayRefreshRateFB = nullptr;
	PFN_xrRequestDisplayRefreshRateFB	 xrRequestDisplayRefreshRateFB = nullptr;

	bool								 bContentProtectEnabled;
	XrCompositionLayerSecureContentFB	 ContentProtect;
	bool								 bSupportPerformanceSettingsEXT;
	PFN_xrPerfSettingsSetPerformanceLevelEXT xrPerfSettingsSetPerformanceLevelEXT = nullptr;

	XrTime                               CurrentDisplayTime;
	class FOpenXRHMD*				     OpenXRHMD = nullptr;
	bool                                 IsSupportsUserPresence;
	EHMDWornState::Type     			 WornState;

	XrSpace								 CurrentBaseSpace = XR_NULL_HANDLE;

private:

	FIntRect GetViewportSize(const FOpenXRLayer::FPerEyeTextureData& EyeData, const IStereoLayers::FLayerDesc& Desc);
	FVector2D GetQuadSize(const FOpenXRLayer::FPerEyeTextureData& EyeData, const IStereoLayers::FLayerDesc& Desc);

	class UWorld*						 World = nullptr;
	void OnWorldCreated(UWorld* NewWorld) { World = NewWorld; }
	FDelegateHandle						 WorldLoadDelegate;

	/*************************** MRC Begin ***************************/
public:
	bool bSupportMRCExtension;
	void CreateMRCLayer(class UTexture* BackgroundRTTexture, class UTexture* ForegroundRTTexture);
	void DestroyMRCLayer();
	bool GetExternalCameraInfo(int32& width, int32& height, float& fov);
	bool GetExternalCameraPose(FTransform& Pose);
	void EnableMRCDebugMode(class UWorld* WorldContext, bool Enable, bool ViewInHMD, bool UseCustomTransform, const FTransform& Pose, bool UseCustomCameraInfo, int Width, int Height, float Fov);

private:
	XrSpace								 ViewTrackingSpace = XR_NULL_HANDLE;
	XrSpace								 MRCSpace = XR_NULL_HANDLE;
	bool								 bIsMRCRunning;
	UPROPERTY()
	class AMRCCameraPICO*		 RCSceneCapture2DPICO = nullptr;
	IStereoLayers::FLayerDesc			 MRCLayerDesc;
	TSharedPtr<FOpenXRLayer, ESPMode::ThreadSafe> MRCLayer;
	XrCompositionLayerQuad				 MRCQuadLayerLeft;
	XrCompositionLayerQuad				 MRCQuadLayerRight;

	struct FMRCDebugModePICO
	{
		FMRCDebugModePICO()
			:EnableExtension(false)
			, ViewInHMD(false)
			, UseCustomTransform(false)
			, Pose(FTransform::Identity)
			, UseCustomCameraInfo(false)
			, Width(256)
			, Height(256)
			, Fov(90.0f)
		{}

		bool EnableExtension;
		bool ViewInHMD;
		bool UseCustomTransform;
		FTransform Pose;
		bool UseCustomCameraInfo;
		int Width;
		int Height;
		float Fov;
	};

	FMRCDebugModePICO					 MRCDebugMode;

	/*************************** MRC End ***************************/

	bool								 bSupportedBDCompositionLayerSettingsExt;

	XrCompositionLayerAlphaBlendFB		 BlendState = { XR_TYPE_COMPOSITION_LAYER_ALPHA_BLEND_FB };

	bool								 bSupportAdaptiveResolution = false;
	
	union FXrCompositionLayerUnionPICO
	{
		XrCompositionLayerBaseHeader Header;
		XrCompositionLayerQuad Quad;
		XrCompositionLayerCylinderKHR Cylinder;
		XrCompositionLayerEquirectKHR Equirect;
		XrCompositionLayerEacBD EAC;
	};

	uint32								 StereoLayerID = 0;
	mutable FCriticalSection			 LayerCritSect;
	TMap<uint32, FOpenXRLayer>			 StereoLayersPICO;
	bool NeedReallocateTexture(bool bLeft, const FOpenXRLayer& Layer);
	void UpdateLayerSwapchainTexture(int64 WaitTimeout, const FOpenXRLayer& Layer, FRHICommandListImmediate& RHICmdList);
public:
	FIntPoint GetDefaultRenderTargetSize();
	void GetCurrentRenderTargetSize(uint32& InOutSizeX, uint32& InOutSizeY);

	uint32 CreateLayer(const IStereoLayers::FLayerDesc& Desc);
	void SetLayerDesc(uint32 ID, const IStereoLayers::FLayerDesc& Desc);
	void DestroyLayer(uint32 ID);
	void MarkTextureForUpdate(uint32 ID);
};