#include "ixr.h"
#include "openxr/openxr.h"
#include "vulkan/vulkan_core.h"
#undef __cplusplus
#include "vulkan/vk_enum_string_helper.h"
#define __cplusplus 202400L
#define XR_USE_GRAPHICS_API_VULKAN
#include "openxr/openxr_platform.h"
#include "tier0/lib.h"
#include "tier1/utlvector.h"
#include "materialsystem/imaterialsystem.h"

struct EyeSwapchain_t
{
	XrSwapchain m_swapchain;
	uint32_t m_uWidth;
	uint32_t m_uHeight;
	uint32_t m_uImageIndex;
	CUtlVector<IImage*> m_pImages;
	IImage *m_pUserImage;
};

class COpenXRManager;
class COpenXRHeadset: public IXRHeadset, public IAppSystem
{
public:
	virtual void Init() override;
	virtual void Shutdown() override;

	virtual uint32_t GetSurfaceCount() override;
	virtual XRRenderSurface_t GetSurface( uint32_t i ) override;
	virtual void SetSurfaceImage( uint32_t i, IImage *pImage ) override;

	void Frame();
	void BeforeRender();
	void AfterRender();
	void CopySwapchain();
	void CreateSwapchainForEye( XrViewConfigurationView *pConf, VkFormat eVkFormat, EyeSwapchain_t *pEye );

	COpenXRManager *m_pManager;
	XrSystemId m_systemId;
	XrSession m_session;
	XrSpace m_space;

	CUtlVector<EyeSwapchain_t> m_eyeSwapchains;
	CUtlVector<XrCompositionLayerProjectionView> m_projviews;
	CUtlVector<XrView> m_views;
	float m_fPredictedTime;
	
};

class COpenXRController: public IXRController
{

};

class COpenXRManager: public IXRManager
{
public:
	virtual void Init() override;
	virtual void Shutdown() override;
	virtual void ConnectInterface( void *pIface, const char *szName ) override;
	virtual void Frame() override;
	virtual void PreRender() override;
	virtual void PostRender() override;
	virtual void CopySwapchain() override;

	virtual uint32_t GetHeadsetCount() override;
	virtual IXRHeadset *GetHeadset( uint32_t i ) override;

	bool bHasXR = false;
	IRenderContext *m_pRenderContext;
	COpenXRHeadset m_headset;

	XrInstance m_instance;
};
EXPOSE_INTERFACE(COpenXRManager, IXRManager, OPEN_XR_INTERFACE_VERSION);

void COpenXRManager::Init()
{
	PFN_xrInitializeLoaderKHR fnInitLoader = NULL;
	XrResult r = xrGetInstanceProcAddr(NULL, "xrInitializeLoaderKHR", (PFN_xrVoidFunction*)&fnInitLoader);
	if (r != XR_SUCCESS)
	{
		V_printf("Failed to load OpenXR, we are not running OpenXR\n");
		return;
	}
	XrApplicationInfo appInfo = {};
	appInfo.apiVersion = XR_API_VERSION_1_1;
	V_strncpy(appInfo.applicationName, "funny", XR_MAX_APPLICATION_NAME_SIZE);
	V_strncpy(appInfo.engineName, "funny", XR_MAX_ENGINE_NAME_SIZE);
	if (r != XR_SUCCESS)
		Plat_FatalErrorFunc("Failed to load OpenXR\n");
	const char *extensions[] = {
		XR_EXT_DEBUG_UTILS_EXTENSION_NAME,
		XR_KHR_VULKAN_ENABLE_EXTENSION_NAME,
	};

	const char *layers[] = {
		"XR_APILAYER_LUNARG_core_validation"
	};


	XrInstanceCreateInfo createInfo = {};
	createInfo.type = XR_TYPE_INSTANCE_CREATE_INFO;
	createInfo.applicationInfo = appInfo;
	createInfo.enabledExtensionCount = sizeof(extensions)/sizeof(*extensions);
	createInfo.enabledExtensionNames = extensions;
	createInfo.enabledApiLayerCount = sizeof(layers)/sizeof(*layers);
	createInfo.enabledApiLayerNames = layers;
	r = xrCreateInstance(&createInfo, &m_instance);
	if (r != XR_SUCCESS)
	{
		V_printf("xrCreateInstance failed, we are not running OpenXR\n");
		return;
	}
	
	m_headset.m_pManager = this;
	m_headset.Init();

	XrInstanceProperties props = {};
	props.type = XR_TYPE_INSTANCE_PROPERTIES;
	xrGetInstanceProperties(m_instance, &props);
	V_printf("OpenXR runtime: %s\n", props.runtimeName);

	bHasXR = true;

}

