// SPDX-License-Identifier: AGPL-3.0-only
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <unistd.h>
#include <vector>
#include <string>
#include <algorithm>
#include <fstream>
#include <sstream>

#include <xkbcommon/xkbcommon.h>
#include <xkbcommon/xkbcommon-keysyms.h>

#include <wayland-server-core.h>

#include "wlr/display.hpp"
#include "wlr/backend.hpp"
#include "wlr/compositor.hpp"
#include "wlr/output.hpp"
#include "wlr/xdg_shell.hpp"
#include "wlr/seat.hpp"
#include "wlr/scene.hpp"
#include "wlr/node.hpp"

struct ToplevelState {
	wlr::XdgSurface xdg;
	wlr::Tree tree;
	bool mapped = false;
	int gx = 0, gy = 0, gw = 0, gh = 0;
	std::string title;
};

struct Server {
	wlr::Display display;
	wlr::Backend backend{display};
	wlr::Renderer renderer{wlr_renderer_autocreate(backend.raw())};
	wlr::Allocator allocator{wlr_allocator_autocreate(backend.raw(), renderer.raw())};
	wlr::Compositor compositor{display, renderer};
	wlr::XdgShell xdg_shell{display};
	wlr::Seat seat{display, "seat0"};
	wlr::Cursor cursor;
	wlr::OutputLayout layout{display};
	wlr::Scene scene;
	wlr::Tree root{wlr_scene_tree_create(scene.root())};

	std::vector<struct wlr_scene_output*> outputs;
	std::vector<ToplevelState*> toplevels;
	ToplevelState* focused = nullptr;
	struct wlr_keyboard* keyboard = nullptr;
	wlr::SignalSink key_sink;
	wl_event_source* scheduler_source = nullptr;
	const char* socket_name = nullptr;

	uint32_t mod = WLR_MODIFIER_ALT;
	std::string terminal = "kitty";
	int caps = 0;
};

static Server* S = nullptr;

static void relayout() {
	struct wlr_box b = S->layout.box();
	int n = 0;
	for (auto* t : S->toplevels)
		if (t->mapped) n++;
	if (n == 0)
		return;
	if (n == 1) {
		for (auto* t : S->toplevels) {
			if (!t->mapped) continue;
			t->gx = b.x; t->gy = b.y; t->gw = b.width; t->gh = b.height;
			t->tree.set_position(b.x, b.y);
			wlr::Toplevel(t->xdg.toplevel()).set_size(b.width, b.height);
		}
		return;
	}
	int mw = b.width / 2;
	int i = 0;
	for (auto* t : S->toplevels) {
		if (!t->mapped) continue;
		if (i == 0) {
			t->gx = b.x; t->gy = b.y; t->gw = mw; t->gh = b.height;
			t->tree.set_position(b.x, b.y);
			wlr::Toplevel(t->xdg.toplevel()).set_size(mw, b.height);
		} else {
			int sh = b.height / (n - 1);
			t->gx = b.x + mw; t->gy = b.y + (i - 1) * sh; t->gw = b.width - mw; t->gh = sh;
			t->tree.set_position(t->gx, t->gy);
			wlr::Toplevel(t->xdg.toplevel()).set_size(b.width - mw, sh);
		}
		i++;
	}
}

static void focus(ToplevelState* t) {
	if (S->focused == t)
		return;
	if (S->focused)
		wlr::Toplevel(S->focused->xdg.toplevel()).set_activated(false);
	S->focused = t;
	if (!t) {
		S->seat.keyboard_notify_enter(nullptr);
		return;
	}
	wlr::Toplevel(t->xdg.toplevel()).set_activated(true);
	t->tree.raise_to_top();
	S->seat.keyboard_notify_enter(t->xdg.surface());
}

static void cycle_focus(int dir) {
	std::vector<ToplevelState*> vis;
	for (auto* t : S->toplevels)
		if (t->mapped) vis.push_back(t);
	if (vis.empty())
		return;
	int idx = 0;
	for (size_t i = 0; i < vis.size(); i++)
		if (vis[i] == S->focused) idx = (int)i;
	idx = (idx + dir + (int)vis.size()) % (int)vis.size();
	focus(vis[idx]);
}

static void spawn() {
	pid_t p = fork();
	if (p == 0) {
		if (S->socket_name)
			setenv("WAYLAND_DISPLAY", S->socket_name, 1);
		execlp(S->terminal.c_str(), S->terminal.c_str(), nullptr);
		_exit(127);
	}
}

