#include "openxr_gl.h"

extern "C" {
#include <wayland-server-core.h>
#include <wlr/backend/interface.h>
#include <wlr/interfaces/wlr_output.h>
#include <wlr/util/log.h>
}

#include <cstring>
#include <memory>

#define GL_GLEXT_PROTOTYPES
#include <GL/gl.h>
#include <GL/glext.h>
#include <GL/glx.h>
#include <X11/Xlib.h>
#include <openxr/openxr.h>
#include <openxr/openxr_platform.h>
#include <vector>
// Internal backend definition
struct openxr_backend {
	wlr_backend base;
	wl_display* display;
	wl_event_loop* event_loop;
	XrInstance instance {};
	XrSession session {};
	XrSpace app_space { XR_NULL_HANDLE };
	XrSwapchain swapchain { XR_NULL_HANDLE };
	int32_t width = 0, height = 0;
	std::vector<XrSwapchainImageOpenGLKHR> swapchain_images;
	std::vector<GLuint> framebuffers;
	bool started = false;
	struct wlr_output* output = nullptr;
};

static void output_destroy(struct wlr_output* wlr_output) { (void)wlr_output; }

static bool output_test(
    struct wlr_output* wlr_output, const struct wlr_output_state* state)
{
	(void)wlr_output;
	(void)state;
	return true;
}

static bool output_commit(
    struct wlr_output* wlr_output, const struct wlr_output_state* state)
{
	// Retrieve our backend
	auto* xr = reinterpret_cast<openxr_backend*>(wlr_output->backend);
	(void)state;
	// Wait for frame
	XrFrameState frameState { XR_TYPE_FRAME_STATE };
	xrWaitFrame(xr->session, nullptr, &frameState);
	xrBeginFrame(xr->session, nullptr);
	// Locate views
	uint32_t viewCount = static_cast<uint32_t>(xr->framebuffers.size());
	std::vector<XrView> views(viewCount, { XR_TYPE_VIEW });
	XrViewLocateInfo viewLocInfo { XR_TYPE_VIEW_LOCATE_INFO };
	viewLocInfo.viewConfigurationType
	    = XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO;
	viewLocInfo.displayTime = frameState.predictedDisplayTime;
	viewLocInfo.space = xr->app_space;
	XrViewState viewState { XR_TYPE_VIEW_STATE };
	uint32_t viewCountOutput;
	xrLocateViews(xr->session, &viewLocInfo, &viewState, viewCount,
	    &viewCountOutput, views.data());
	// Prepare projection views
	std::vector<XrCompositionLayerProjectionView> projViews(
	    viewCountOutput, { XR_TYPE_COMPOSITION_LAYER_PROJECTION_VIEW });
	for (uint32_t i = 0; i < viewCountOutput; ++i) {
		// Acquire swapchain image
		XrSwapchainImageAcquireInfo acqInfo {
			XR_TYPE_SWAPCHAIN_IMAGE_ACQUIRE_INFO
		};
		uint32_t imgIndex;
		xrAcquireSwapchainImage(xr->swapchain, &acqInfo, &imgIndex);
		XrSwapchainImageWaitInfo waitInfo { XR_TYPE_SWAPCHAIN_IMAGE_WAIT_INFO };
		waitInfo.timeout = XR_INFINITE_DURATION;
		xrWaitSwapchainImage(xr->swapchain, &waitInfo);
		// Bind framebuffer and clear
		glBindFramebuffer(GL_FRAMEBUFFER, xr->framebuffers[imgIndex]);
		glViewport(0, 0, xr->width, xr->height);
		glClearColor(0.1f, 0.1f, 0.1f, 1.0f);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
		// TODO: render scene (e.g., wlr_scene_render)
		// Release swapchain image
		{
			XrSwapchainImageReleaseInfo relInfo {
				XR_TYPE_SWAPCHAIN_IMAGE_RELEASE_INFO
			};
			xrReleaseSwapchainImage(xr->swapchain, &relInfo);
		}
		// Setup projection view
		projViews[i].pose = views[i].pose;
		projViews[i].fov = views[i].fov;
		projViews[i].subImage.swapchain = xr->swapchain;
		projViews[i].subImage.imageArrayIndex = imgIndex;
		// Set image rectangle offset and extent
		projViews[i].subImage.imageRect.offset.x = 0;
		projViews[i].subImage.imageRect.offset.y = 0;
		projViews[i].subImage.imageRect.extent.width = xr->width;
		projViews[i].subImage.imageRect.extent.height = xr->height;
	}
	// Unbind
	glBindFramebuffer(GL_FRAMEBUFFER, 0);
	// End frame
	XrCompositionLayerProjection layer { XR_TYPE_COMPOSITION_LAYER_PROJECTION };
	layer.space = xr->app_space;
	layer.viewCount = static_cast<uint32_t>(projViews.size());
	layer.views = projViews.data();
	XrCompositionLayerBaseHeader const* layers[]
	    = { reinterpret_cast<XrCompositionLayerBaseHeader const*>(&layer) };
	XrFrameEndInfo frameEndInfo { XR_TYPE_FRAME_END_INFO };
	frameEndInfo.displayTime = frameState.predictedDisplayTime;
	frameEndInfo.environmentBlendMode = XR_ENVIRONMENT_BLEND_MODE_OPAQUE;
	frameEndInfo.layerCount = 1;
	frameEndInfo.layers = layers;
	xrEndFrame(xr->session, &frameEndInfo);
	return true;
}