void COpenXRManager::Shutdown()
{

}
void COpenXRManager::Frame()
{
	if (bHasXR)
		m_headset.Frame();
}

void COpenXRManager::PreRender()
{
	if (bHasXR)
		m_headset.BeforeRender();
}

void COpenXRManager::PostRender()
{
	if (bHasXR)
		m_headset.AfterRender();
}
void COpenXRManager::CopySwapchain()
{
	if (bHasXR)
		m_headset.CopySwapchain();

}


void COpenXRHeadset::Init()
{

	XrSystemGetInfo sgi = {};
	XrResult r; 
	sgi.type = XR_TYPE_SYSTEM_GET_INFO;
	sgi.formFactor = XR_FORM_FACTOR_HEAD_MOUNTED_DISPLAY;
	r = xrGetSystem(m_pManager->m_instance, &sgi, &m_systemId);
	if (r != XR_SUCCESS)
		Plat_FatalErrorFunc("xrGetSystem\n");


	XrGraphicsRequirementsVulkanKHR requirements = {};
	requirements.type = XR_TYPE_GRAPHICS_REQUIREMENTS_VULKAN_KHR;
	PFN_xrGetVulkanGraphicsRequirementsKHR fnVulkanRequirements = NULL;
	xrGetInstanceProcAddr(m_pManager->m_instance, "xrGetVulkanGraphicsRequirementsKHR", (PFN_xrVoidFunction*)&fnVulkanRequirements);
	if (r != XR_SUCCESS)
		Plat_FatalErrorFunc("Failed to get xrGetVulkanGraphicsRequirementsKHR\n");
	fnVulkanRequirements( m_pManager->m_instance, m_systemId, &requirements);


	PFN_xrGetVulkanGraphicsDeviceKHR fnVulkanDevice = NULL;
	xrGetInstanceProcAddr(m_pManager->m_instance, "xrGetVulkanGraphicsDeviceKHR", (PFN_xrVoidFunction*)&fnVulkanDevice);
	if (r != XR_SUCCESS)
		Plat_FatalErrorFunc("Failed to get xrGetVulkanGraphicsDeviceKHR\n");
	VkPhysicalDevice d = NULL;
	fnVulkanDevice(m_pManager->m_instance, m_systemId, (VkInstance)m_pManager->m_pRenderContext->GetVulkanInstance(), &d);

	XrGraphicsBindingVulkan2KHR vkBinding = {};
	vkBinding.type = XR_TYPE_GRAPHICS_BINDING_VULKAN2_KHR;
	vkBinding.instance = (VkInstance)m_pManager->m_pRenderContext->GetVulkanInstance();
	vkBinding.physicalDevice = (VkPhysicalDevice)m_pManager->m_pRenderContext->GetVulkanPhysicalDevice();
	vkBinding.device = (VkDevice)m_pManager->m_pRenderContext->GetVulkanDevice();

	XrSessionCreateInfo sessionCreateInfo = {};
	sessionCreateInfo.type = XR_TYPE_SESSION_CREATE_INFO;
	sessionCreateInfo.next = &vkBinding;
	sessionCreateInfo.systemId = m_systemId;
	r = xrCreateSession(m_pManager->m_instance, &sessionCreateInfo, &m_session);
	if (r != XR_SUCCESS)
		Plat_FatalErrorFunc("xrCreateSession\n");

	/* create swapchains for rendering pipeline */

	uint32_t formatCount = 0;
	xrEnumerateSwapchainFormats(m_session, 0, &formatCount, NULL);

	CUtlVector<int64_t> formats = formatCount;
	xrEnumerateSwapchainFormats(m_session, formatCount, &formatCount, formats.GetData());

	const VkFormat preferedSurfaceFormats[] = {
		VK_FORMAT_B8G8R8A8_UNORM,
		VK_FORMAT_R8G8B8A8_UNORM,
		VK_FORMAT_R8G8B8A8_SRGB,
	};

	VkFormat eSelectedFormat = preferedSurfaceFormats[0];
	for (auto &format: formats)
	{
		V_printf("Format: %s\n", string_VkFormat((VkFormat)format));
		for (int i = 0; i < sizeof(preferedSurfaceFormats)/sizeof(VkFormat); i++)
		{
			if (format == preferedSurfaceFormats[i])
			{
				eSelectedFormat = preferedSurfaceFormats[i];
				goto formatPicked;
			}
		}
	}
formatPicked:

	uint32_t viewCount = 0;
	xrEnumerateViewConfigurationViews(m_pManager->m_instance, m_systemId, XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO, 0, &viewCount, NULL);
	CUtlVector<XrViewConfigurationView> views = viewCount;
	for ( auto &v: views )
	{
		v.type = XR_TYPE_VIEW_CONFIGURATION_VIEW;
	}
	xrEnumerateViewConfigurationViews(m_pManager->m_instance, m_systemId, XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO, viewCount, &viewCount, views.GetData());

	for ( auto &v: views )
	{
		m_eyeSwapchains.AppendTail((EyeSwapchain_t){});
		CreateSwapchainForEye(&v, eSelectedFormat, m_eyeSwapchains.GetData() + m_eyeSwapchains.GetSize() - 1);
	}

	XrReferenceSpaceCreateInfo spaceInfo = {};
	spaceInfo.type = XR_TYPE_REFERENCE_SPACE_CREATE_INFO;
	spaceInfo.referenceSpaceType = XR_REFERENCE_SPACE_TYPE_LOCAL;
	spaceInfo.poseInReferenceSpace = {{0.0f, 0.0f, 0.0f, 1.0f}, {0.0f, 0.0f, 0.0f}};
	xrCreateReferenceSpace(m_session, &spaceInfo, &m_space);
	m_fPredictedTime = 0;
	m_views = m_eyeSwapchains.GetSize();
	m_projviews = m_eyeSwapchains.GetSize();
	for ( auto &v: m_views )
	{
		v = {};
		v.type = XR_TYPE_VIEW;
	}
	for ( auto &v: m_projviews )
	{
		v = {};
		v.type = XR_TYPE_COMPOSITION_LAYER_PROJECTION_VIEW;
	}

	XrViewState viewState = {};
	viewState.type = XR_TYPE_VIEW_STATE;
	XrViewLocateInfo viewLocale = {};
	viewLocale.type = XR_TYPE_VIEW_LOCATE_INFO;
	viewLocale.displayTime = m_fPredictedTime;
	viewLocale.space = m_space;
	viewLocale.viewConfigurationType = XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO;
	uint32_t vc = m_views.GetSize();
	xrLocateViews(m_session, &viewLocale, &viewState, vc, &vc, m_views.GetData());

}
void COpenXRHeadset::Shutdown()
{

}