static void on_key(void* data) {
	auto* ev = static_cast<struct wlr_keyboard_key_event*>(data);
	if (ev->state != WL_KEYBOARD_KEY_STATE_PRESSED) {
		if (S->focused)
			S->seat.keyboard_notify_key(ev->time_msec, ev->keycode, ev->state);
		return;
	}
	struct wlr_keyboard* kb = S->keyboard;
	if (!kb) {
		if (S->focused)
			S->seat.keyboard_notify_key(ev->time_msec, ev->keycode, ev->state);
		return;
	}
	uint32_t mods = wlr_keyboard_get_modifiers(kb);
	xkb_keysym_t sym = xkb_state_key_get_one_sym(kb->xkb_state, ev->keycode + 8);
	if (mods & S->mod) {
		if (sym == XKB_KEY_Return) { spawn(); return; }
		if (sym == XKB_KEY_q) { if (S->focused) wlr::Toplevel(S->focused->xdg.toplevel()).close(); return; }
		if (sym == XKB_KEY_j) { cycle_focus(1); return; }
		if (sym == XKB_KEY_k) { cycle_focus(-1); return; }
		if (sym == XKB_KEY_Tab) { cycle_focus(1); return; }
	}
	if (S->focused)
		S->seat.keyboard_notify_key(ev->time_msec, ev->keycode, ev->state);
}

static void on_motion(void* data) {
	auto* ev = static_cast<struct wlr_pointer_motion_event*>(data);
	S->seat.pointer_notify_motion(ev->time_msec, S->cursor.raw()->x, S->cursor.raw()->y);
}

static void on_button(void* data) {
	auto* ev = static_cast<struct wlr_pointer_button_event*>(data);
	S->seat.pointer_notify_button(ev->time_msec, ev->button, ev->state);
	if (ev->state == WL_POINTER_BUTTON_STATE_PRESSED) {
		double cx = S->cursor.raw()->x, cy = S->cursor.raw()->y;
		for (int i = (int)S->toplevels.size() - 1; i >= 0; i--) {
			ToplevelState* t = S->toplevels[i];
			if (t->mapped && cx >= t->gx && cx < t->gx + t->gw &&
					cy >= t->gy && cy < t->gy + t->gh) {
				focus(t);
				break;
			}
		}
	}
}

static void on_axis(void* data) {
	auto* ev = static_cast<struct wlr_pointer_axis_event*>(data);
	S->seat.pointer_notify_axis(ev->time_msec, ev->orientation, ev->delta, ev->delta_discrete, ev->source);
}

static void on_frame(void*) {
	S->seat.pointer_notify_frame();
}

static void on_new_toplevel(void* data) {
	auto* t = static_cast<struct wlr_xdg_toplevel*>(data);
	ToplevelState* tn = new ToplevelState{
		wlr::XdgSurface{t->base},
		wlr::Tree(wlr_scene_xdg_surface_create(S->root.raw(), t->base)),
	};
	S->toplevels.push_back(tn);

	tn->xdg.on_map([tn](void*) {
		tn->mapped = true;
		relayout();
		focus(tn);
	});
	tn->xdg.on_unmap([tn](void*) {
		tn->mapped = false;
		relayout();
		if (S->focused == tn)
			cycle_focus(1);
	});
	tn->xdg.on_destroy([tn](void*) {
		auto it = std::find(S->toplevels.begin(), S->toplevels.end(), tn);
		if (it != S->toplevels.end())
			S->toplevels.erase(it);
		if (S->focused == tn)
			S->focused = nullptr;
		if (tn->tree)
			wlr_scene_node_destroy(&tn->tree.raw()->node);
		delete tn;
		relayout();
	});

	wlr::Toplevel(t).on_set_title([tn](void* d) {
		const char* s = static_cast<struct wlr_xdg_toplevel*>(d)->title;
		tn->title = s ? s : "";
	});

	wl_event_loop_add_idle(S->display.event_loop(),
		[](void* d) {
			auto* s = static_cast<struct wlr_xdg_surface*>(d);
			if (s->initialized)
				wlr_xdg_surface_schedule_configure(s);
		}, t->base);
}

