#include "LunarWM_core.h"

#include "LunarWM_render.h"
#include "LunarWM_wayland.h"
#include "LunarWM_xr.h"
#include "RayExt.h"
#include "common.h"
#include "vec.h"

#include <assert.h>
#include <math.h>
#include <spawn.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#include <lauxlib.h>
#include <lua.h>
#include <lualib.h>

extern char **environ;

static void cleanup_raylib_egl(LunarWM *wm)
{
	if (IsWindowReady()) {
		CloseWindow();
	}
}

static void cleanup_lua_cfg(LunarWM *wm)
{
	if (wm->cman) {
		config_manager_destroy(wm->cman);
		wm->cman = NULL;
	}
}

void LunarWM_set_recenter_from_camera(LunarWM *wm)
{
	Vector3 pos = wm->renderer.camera.position;
	Vector3 fwd
	    = Vector3Normalize(Vector3Subtract(wm->renderer.camera.target, pos));

	float len_xz = sqrtf(fwd.x * fwd.x + fwd.z * fwd.z);
	float yaw = (len_xz > 1e-6f) ? atan2f(fwd.x, fwd.z) : 0.0f;

	Quaternion q_step = QuaternionFromAxisAngle((Vector3) { 0, 1, 0 }, -yaw);
	Vector3 t_step = Vector3Negate(Vector3RotateByQuaternion(pos, q_step));

	Quaternion q_total = QuaternionMultiply(q_step, wm->xr.recenter_rot);
	Vector3 t_total = Vector3Add(
	    Vector3RotateByQuaternion(wm->xr.recenter_trans, q_step), t_step);

	wm->xr.recenter_rot = q_total;
	wm->xr.recenter_trans = t_total;
	wm->xr.recenter_active = true;
}

static void sync_config(LunarWM *wm)
{
	if (wm->cman->cfg.cubemap) {
		Skybox_init(&wm->renderer.skybox, wm->cman->cfg.cubemap);
	} else {
		Skybox_destroy(&wm->renderer.skybox);
	}

	if (IsTextureValid(wm->renderer.hud_rt.texture)) {
		UnloadTexture(wm->renderer.hud_rt.texture);
		wm->renderer.hud_rt.texture.id = 0;
		wm->renderer.hud_rt.texture.width = 0;
		wm->renderer.hud_rt.texture.height = 0;
	}

	int vw = (int)wm->cman->cfg.displays.virtual.resolution.x;
	int vh = (int)wm->cman->cfg.displays.virtual.resolution.y;
	int hud = wm->cman->cfg.displays.hud.size;

	LunarWM_wayland_update_virtual_outputs(wm, vw, vh, hud);
}

static int l_exec(lua_State *L)
{
	char const *cmd = luaL_checkstring(L, 1);

	pid_t pid;
	char *argv[] = { (char *)"sh", (char *)"-c", (char *)cmd, NULL };

	int rc = posix_spawnp(&pid, "sh", NULL, NULL, argv, environ);
	if (rc != 0) {
		lua_pushnil(L);
		lua_pushfstring(L, "posix_spawnp failed: %s", strerror(rc));
		return 2;
	}

	lua_pushinteger(L, (lua_Integer)pid);
	return 1;
}

static int l_cycle_next(lua_State *L)
{
	LunarWM *wm = &g_wm;
	if (vector_size(wm->wayland.v_toplevels) == 0) {
		lua_pushnil(L);
		return 1;
	}

	wm->wayland.current_focus++;
	if (wm->wayland.current_focus >= vector_size(wm->wayland.v_toplevels)) {
		wm->wayland.current_focus = 0;
	}

	LunarWM_Toplevel *tl = wm->wayland.v_toplevels[wm->wayland.current_focus];
	LunarWM_Toplevel_focus(tl);

	lua_pushnil(L);
	return 1;
}

static int l_recenter(lua_State *L)
{
	(void)L;

	if (g_wm.xr.available) {
		LunarWM_set_recenter_from_camera(&g_wm);
	}

	lua_pushnil(L);
	return 1;
}