void COpenXRHeadset::Frame()
{
	while (true)
	{
		XrEventDataBuffer event = {};
		event.type = XR_TYPE_EVENT_DATA_BUFFER;
		XrResult r = xrPollEvent(m_pManager->m_instance, &event);
		if ( r != XR_SUCCESS )
			break;

		switch (event.type)
		{
		case XR_TYPE_EVENT_DATA_SESSION_STATE_CHANGED:
		{
			XrEventDataSessionStateChanged *e = (XrEventDataSessionStateChanged*)&event;
			if ( e->state == XR_SESSION_STATE_READY )
			{
				XrSessionBeginInfo beginInfo = {};
				beginInfo.type = XR_TYPE_SESSION_BEGIN_INFO;
				beginInfo.primaryViewConfigurationType = XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO;
				xrBeginSession(m_session, &beginInfo);	
			}
			if ( e->state == XR_SESSION_STATE_STOPPING )
			{
				xrEndSession(m_session);
			}
			break;
		}
		default:
			break;
		}

	}

	XrFrameWaitInfo frameWait = {};
	XrFrameState frameState = {};
	frameState.type = XR_TYPE_FRAME_STATE;
	frameWait.type = XR_TYPE_FRAME_WAIT_INFO;
	xrWaitFrame(m_session, &frameWait, &frameState);
	m_fPredictedTime = frameState.predictedDisplayTime;


	m_views = m_eyeSwapchains.GetSize();
	m_projviews = m_eyeSwapchains.GetSize();
	for ( auto &v: m_views )
	{
		v = {};
		v.type = XR_TYPE_VIEW;
	}
	for ( auto &v: m_projviews )
	{
		v = {};
		v.type = XR_TYPE_COMPOSITION_LAYER_PROJECTION_VIEW;
	}

	XrViewState viewState = {};
	viewState.type = XR_TYPE_VIEW_STATE;
	XrViewLocateInfo viewLocale = {};
	viewLocale.type = XR_TYPE_VIEW_LOCATE_INFO;
	viewLocale.displayTime = m_fPredictedTime;
	viewLocale.space = m_space;
	viewLocale.viewConfigurationType = XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO;
	uint32_t viewCount = m_views.GetSize();
	xrLocateViews(m_session, &viewLocale, &viewState, viewCount, &viewCount, m_views.GetData());
}

