module;

#include <cassert>

#include <EGL/egl.h>
#include <EGL/eglext.h>
#include <GLES2/gl2.h>
#include <GLES2/gl2ext.h>
#include <openxr/openxr.h>
#include <openxr/openxr_platform.h>

#include <wayland-client-protocol.h>
#include <wayland-egl.h>
#include <wayland-server-core.h>

extern "C" {
#include <wlr/backend/wayland.h>
#include <wlr/render/allocator.h>
#include <wlr/render/gles2.h>
#include <wlr/types/wlr_compositor.h>
#include <wlr/types/wlr_data_device.h>
#include <wlr/types/wlr_subcompositor.h>
#include <wlr/util/log.h>
#include <wlr/render/egl.h>
}

export module LunarWM.LunarWM;

import std;

namespace LunarWM {

export struct LunarWM {
  LunarWM() = default;
  ~LunarWM();

  void init();

  void run();
  void terminate();

private:
  void init_wayland();
  void init_xr();

	void poll_events_xr();

  bool m_initialized{};

  struct {
    wl_display *display{};
    wl_event_loop *event_loop{};

    wlr_backend *backend{};
    wlr_renderer *renderer{};

    wlr_egl *egl{};
    EGLDisplay egl_display;
    EGLContext egl_context;
    EGLConfig egl_config;

    wlr_allocator *allocator{};
    wlr_compositor *compositor{};
    wlr_subcompositor *subcompositor{};
    wlr_data_device_manager *data_device_manager{};

    wlr_seat *seat{};
    wl_list keyboards;
    wl_listener new_input_listener{};
  } m_wayland;

  struct {
    std::optional<XrInstance> instance;
		std::optional<XrSystemId> system_id;
		XrSession session = XR_NULL_HANDLE;
		XrSessionState session_state = XR_SESSION_STATE_UNKNOWN;
  } m_xr;