static int l_reload_config(lua_State *L)
{
	LunarWM *wm = &g_wm;
	ConfigManager *cm = wm->cman;
	if (!cm) {
		lua_pushnil(L);
		return 1;
	}

	config_manager_reload(cm);
	sync_config(wm);
	lua_pushnil(L);
	return 1;
}

static int l_quit_compositor(lua_State *L)
{
	(void)L;
	LunarWM_terminate(&g_wm);
	lua_pushnil(L);
	return 1;
}

bool LunarWM_init(LunarWM *wm)
{
	memset(wm, 0, sizeof(*wm));
	wm->xr.session = XR_NULL_HANDLE;
	wm->xr.session_state = XR_SESSION_STATE_UNKNOWN;
	wm->xr.environment_blend_mode = XR_ENVIRONMENT_BLEND_MODE_MAX_ENUM;
	wm->xr.local_space = wm->xr.view_space = XR_NULL_HANDLE;
	wm->xr.hand_tracking_system_properties.type
	    = XR_TYPE_SYSTEM_HAND_TRACKING_PROPERTIES_EXT;
	wm->xr.hand_tracking_system_properties.next = NULL;
	wm->xr.hand_tracking_system_properties.supportsHandTracking = XR_FALSE;
	wm->xr.hand_tracking_enabled = false;
	wm->renderer.camera.position = (Vector3) { 0, 0, 0 };
	wm->renderer.camera.target = (Vector3) { 0, 0, 1 };
	wm->renderer.camera.up = (Vector3) { 0, 1, 0 };
	wm->renderer.camera.fovy = 45;
	wm->renderer.camera.projection = CAMERA_PERSPECTIVE;
	wm->xr.recenter_rot = (Quaternion) { 0, 0, 0, 1 };
	wm->xr.recenter_trans = (Vector3) { 0, 0, 0 };
	wm->xr.recenter_active = false;
	wm->counter = 0;

	wm->wm.active_workspace = 0;
	for (size_t i = 0; i < ARRAY_SZ(wm->wm.workspaces); i++) {
		wm->wm.workspaces[i].v_windows = vector_create();
	}

	wm->cman = config_manager_create(get_config_path());
	assert(wm->cman);

	{
		lua_State *L = wm->cman->L;

		lua_newtable(L);
		lua_pushcfunction(L, l_quit_compositor);
		lua_setfield(L, -2, "quit_compositor");
		lua_pushcfunction(L, l_reload_config);
		lua_setfield(L, -2, "reload_config");
		lua_pushcfunction(L, l_recenter);
		lua_setfield(L, -2, "recenter");
		lua_pushcfunction(L, l_cycle_next);
		lua_setfield(L, -2, "cycle_next");
		lua_pushcfunction(L, l_exec);
		lua_setfield(L, -2, "exec");
		lua_setglobal(L, "lunar");

		config_manager_reload(wm->cman);
	}

	// if (getenv("DISPLAY") != NULL || getenv("WAYLAND_DISPLAY") != NULL) {
	//	wlr_log(WLR_ERROR, "This compositor can only be ran in DRM mode.");
	//	return false;
	// }

	if (!LunarWM_wayland_init(wm)) {
		wlr_log(WLR_ERROR, "Failed to initialize wlroots");
		return false;
	}

	EGLSurface draw = eglGetCurrentSurface(EGL_DRAW);
	EGLSurface read = eglGetCurrentSurface(EGL_READ);
	if (eglMakeCurrent(wm->wayland.egl_display, EGL_NO_SURFACE, EGL_NO_SURFACE,
	        wm->wayland.egl_context)
	    == EGL_FALSE) {
		wlr_log(WLR_ERROR, "Failed to eglMakeCurrent");
		return false;
	}

	wm->xr.available = false;
	{
		char *no_xr = getenv("LWM_NO_XR");
		bool xr = true;

		if (no_xr != NULL && no_xr[0] != '\0')
			xr = false;

		if (xr) {
			if (!LunarWM_xr_init(wm)) {
				wlr_log(
				    WLR_ERROR, "Failed to initialize OpenXR! Disabling XR...");
				LunarWM_xr_cleanup(wm);
			} else {
				wm->xr.available = true;
			}
		}
	}

	wlr_log(WLR_INFO, "OpenGL ES version: %s", glGetString(GL_VERSION));
	InitWindow(0, 0, "");

	if (eglMakeCurrent(
	        wm->wayland.egl_display, draw, read, wm->wayland.egl_context)
	    == EGL_FALSE) {
		wlr_log(WLR_ERROR, "Failed to eglMakeCurrent");
		return false;
	}

	sync_config(wm);

	wm->initialized = true;

	return true;
}