void COpenXRHeadset::BeforeRender()
{

	XrFrameBeginInfo begin = {};
	begin.type = XR_TYPE_FRAME_BEGIN_INFO;
	xrBeginFrame(m_session, &begin);
	for ( uint32_t i = 0; i < m_views.GetSize(); i++ )
	{
		XrSwapchainImageAcquireInfo acq = {};
		acq.type = XR_TYPE_SWAPCHAIN_IMAGE_ACQUIRE_INFO;
		uint32_t idx = 0;
		XrResult r = xrAcquireSwapchainImage(m_eyeSwapchains[i].m_swapchain, &acq, &idx);
		m_eyeSwapchains[i].m_uImageIndex = idx;
		if (r) return;
		XrSwapchainImageWaitInfo wait = {};
		wait.type = XR_TYPE_SWAPCHAIN_IMAGE_WAIT_INFO;
		wait.timeout = XR_INFINITE_DURATION;
		xrWaitSwapchainImage(m_eyeSwapchains[i].m_swapchain, &wait);

		m_projviews[i].fov = m_views[i].fov;
		m_projviews[i].pose = m_views[i].pose;
		m_projviews[i].subImage.swapchain = m_eyeSwapchains[i].m_swapchain;
		m_projviews[i].subImage.imageRect.extent.height = m_eyeSwapchains[i].m_pImages[idx]->GetImageWidth();
		m_projviews[i].subImage.imageRect.extent.width = m_eyeSwapchains[i].m_pImages[idx]->GetImageHeight();
		m_projviews[i].subImage.imageArrayIndex = 0;
	}

}
void COpenXRHeadset::CopySwapchain()
{
	IRenderContext *pCtx = m_pManager->m_pRenderContext;
	IRenderCommandList *pList = pCtx->CreateCommandList();
	pList->StartRecording();
	for ( auto &eye: m_eyeSwapchains)
	{
		ImageSector_t sector = {};
		sector.m_iWidth = eye.m_uWidth;
		sector.m_iHeight = eye.m_uHeight;
		pList->BlitImageToImage(eye.m_pUserImage, sector, eye.m_pImages[eye.m_uImageIndex], sector);
	}
	pList->EndRecording();
	pCtx->SubmitCommandList(pList);
	pCtx->DestroyCommandList(pList);
}

void COpenXRHeadset::AfterRender()
{
	for ( uint32_t i = 0; i < m_views.GetSize(); i++ )
	{
		XrSwapchainImageReleaseInfo rel = {};
		rel.type = XR_TYPE_SWAPCHAIN_IMAGE_RELEASE_INFO;
		xrReleaseSwapchainImage(m_eyeSwapchains[i].m_swapchain, &rel);
	}

	XrCompositionLayerProjection proj = {};
	proj.type = XR_TYPE_COMPOSITION_LAYER_PROJECTION;
	proj.space = m_space;
	proj.viewCount = m_projviews.GetSize();
	proj.views = m_projviews.GetData();
	XrCompositionLayerBaseHeader *compositions[] = {
		(XrCompositionLayerBaseHeader*)&proj
	};

	XrFrameEndInfo end = {};
	end.type = XR_TYPE_FRAME_END_INFO;
	end.environmentBlendMode = XR_ENVIRONMENT_BLEND_MODE_OPAQUE;
	end.layerCount = 1;
	end.layers = compositions;
	end.displayTime = m_fPredictedTime;
	xrEndFrame(m_session, &end);
}