  bool m_running{};
  bool m_session_running{};
  bool m_session_state{};
};

void LunarWM::init() {
  this->init_wayland();
  this->init_xr();

  m_initialized = true;
}

void LunarWM::init_wayland() {
  wlr_log_init(WLR_DEBUG, NULL);

  m_wayland.display = wl_display_create();
  if (!m_wayland.display) {
    throw std::runtime_error("Failed to create wayland display");
  }

  m_wayland.event_loop = wl_display_get_event_loop(m_wayland.display);
  if (!m_wayland.event_loop) {
    throw std::runtime_error("Failed to get wayland event loop");
  }

  m_wayland.backend = wlr_backend_autocreate(m_wayland.event_loop, nullptr);
  if (!m_wayland.backend) {
    throw std::runtime_error("Failed to create wlroots backend");
  }

	setenv("WLR_RENDERER", "gles2", 1);
  m_wayland.renderer = wlr_renderer_autocreate(m_wayland.backend);
  if (!m_wayland.renderer) {
    throw std::runtime_error("Failed to create wlroots renderer");
  }

  m_wayland.egl = wlr_gles2_renderer_get_egl(m_wayland.renderer);
  if (!m_wayland.egl) {
    throw std::runtime_error("Failed to get egl information from renderer");
  }

  m_wayland.egl_display = wlr_egl_get_display(m_wayland.egl);
  m_wayland.egl_context = wlr_egl_get_context(m_wayland.egl);
  m_wayland.egl_config = EGL_NO_CONFIG_KHR;

  if (!wlr_renderer_init_wl_display(m_wayland.renderer, m_wayland.display)) {
    throw std::runtime_error(
        "Failed to initialize renderer with wayland display");
  }

  m_wayland.allocator =
      wlr_allocator_autocreate(m_wayland.backend, m_wayland.renderer);
  if (!m_wayland.allocator) {
    throw std::runtime_error("Failed to create wlroots allocator");
  }

  m_wayland.compositor =
      wlr_compositor_create(m_wayland.display, 5, m_wayland.renderer);
  if (!m_wayland.compositor) {
    throw std::runtime_error("Failed to create wlroots compositor");
  }

  m_wayland.subcompositor = wlr_subcompositor_create(m_wayland.display);
  if (!m_wayland.subcompositor) {
    throw std::runtime_error("Failed to create wlroots subcompositor");
  }

  m_wayland.data_device_manager =
      wlr_data_device_manager_create(m_wayland.display);
  if (!m_wayland.data_device_manager) {
    throw std::runtime_error("Failed to create wlroots data device manager");
  }

  wl_list_init(&m_wayland.keyboards);

  m_wayland.new_input_listener.notify = [](wl_listener *listener, void *data) {
    auto wm = reinterpret_cast<LunarWM *>(
        wl_container_of(listener, static_cast<LunarWM *>(nullptr),
                        m_wayland.new_input_listener));
    auto dev = reinterpret_cast<wlr_input_device *>(data);
    if (dev->type == WLR_INPUT_DEVICE_KEYBOARD) {
      // FIXME: Implement
    } else if (dev->type == WLR_INPUT_DEVICE_POINTER) {
      // FIXME: Implement
    }

    std::uint32_t caps = WL_SEAT_CAPABILITY_POINTER;
    if (!wl_list_empty(&wm->m_wayland.keyboards)) {
      caps |= WL_SEAT_CAPABILITY_KEYBOARD;
    }
    assert(wm->m_wayland.seat);
    wlr_seat_set_capabilities(wm->m_wayland.seat, caps);
  };
  wl_signal_add(&m_wayland.backend->events.new_input,
                &m_wayland.new_input_listener);
  m_wayland.seat = wlr_seat_create(m_wayland.display, "seat0");
  if (!m_wayland.seat) {
    throw std::runtime_error("Failed to create wlroots seat");
  }
}

void LunarWM::init_xr() {
  XrApplicationInfo app_info{
		.applicationVersion = 1,
		.engineVersion = 1,
		.apiVersion = XR_CURRENT_API_VERSION,
  };
	strncpy(app_info.applicationName, "LunarWM", XR_MAX_APPLICATION_NAME_SIZE);
	strncpy(app_info.engineName, "LunarWM Engine", XR_MAX_ENGINE_NAME_SIZE);

	std::vector<std::string> instance_extensions {
		XR_EXT_DEBUG_UTILS_EXTENSION_NAME,
		XR_MNDX_EGL_ENABLE_EXTENSION_NAME,
		XR_KHR_OPENGL_ES_ENABLE_EXTENSION_NAME,
	};
	std::vector<std::string> apiLayers;
	std::vector<char const *> active_instance_extensions;
	std::vector<char const *> activeAPILayers;

	uint32_t apiLayerCount = 0;
	std::vector<XrApiLayerProperties> apiLayerProperties;
	if (xrEnumerateApiLayerProperties(0, &apiLayerCount, nullptr) != XR_SUCCESS) {
		throw std::runtime_error("Failed to enumerate API layer properties");
	}
	apiLayerProperties.resize(apiLayerCount, {XR_TYPE_API_LAYER_PROPERTIES});
	if (xrEnumerateApiLayerProperties(apiLayerCount, &apiLayerCount, apiLayerProperties.data()) != XR_SUCCESS) {
		throw std::runtime_error("Failed to enumerate API layer properties");
	}

	for (auto &requestLayer : apiLayers) {
		for (auto &layerProperty : apiLayerProperties) {
			if (strcmp(requestLayer.c_str(), layerProperty.layerName) != 0) {
				continue;
			} else {
					activeAPILayers.push_back(requestLayer.c_str());
				break;
			}
		}
	}

	uint32_t extensionCount = 0;
	std::vector<XrExtensionProperties> extensionProperties;
	if (xrEnumerateInstanceExtensionProperties(nullptr, 0, &extensionCount, nullptr) != XR_SUCCESS) {
		throw std::runtime_error("Failed to enumerate OpenXR instance extension properties");
	}
	extensionProperties.resize(extensionCount, {XR_TYPE_EXTENSION_PROPERTIES});
	if (xrEnumerateInstanceExtensionProperties(nullptr, extensionCount, &extensionCount, extensionProperties.data()) != XR_SUCCESS) {
		throw std::runtime_error("Failed to enumerate OpenXR instance extension properties");
	}

	for (auto &requestedInstanceExtension : instance_extensions) {
		bool found = false;
		for (auto &extensionProperty : extensionProperties) {
			if (strcmp(requestedInstanceExtension.c_str(), extensionProperty.extensionName) != 0) {
				continue;
			} else {
				active_instance_extensions.push_back(requestedInstanceExtension.c_str());
				found = true;
				break;
			}
		}
		if (!found) {
			throw std::runtime_error(std::format("Failed to find OpenXR instance extension: {}", requestedInstanceExtension));
		}
	}

	{
		XrInstanceCreateInfo ci {
			.type = XR_TYPE_INSTANCE_CREATE_INFO,
			.next = NULL,
			.createFlags = 0,
			.applicationInfo = app_info,
			.enabledApiLayerCount = static_cast<std::uint32_t>(activeAPILayers.size()),
			.enabledApiLayerNames = activeAPILayers.data(),
			.enabledExtensionCount = static_cast<std::uint32_t>(active_instance_extensions.size()),
			.enabledExtensionNames = active_instance_extensions.data(),
		};

		XrInstance instance;
		if (xrCreateInstance(&ci, &instance) != XR_SUCCESS) {
			throw std::runtime_error("Failed to create OpenXR instance");
		}
		m_xr.instance = instance;
	}

	{
		XrSystemGetInfo gi {
			.type = XR_TYPE_SYSTEM_GET_INFO,
			.next = nullptr,
		};

		XrFormFactor factors[] {
			XR_FORM_FACTOR_HEAD_MOUNTED_DISPLAY,
			XR_FORM_FACTOR_HANDHELD_DISPLAY,
		};
		for (auto const factor : factors) {
			gi.formFactor = factor;
			XrSystemId system_id;
			if (xrGetSystem(*m_xr.instance, &gi, &system_id) == XR_SUCCESS) {
				m_xr.system_id = system_id;
				break;
			}
		}

		if (!m_xr.system_id) {
			throw std::runtime_error("Failed to find valid form factor");
		}
	}

	XrGraphicsRequirementsOpenGLESKHR reqs {
		.type = XR_TYPE_GRAPHICS_REQUIREMENTS_OPENGL_ES_KHR,
		.next = nullptr,
	};
	PFN_xrGetOpenGLESGraphicsRequirementsKHR xrGetOpenGLESGraphicsRequirementsKHR;
	xrGetInstanceProcAddr(*m_xr.instance, "xrGetOpenGLESGraphicsRequirementsKHR", (PFN_xrVoidFunction*)&xrGetOpenGLESGraphicsRequirementsKHR);
	if (xrGetOpenGLESGraphicsRequirementsKHR(*m_xr.instance, *m_xr.system_id, &reqs) != XR_SUCCESS) {
		throw std::runtime_error("Failed to get GLES graphics requirements");
	}
	printf("OpenGL ES range: %d.%d.%d  –  %d.%d.%d\n",
		XR_VERSION_MAJOR(reqs.minApiVersionSupported),
		XR_VERSION_MINOR(reqs.minApiVersionSupported),
		XR_VERSION_PATCH(reqs.minApiVersionSupported),
		XR_VERSION_MAJOR(reqs.maxApiVersionSupported),
		XR_VERSION_MINOR(reqs.maxApiVersionSupported),
		XR_VERSION_PATCH(reqs.maxApiVersionSupported));

	glEnable(GL_DEBUG_OUTPUT_KHR);
	glEnable(GL_DEBUG_OUTPUT_SYNCHRONOUS_KHR);

	{
		XrGraphicsBindingEGLMNDX gbind = {
			.type = XR_TYPE_GRAPHICS_BINDING_EGL_MNDX,
			.next = nullptr,
			.getProcAddress = eglGetProcAddress,
			.display = m_wayland.egl_display,
			.config = m_wayland.egl_config,
			.context = m_wayland.egl_context,
		};

		XrSessionCreateInfo ci {
			.type = XR_TYPE_SESSION_CREATE_INFO,
			.next = &gbind,
			.createFlags = 0,
			.systemId = *m_xr.system_id,
		};

		if (xrCreateSession(*m_xr.instance, &ci, &m_xr.session) != XR_SUCCESS) {
			throw std::runtime_error("Failed to create OpenXR session");
		}
	}
}

void LunarWM::poll_events_xr() {
	XrEventDataBuffer event_data{XR_TYPE_EVENT_DATA_BUFFER};
	auto XrPollEvents = [&]() -> bool {
		event_data = {XR_TYPE_EVENT_DATA_BUFFER};
		return xrPollEvent(*m_xr.instance, &event_data) == XR_SUCCESS;
	};

	while (XrPollEvents()) {
		switch (event_data.type) {
		case XR_TYPE_EVENT_DATA_EVENTS_LOST: {
			XrEventDataEventsLost *eventsLost = reinterpret_cast<XrEventDataEventsLost *>(&event_data);
			wlr_log(WLR_INFO, "OPENXR: Events Lost: %d", eventsLost->lostEventCount);
			break;
		}
		case XR_TYPE_EVENT_DATA_INSTANCE_LOSS_PENDING: {
			XrEventDataInstanceLossPending *instanceLossPending = reinterpret_cast<XrEventDataInstanceLossPending *>(&event_data);
			wlr_log(WLR_INFO, "OPENXR: Instance Loss Pending at: %ld", instanceLossPending->lossTime);
			m_session_running = false;
			m_running = false;
			break;
		}
		case XR_TYPE_EVENT_DATA_INTERACTION_PROFILE_CHANGED: {
			XrEventDataInteractionProfileChanged *interactionProfileChanged = reinterpret_cast<XrEventDataInteractionProfileChanged *>(&event_data);
			wlr_log(WLR_INFO, "OPENXR: Interaction Profile changed for Session: %p", interactionProfileChanged->session);
			if (interactionProfileChanged->session != m_xr.session) {
				wlr_log(WLR_ERROR, "XrEventDataInteractionProfileChanged for unknown Session");
				break;
			}
			break;
		}
		case XR_TYPE_EVENT_DATA_REFERENCE_SPACE_CHANGE_PENDING: {
			XrEventDataReferenceSpaceChangePending *referenceSpaceChangePending = reinterpret_cast<XrEventDataReferenceSpaceChangePending *>(&event_data);
			wlr_log(WLR_INFO, "OPENXR: Reference Space Change pending for Session: %p", referenceSpaceChangePending->session);
			if (referenceSpaceChangePending->session != m_xr.session) {
				wlr_log(WLR_ERROR, "XrEventDataReferenceSpaceChangePending for unknown Session");
				break;
			}
			break;
		}
		case XR_TYPE_EVENT_DATA_SESSION_STATE_CHANGED: {
			XrEventDataSessionStateChanged *sessionStateChanged = reinterpret_cast<XrEventDataSessionStateChanged *>(&event_data);
			if (sessionStateChanged->session != m_xr.session) {
				wlr_log(WLR_ERROR, "XrEventDataSessionStateChanged for unknown Session");
				break;
			}

			if (sessionStateChanged->state == XR_SESSION_STATE_READY) {
				XrSessionBeginInfo sessionBeginInfo{XR_TYPE_SESSION_BEGIN_INFO};
				sessionBeginInfo.primaryViewConfigurationType = XR_VIEW_CONFIGURATION_TYPE_PRIMARY_STEREO;
				if (xrBeginSession(m_xr.session, &sessionBeginInfo) != XR_SUCCESS) {
					throw std::runtime_error("Failed to begin session");
				}
				m_session_running = true;
			}
			if (sessionStateChanged->state == XR_SESSION_STATE_STOPPING) {
				if (xrEndSession(m_xr.session) != XR_SUCCESS) {
					throw std::runtime_error("Failed to end session");
				}
				m_session_running = false;
			}
			if (sessionStateChanged->state == XR_SESSION_STATE_EXITING) {
				m_session_running = false;
				m_running = false;
			}
			if (sessionStateChanged->state == XR_SESSION_STATE_LOSS_PENDING) {
				m_session_running = false;
				m_running = false;
			}
			m_xr.session_state = sessionStateChanged->state;
			break;
		}
		default: {
			break;
		}
		}
	}
}

LunarWM::~LunarWM() {
  assert(m_initialized);

	puts("1");

	if (m_xr.session != XR_NULL_HANDLE) {
		xrDestroySession(m_xr.session);
	}
	puts("2");

	if (m_xr.instance) {
		xrDestroyInstance(*m_xr.instance);
	}
	puts("3");

  wl_list_remove(&m_wayland.keyboards);

	puts("4");
  wl_display_destroy_clients(m_wayland.display);
  if (!m_wayland.allocator) {
    wlr_allocator_destroy(m_wayland.allocator);
  }
	puts("5");
  if (!m_wayland.renderer) {
    wlr_renderer_destroy(m_wayland.renderer);
  }
	puts("6");
  if (!m_wayland.backend) {
    wlr_backend_destroy(m_wayland.backend);
  }
	puts("7");
  if (!m_wayland.display) {
    wl_display_destroy(m_wayland.display);
  }
}

void LunarWM::run() {
  if (!wlr_backend_start(m_wayland.backend)) {
    throw std::runtime_error("Failed to start backend");
  }

  auto const *socket = wl_display_add_socket_auto(m_wayland.display);
  if (!socket) {
    throw std::runtime_error("Failed to add wayland socket to display");
  }

  setenv("WAYLAND_DISPLAY", socket, true);
  wlr_log(WLR_INFO, "Running compositor on WAYLAND_DISPLAY=%s", socket);

  m_running = true;
  while (m_running) {
    wl_display_flush_clients(m_wayland.display);
    wl_event_loop_dispatch(m_wayland.event_loop, 0);

		poll_events_xr();
  }
}

void LunarWM::terminate() {
	wlr_log(WLR_INFO, "Stopping compositor");
	m_running = false;
}

} // namespace LunarWM