static const struct wlr_drm_format_set* output_get_primary_formats(
    struct wlr_output* wlr_output, uint32_t buffer_caps)
{
	(void)wlr_output;
	(void)buffer_caps;
	return NULL;
}

static const struct wlr_output_impl output_impl = {
	.destroy = output_destroy,
	.test = output_test,
	.commit = output_commit,
	.get_primary_formats = output_get_primary_formats,
};

static bool backend_start(wlr_backend* backend)
{
	auto* xr = reinterpret_cast<openxr_backend*>(backend);
	if (xr->started)
		return true;

	XrApplicationInfo ai {};
	std::strncpy(
	    ai.applicationName, "LunarWM", XR_MAX_APPLICATION_NAME_SIZE - 1);
	ai.applicationVersion = 1;
	std::strncpy(ai.engineName, "LunarWM", XR_MAX_ENGINE_NAME_SIZE - 1);
	ai.engineVersion = 1;
	ai.apiVersion = XR_CURRENT_API_VERSION;

	char const* exts[] = {
		XR_EXT_DEBUG_UTILS_EXTENSION_NAME,
		XR_KHR_OPENGL_ENABLE_EXTENSION_NAME,
	};

	XrInstanceCreateInfo ic { XR_TYPE_INSTANCE_CREATE_INFO };
	ic.applicationInfo = ai;
	ic.enabledExtensionCount = sizeof(exts) / sizeof(exts[0]);
	ic.enabledExtensionNames = exts;

	if (xrCreateInstance(&ic, &xr->instance) != XR_SUCCESS) {
		wlr_log(WLR_ERROR, "Failed to create OpenXR instance");
		return false;
	}

	XrSystemGetInfo sgi { XR_TYPE_SYSTEM_GET_INFO };
	sgi.formFactor = XR_FORM_FACTOR_HEAD_MOUNTED_DISPLAY;
	XrSystemId system_id;
	if (xrGetSystem(xr->instance, &sgi, &system_id) != XR_SUCCESS) {
		wlr_log(WLR_ERROR, "xrGetSystem failed");
		return false;
	}

	PFN_xrGetOpenGLGraphicsRequirementsKHR get_reqs;
	xrGetInstanceProcAddr(xr->instance, "xrGetOpenGLGraphicsRequirementsKHR",
	    reinterpret_cast<PFN_xrVoidFunction*>(&get_reqs));
	XrGraphicsRequirementsOpenGLKHR gl_reqs {
		XR_TYPE_GRAPHICS_REQUIREMENTS_OPENGL_KHR
	};
	get_reqs(xr->instance, system_id, &gl_reqs);

	Display* dpy = XOpenDisplay(nullptr);
	if (!dpy) {
		wlr_log(WLR_ERROR, "Failed to open X display");
		return false;
	}
	int screen = DefaultScreen(dpy);
	static int vis_attrs[] = { GLX_X_RENDERABLE, True, GLX_DRAWABLE_TYPE,
		GLX_WINDOW_BIT, GLX_RENDER_TYPE, GLX_RGBA_BIT, GLX_X_VISUAL_TYPE,
		GLX_TRUE_COLOR, GLX_RED_SIZE, 8, GLX_GREEN_SIZE, 8, GLX_BLUE_SIZE, 8,
		GLX_DEPTH_SIZE, 24, GLX_DOUBLEBUFFER, True, None };
	int fbcount;
	GLXFBConfig* fbcs = glXChooseFBConfig(dpy, screen, vis_attrs, &fbcount);
	if (!fbcs || !fbcount) {
		wlr_log(WLR_ERROR, "No GLXFBConfig found");
		return false;
	}
	GLXFBConfig fbc = fbcs[0];
	XFree(fbcs);

	XVisualInfo* vi = glXGetVisualFromFBConfig(dpy, fbc);
	Window root = RootWindow(dpy, screen);
	XSetWindowAttributes swa;
	swa.colormap = XCreateColormap(dpy, root, vi->visual, AllocNone);
	swa.event_mask = ExposureMask;
	Window win = XCreateWindow(dpy, root, 0, 0, 16, 16, 0, vi->depth,
	    InputOutput, vi->visual, CWColormap | CWEventMask, &swa);
	GLXContext ctx
	    = glXCreateNewContext(dpy, fbc, GLX_RGBA_TYPE, nullptr, True);
	glXMakeContextCurrent(dpy, win, win, ctx);

	XrGraphicsBindingOpenGLXlibKHR bind {
		XR_TYPE_GRAPHICS_BINDING_OPENGL_XLIB_KHR
	};
	bind.xDisplay = dpy;
	bind.visualid = vi->visualid;
	bind.glxFBConfig = fbc;
	bind.glxDrawable = win;
	bind.glxContext = ctx;

	XrSessionCreateInfo sci { XR_TYPE_SESSION_CREATE_INFO };
	sci.next = &bind;
	sci.systemId = system_id;
	if (xrCreateSession(xr->instance, &sci, &xr->session) != XR_SUCCESS) {
		wlr_log(WLR_ERROR, "xrCreateSession failed");
		return false;
	}
	// Create reference space
	XrReferenceSpaceCreateInfo spaceInfo {
		XR_TYPE_REFERENCE_SPACE_CREATE_INFO
	};
	spaceInfo.referenceSpaceType = XR_REFERENCE_SPACE_TYPE_LOCAL;
	spaceInfo.poseInReferenceSpace = { { 0, 0, 0, 1 }, { 0, 0, 0 } };
	if (xrCreateReferenceSpace(xr->session, &spaceInfo, &xr->app_space)
	    != XR_SUCCESS) {
		wlr_log(WLR_ERROR, "xrCreateReferenceSpace failed");
		return false;
	}
	// Get recommended view configuration
	uint32_t viewCount;
	xrEnumerateViewConfigurationViews(xr->instance, system_id,
	    XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO, 0, &viewCount, nullptr);
	std::vector<XrViewConfigurationView> viewConfigs(
	    viewCount, { XR_TYPE_VIEW_CONFIGURATION_VIEW });
	xrEnumerateViewConfigurationViews(xr->instance, system_id,
	    XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO, viewCount, &viewCount,
	    viewConfigs.data());
	xr->width = viewConfigs[0].recommendedImageRectWidth;
	xr->height = viewConfigs[0].recommendedImageRectHeight;
	// Create swapchain
	uint32_t formatCount;
	xrEnumerateSwapchainFormats(xr->session, 0, &formatCount, nullptr);
	std::vector<int64_t> formats(formatCount);
	xrEnumerateSwapchainFormats(
	    xr->session, formatCount, &formatCount, formats.data());
	int64_t swapFormat = formats.empty() ? 0 : formats[0];
	XrSwapchainCreateInfo swapchain_ci { XR_TYPE_SWAPCHAIN_CREATE_INFO };
	swapchain_ci.usageFlags = XR_SWAPCHAIN_USAGE_COLOR_ATTACHMENT_BIT;
	swapchain_ci.format = swapFormat;
	swapchain_ci.sampleCount = 1;
	swapchain_ci.width = xr->width;
	swapchain_ci.height = xr->height;
	swapchain_ci.faceCount = 1;
	swapchain_ci.arraySize = viewCount;
	swapchain_ci.mipCount = 1;
	if (xrCreateSwapchain(xr->session, &swapchain_ci, &xr->swapchain)
	    != XR_SUCCESS) {
		wlr_log(WLR_ERROR, "xrCreateSwapchain failed");
		return false;
	}
	// Enumerate swapchain images
	uint32_t imageCount;
	xrEnumerateSwapchainImages(xr->swapchain, 0, &imageCount, nullptr);
	xr->swapchain_images.resize(
	    imageCount, { XR_TYPE_SWAPCHAIN_IMAGE_OPENGL_KHR });
	xrEnumerateSwapchainImages(xr->swapchain, imageCount, &imageCount,
	    reinterpret_cast<XrSwapchainImageBaseHeader*>(
	        xr->swapchain_images.data()));
	// Create framebuffers for each image
	xr->framebuffers.resize(imageCount);
	for (uint32_t i = 0; i < imageCount; ++i) {
		GLuint fbo;
		glGenFramebuffers(1, &fbo);
		glBindFramebuffer(GL_FRAMEBUFFER, fbo);
		glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
		    GL_TEXTURE_2D, xr->swapchain_images[i].image, 0);
		glBindFramebuffer(GL_FRAMEBUFFER, 0);
		xr->framebuffers[i] = fbo;
	}

	xr->output = static_cast<wlr_output*>(calloc(1, sizeof(wlr_output)));
	if (xr->output) {
		wlr_output_init(
		    xr->output, &xr->base, &output_impl, xr->event_loop, nullptr);
		wlr_output_set_name(xr->output, "OpenXR");
		wlr_output_create_global(xr->output, xr->display);
	}

	xr->started = true;
	wlr_log(WLR_INFO, "OpenXR backend started");
	return true;
}