void LunarWM_destroy(LunarWM *wm)
{
	if (!wm)
		return;

	eglMakeCurrent(wm->wayland.egl_display, EGL_NO_SURFACE, EGL_NO_SURFACE,
	    wm->wayland.egl_context);

	LunarWM_xr_cleanup(wm);
	cleanup_raylib_egl(wm);
	LunarWM_wayland_cleanup(wm);
	cleanup_lua_cfg(wm);

	for (size_t i = 0; i < ARRAY_SZ(wm->wm.workspaces); ++i) {
		if (wm->wm.workspaces[i].v_windows) {
			vector_free(wm->wm.workspaces[i].v_windows);
			wm->wm.workspaces[i].v_windows = NULL;
		}
	}

	memset(wm, 0, sizeof(*wm));
}

void LunarWM_terminate(LunarWM *wm)
{
	wlr_log(WLR_INFO, "Stopping compositor");
	wm->running = false;
}

void LunarWM_run(LunarWM *wm)
{
	assert(wm);
	assert(wm->initialized);

	wm->renderer.first_frame = true;

	if (!wlr_backend_start(wm->wayland.backend)) {
		wlr_log(WLR_ERROR, "Failed to start backend");
		return;
	}

	char const *socket = wl_display_add_socket_auto(wm->wayland.display);
	if (socket == NULL) {
		wlr_log(WLR_ERROR, "Failed to add wayland socket to display");
		return;
	}

	setenv("WAYLAND_DISPLAY", socket, 1);
	wlr_log(LOG_INFO, "Running on WAYLAND_DISPLAY=%s", socket);

	wm->running = true;

	struct timespec last, now;
	clock_gettime(CLOCK_MONOTONIC, &last);
	while (wm->running) {
		clock_gettime(CLOCK_MONOTONIC, &now);
		float dt = (now.tv_sec - last.tv_sec)
		    + (now.tv_nsec - last.tv_nsec) / 1000000000.0f;
		last = now;

		for (size_t i = 0; i < vector_size(wm->wayland.v_toplevels); i++) {
			LunarWM_Toplevel *tl = wm->wayland.v_toplevels[i];
			if (tl->surface) {
				wlr_surface_send_frame_done(tl->surface, &now);
			}
		}

		wl_display_flush_clients(wm->wayland.display);
		wl_event_loop_dispatch(wm->wayland.event_loop, 0);

		EGLSurface draw = eglGetCurrentSurface(EGL_DRAW);
		EGLSurface read = eglGetCurrentSurface(EGL_READ);
		if (eglMakeCurrent(wm->wayland.egl_display, EGL_NO_SURFACE,
		        EGL_NO_SURFACE, wm->wayland.egl_context)
		    == EGL_FALSE) {
			wlr_log(WLR_ERROR, "Failed to eglMakeCurrent");
			return;
		}

		WindowShouldClose();
		BeginDrawing();

		if (wm->xr.available) {
			LunarWM_xr_poll_events(wm);

			if (!wm->xr.session_running) {
				EndDrawing();
				continue;
			}

			XrFrameState frame_state = {
				.type = XR_TYPE_FRAME_STATE,
			};
			XrFrameWaitInfo frame_wait_info = {
				.type = XR_TYPE_FRAME_WAIT_INFO,
			};
			if (xrWaitFrame(wm->xr.session, &frame_wait_info, &frame_state)
			    != XR_SUCCESS) {
				wlr_log(WLR_ERROR, "Failed to wait for OpenXR frame");
				return;
			}

			XrFrameBeginInfo frame_begin_info = {
				.type = XR_TYPE_FRAME_BEGIN_INFO,
			};
			XrResult res = xrBeginFrame(wm->xr.session, &frame_begin_info);
			if (res != XR_FRAME_DISCARDED && res != XR_SUCCESS) {
				wlr_log(WLR_ERROR, "Failed to begin the OpenXR Frame: %d", res);
				return;
			}

			if (wm->xr.hand_tracking_enabled) {
				for (size_t i = 0; i < 2; i++) {
					LunarWM_Hand *hand = &wm->xr.hands[i];
					bool const unobstructed = true;
					if (!wm->xr.LocateHandJointsEXT
					    || hand->hand_tracker == XR_NULL_HANDLE) {
						continue;
					}

					XrHandJointsMotionRangeInfoEXT mri = {
						.type = XR_TYPE_HAND_JOINTS_MOTION_RANGE_INFO_EXT,
					};
					mri.handJointsMotionRange = unobstructed
					    ? XR_HAND_JOINTS_MOTION_RANGE_UNOBSTRUCTED_EXT
					    : XR_HAND_JOINTS_MOTION_RANGE_CONFORMING_TO_CONTROLLER_EXT;

					XrHandJointsLocateInfoEXT li = {
						.type = XR_TYPE_HAND_JOINTS_LOCATE_INFO_EXT,
						.next = &mri,
						.baseSpace = wm->xr.local_space,
						.time = frame_state.predictedDisplayTime,
					};

					XrHandJointLocationsEXT hji = {
						.type = XR_TYPE_HAND_JOINT_LOCATIONS_EXT,
						.jointCount = XR_HAND_JOINT_COUNT_EXT,
						.jointLocations = hand->joint_locations,
					};

					if (wm->xr.LocateHandJointsEXT(
					        hand->hand_tracker, &li, &hji)
					    != XR_SUCCESS) {
						wlr_log(WLR_ERROR, "Failed to locate hand joints");
						return;
					}
				}
			}

			LunarWM_RenderLayerInfo render_layer_info = {
				.predicted_display_time = frame_state.predictedDisplayTime,
				.layer_projection.type = XR_TYPE_COMPOSITION_LAYER_PROJECTION,
			};
			bool const session_active
			    = (wm->xr.session_state == XR_SESSION_STATE_SYNCHRONIZED)
			    || (wm->xr.session_state == XR_SESSION_STATE_VISIBLE)
			    || (wm->xr.session_state == XR_SESSION_STATE_FOCUSED);
			if (session_active && (frame_state.shouldRender != 0u)) {
				bool rendered
				    = LunarWM_render_layer(wm, &render_layer_info, dt);
				if (rendered) {
					render_layer_info.layers[render_layer_info.layers_count]
					    = (XrCompositionLayerBaseHeader *)&render_layer_info
					          .layer_projection;
				}
			}

			XrFrameEndInfo frame_end_info = {
				.type = XR_TYPE_FRAME_END_INFO,
				.displayTime = frame_state.predictedDisplayTime,
				.environmentBlendMode = wm->xr.environment_blend_mode,
				.layerCount = render_layer_info.layers_count,
				.layers = (XrCompositionLayerBaseHeader const **)
				    render_layer_info.layers,
			};
			if (xrEndFrame(wm->xr.session, &frame_end_info) != XR_SUCCESS) {
				wlr_log(WLR_ERROR, "Failed to end OpenXR frame");
				return;
			}
		} else {
			wm->renderer.camera.fovy = 75;
		}

		EndDrawing();
	}
}