static void on_new_output(void* data) {
	auto* o = static_cast<struct wlr_output*>(data);
	wlr::Output out(o);
	out.set_preferred_mode();
	wlr_output_init_render(o, S->allocator.raw(), S->renderer.raw());
	wlr_output_create_global(o, S->display.raw());
	struct wlr_scene_output* so = wlr_scene_output_create(S->scene.raw(), o);
	wlr_scene_output_set_position(so, 0, 0);
	S->outputs.push_back(so);
	S->layout.add_auto_output(o);
	wl_event_source_timer_update(S->scheduler_source, 0);
	static bool spawned = false;
	if (!spawned) {
		spawned = true;
		spawn();
	}
}

static void on_new_input(void* data) {
	auto* dev = static_cast<struct wlr_input_device*>(data);
	if (dev->type == WLR_INPUT_DEVICE_KEYBOARD) {
		struct wlr_keyboard* kb = wlr_keyboard_from_input_device(dev);
		wlr::Keyboard k(kb);
		k.set_default_keymap();
		S->keyboard = kb;
		S->seat.set_keyboard(k);
		S->key_sink.connect(&kb->events.key, on_key);
		S->caps |= WL_SEAT_CAPABILITY_KEYBOARD;
	} else {
		S->cursor.attach_input_device(dev);
		if (dev->type == WLR_INPUT_DEVICE_POINTER)
			S->caps |= WL_SEAT_CAPABILITY_POINTER;
		else if (dev->type == WLR_INPUT_DEVICE_TOUCH)
			S->caps |= WL_SEAT_CAPABILITY_TOUCH;
	}
	S->seat.set_capabilities((uint32_t)S->caps);
}

static void commit_outputs() {
	for (struct wlr_scene_output* so : S->outputs)
		wlr_scene_output_commit(so, nullptr);
}

static int frame_tick(void* data) {
	(void)data;
	commit_outputs();
	wl_event_source_timer_update(S->scheduler_source, 16);
	return 0;
}

static void load_config() {
	const char* home = getenv("HOME");
	std::string path = std::string(home ? home : "") + "/.config/testwm/testwm.conf";
	if (getenv("TESTWM_CONFIG"))
		path = getenv("TESTWM_CONFIG");
	std::ifstream f(path);
	if (!f)
		return;
	std::string line;
	while (std::getline(f, line)) {
		line.erase(0, line.find_first_not_of(" \t"));
		if (line.empty() || line[0] == '#')
			continue;
		auto eq = line.find('=');
		if (eq == std::string::npos)
			continue;
		std::string k = line.substr(0, eq), v = line.substr(eq + 1);
		k.erase(k.find_last_not_of(" \t") + 1);
		v.erase(0, v.find_first_not_of(" \t"));
		v.erase(v.find_last_not_of(" \t") + 1);
		if (k == "terminal") S->terminal = v;
		else if (k == "modifier") S->mod = (v == "super" || v == "win" || v == "logo") ? WLR_MODIFIER_LOGO : WLR_MODIFIER_ALT;
	}
}

int main() {
	Server server;
	S = &server;

	load_config();

	if (!S->renderer.init_wl_display(S->display))
		return 1;

	S->cursor.attach_output_layout(S->layout.raw());
	S->cursor.on_motion(on_motion);
	S->cursor.on_button(on_button);
	S->cursor.on_axis(on_axis);
	S->cursor.on_frame(on_frame);

	float bg_color[4] = {0.08f, 0.08f, 0.12f, 1.0f};
	wlr::Rect bg(wlr_scene_rect_create(S->scene.root(), 4096, 4096, bg_color));
	(void)bg;

	S->backend.on_new_output(on_new_output);
	S->backend.on_new_input(on_new_input);
	S->xdg_shell.on_new_toplevel(on_new_toplevel);

	S->scheduler_source = wl_event_loop_add_timer(S->display.event_loop(), frame_tick, nullptr);

	const char* socket = S->display.add_socket_auto();
	S->socket_name = socket;
	if (!socket)
		return 1;
	fprintf(stderr,
		"testWM: WAYLAND_DISPLAY=%s\n"
		"testWM: bindings: Alt+Enter spawn '%s', Alt+Q close, Alt+J/Alt+K cycle focus, click to focus\n",
		socket, S->terminal.c_str());
	fflush(stderr);

	if (!S->backend.start())
		return 1;

	S->display.run();
	return 0;
}