static void backend_destroy(wlr_backend* backend)
{
	auto* xr = reinterpret_cast<openxr_backend*>(backend);
	// destroy swapchain and framebuffers
	if (xr->swapchain) {
		xrDestroySwapchain(xr->swapchain);
		xr->swapchain = XR_NULL_HANDLE;
	}
	for (auto fbo : xr->framebuffers) {
		glDeleteFramebuffers(1, &fbo);
	}
	// destroy reference space
	if (xr->app_space) {
		xrDestroySpace(xr->app_space);
		xr->app_space = XR_NULL_HANDLE;
	}
	// destroy XR session and instance
	if (xr->session) {
		xrDestroySession(xr->session);
	}
	if (xr->instance) {
		xrDestroyInstance(xr->instance);
	}
	free(xr);
}

static int backend_get_drm_fd(wlr_backend* backend)
{
	(void)backend;
	return -1;
}

static wlr_backend_impl const backend_impl = {
	.start = backend_start,
	.destroy = backend_destroy,
	.get_drm_fd = backend_get_drm_fd,
};

wlr_backend* wlr_openxr_backend_create(wl_display* display, wl_event_loop* loop)
{
	openxr_backend* b = static_cast<openxr_backend*>(calloc(1, sizeof(*b)));

	if (!b)
		return nullptr;
	b->display = display;
	b->event_loop = loop;
	wlr_backend_init(&b->base, &backend_impl);
	return &b->base;
}

bool wlr_backend_is_openxr(wlr_backend* backend)
{
	return backend->impl == &backend_impl;
}