uint32_t COpenXRHeadset::GetSurfaceCount()
{
	V_printf("%i\n", m_eyeSwapchains.GetSize());
	return m_eyeSwapchains.GetSize();
}

XRRenderSurface_t COpenXRHeadset::GetSurface( uint32_t i )
{
	XRRenderSurface_t surface = {};
	surface.m_fFov = m_projviews[i].fov.angleLeft+m_projviews[i].fov.angleRight;
	surface.m_uWidth = m_eyeSwapchains[i].m_uWidth;
	surface.m_uHeight = m_eyeSwapchains[i].m_uHeight;
	return surface;
}

void COpenXRHeadset::SetSurfaceImage( uint32_t i, IImage *pImage )
{
	m_eyeSwapchains[i].m_pUserImage = pImage;

}

void COpenXRHeadset::CreateSwapchainForEye( XrViewConfigurationView *pConf, VkFormat eVkFormat, EyeSwapchain_t *pEye )
{
	XrSwapchainCreateInfo createInfo = {};
	createInfo.type = XR_TYPE_SWAPCHAIN_CREATE_INFO;
	createInfo.usageFlags = XR_SWAPCHAIN_USAGE_COLOR_ATTACHMENT_BIT | XR_SWAPCHAIN_USAGE_SAMPLED_BIT | XR_SWAPCHAIN_USAGE_TRANSFER_DST_BIT;
	createInfo.format = eVkFormat;
	createInfo.sampleCount = pConf->recommendedSwapchainSampleCount;
	createInfo.width = pConf->recommendedImageRectWidth;
	createInfo.height = pConf->recommendedImageRectHeight;
	createInfo.faceCount = 1;
	createInfo.arraySize = 1;
	createInfo.mipCount = 1;
	XrResult r = xrCreateSwapchain(m_session, &createInfo, &pEye->m_swapchain);
	if (r != XR_SUCCESS)
		Plat_FatalErrorFunc("xrCreateSwapchain\n");
	pEye->m_uWidth = pConf->recommendedImageRectWidth;
	pEye->m_uHeight = pConf->recommendedImageRectHeight;





	uint32_t imageCount = 0;
	xrEnumerateSwapchainImages(pEye->m_swapchain, 0, &imageCount, NULL);
	CUtlVector<XrSwapchainImageVulkan2KHR> images = imageCount;
	for ( auto &i: images )
	{
		i.type = XR_TYPE_SWAPCHAIN_IMAGE_VULKAN2_KHR;
	}
	xrEnumerateSwapchainImages(pEye->m_swapchain, imageCount, &imageCount, (XrSwapchainImageBaseHeader*)images.GetData());
	pEye->m_pImages = imageCount;

	EImageFormat eFormat = IMAGE_FORMAT_RGBA8_UNORM;
	switch(eVkFormat)
	{
	case VK_FORMAT_R8G8B8A8_UNORM: eFormat = IMAGE_FORMAT_RGBA8_UNORM; break;
	case VK_FORMAT_B8G8R8A8_UNORM: eFormat = IMAGE_FORMAT_BGRA8_UNORM; break;
	case VK_FORMAT_R8G8B8A8_SRGB: eFormat = IMAGE_FORMAT_RGBA8_SRGB; break;
	default:
		break;

	}
	for ( int i = 0; i < imageCount; i++ )
	{
		pEye->m_pImages[i] = m_pManager->m_pRenderContext->CreateImageFromVkImage(
				images[i].image, 
				createInfo.width,
				createInfo.height,
				VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
				eFormat,
				MULTISAMPLE_TYPE_1_SAMPLES);
	}


}

uint32_t COpenXRManager::GetHeadsetCount()
{
	return 1;
}

IXRHeadset *COpenXRManager::GetHeadset( uint32_t i )
{
	if (i)
		return 0;
	if(bHasXR)
		return 0;

	return &m_headset;
}

void COpenXRManager::ConnectInterface( void *pIface, const char *szName )
{
	if (!V_strcmp(szName, RENDER_CONTEXT_INTERFACE_VERSION))
		m_pRenderContext = (IRenderContext*)pIface;
}