#define _POSIX_C_SOURCE 200112L
#include <stdio.h>
#include <string.h>
#include <sys/mman.h>
#include <wayland-client.h>
#include <xkbcommon/xkbcommon.h>
#include <malloc.h>
#include <unistd.h>
#include <poll.h>
#include <time.h>
#include <sys/timerfd.h>
#include <errno.h>
#include <stdlib.h>

#include "xdg-shell-protocol.h"
#include "dwl-bar-ipc-protocol.h"
#include "wlr-layer-shell-protocol.h"

#include "drw.h"
#include "util.h"
#include "compositor.h"

/* macro definitions */
#define MAX_LINE_LENGTH 1024
#define match_then_bind(obj, inter, ver) \
	if (strcmp(interface, inter.name) == 0) { \
		obj = wl_registry_bind(wl_registry, name, &inter, ver);

#define end_match }
#define or_match } else

/* struct definitions */
typedef struct {
	struct wl_display *display;
	struct wl_registry *registry;
	struct wl_compositor *compositor;
	struct wl_shm *shm;
	struct zdwl_manager_v1 *dwl_manager;
	struct zwlr_layer_shell_v1 *layer;
	struct wl_seat *seat;
	struct wl_keyboard *keyboard;
	struct xkb_context *xkb_context;
	struct xkb_state *xkb_state;
	struct xkb_keymap *xkb_keymap;
	struct {
		uint32_t control;
		uint32_t alt;
		uint32_t shift;
		uint32_t sym;
		uint32_t state;
	} keys;
	struct {
		int timer;
		int delay;
		int period;
	} repeat;
} Client;

typedef struct Monitor Monitor;
struct Monitor {
	uint32_t width;
	Canvas *canvas;
	Font *font;
	struct zdwl_output_v1 *dwl_output;
	struct wl_output *output;
	struct wl_surface *surface;
	struct zwlr_layer_surface_v1 *wlr_surface;
	Monitor *next;
};

typedef struct Option Option;
struct Option {
	char *name;
	Option *next;
	Option *previous;
	Option *next_match;
	Option *previous_match;
};

typedef struct Options Options;
struct Options {
	Option *first;
	Option *last;
	Option *first_match;
	Option *last_match;
	Option *current;
	Option *selection;
} options = { 0 };

/* function definition */
static void calculate_scroll();
static void append_match(Option *item, Option **first, Option **last);
static void generate_matches();
static void draw_dmenu(Monitor *monitor);
static void setup();
static void readstdin();
static Monitor *add_monitor();
static void registry_global(void *data, struct wl_registry *wl_registry, uint32_t name, const char *interface, uint32_t version);
static void seat_capabilities(void *data, struct wl_seat *seat, uint32_t capabilities);
static void dwl_output_active(void *data, struct zdwl_output_v1 *output, uint32_t active);
static void wlr_layer_surface_configure(void *data, struct zwlr_layer_surface_v1 *surface, uint32_t serial, uint32_t w, uint32_t h);
static void keyboard_keymap(void *data, struct wl_keyboard *keyboard, uint32_t format, int32_t fd, uint32_t size);
static void keyboard_key(void *data, struct wl_keyboard *keyboard, uint32_t serial, uint32_t time, uint32_t key, uint32_t state);
static void keyboard_modifiers(void *data, struct wl_keyboard *keyboard, uint32_t serial, uint32_t depressed, uint32_t latched, uint32_t locked, uint32_t group);
static void keyboard_repeat_info(void *data, struct wl_keyboard *keyboard, int32_t rate, int32_t delay);
static void handle_keyboard_event();
static void dummy() {}

/* global variables */
static const struct wl_registry_listener registry_listnener = {
	.global = registry_global,
	.global_remove = dummy
};

static const struct zdwl_output_v1_listener dwl_output_listener = {
	.active = dwl_output_active,
	.appid = dummy,
	.frame = dummy,
	.layout = dummy,
	.tag = dummy,
	.title = dummy,
	.toggle_visibility = dummy,
};

static const struct zwlr_layer_surface_v1_listener wlr_layer_surface_listener = {
	.closed = dummy,
	.configure = wlr_layer_surface_configure,
};

static const struct wl_seat_listener seat_listener = {
	.capabilities = seat_capabilities,
	.name = dummy,
};

static const struct wl_keyboard_listener keyboard_listener = {
	.keymap = keyboard_keymap,
	.enter = dummy,
	.key = keyboard_key,
	.leave = dummy,
	.modifiers = keyboard_modifiers,
	.repeat_info = keyboard_repeat_info,
};

Client client = { 0 };
Monitor *monitors = 0;
Monitor *active_monitor = 0;
uint32_t height = 50;
uint32_t numitems = 0;
char input_field[MAX_LINE_LENGTH] = "";
char *prompt = 0;
int running = 1;

#include "config.h"

/* function implementations */
void
calculate_scroll()
{
	unsigned i = 0;
	Option *match = options.first_match;
	options.current = match;

	for (; match && match != options.selection; match = match->next_match) {
		if (i % lines == lines - 1)
			options.current = match->next_match;
		++i;
	}
}


void
append_match(Option *item, Option **first, Option **last)
{
	if (*last)
		(*last)->next_match = item;
	else
		*first = item;

	item->previous_match = *last;
	item->next_match = 0;
	*last = item;
}


void
generate_matches()
{
	static char **tokv = NULL;
	static int tokn = 0;

	char buf[sizeof(input_field)], *s;
	int i, tokc = 0;
	size_t len, textsize;
	Option *item, *lprefix, *lsubstr, *prefixend, *substrend;


	strcpy(buf, input_field);
	/* separate input text into tokens to be matched individually */
	for (s = strtok(buf, " "); s; tokv[tokc - 1] = s, s = strtok(0, " "))
		if (++tokc > tokn && !(tokv = realloc(tokv, ++tokn * sizeof *tokv)))
			die("cannot realloc bytes");

	len = tokc ? strlen(tokv[0]) : 0;

	options.first_match = lprefix = lsubstr = options.last_match = prefixend = substrend = 0;
	textsize = strlen(input_field) + 1;

	for (item = options.first; item && item->name; item = item->next) {
		for (i = 0; i < tokc; i++)
			if (!strstr(item->name, tokv[i]))
				break;
		if (i != tokc) /* not all tokens match */
			continue;
		/* exact matches go first, then prefixes, then substrings */
		if (!tokc || !strncmp(input_field, item->name, textsize))
			append_match(item, &options.first_match, &options.last_match);
		else if (!strncmp(tokv[0], item->name, len))
			append_match(item, &lprefix, &prefixend);
		else
			append_match(item, &lsubstr, &substrend);
	}

	if (lprefix) {
		if (options.first_match) {
			options.last_match->next_match = lprefix;
			lprefix->previous_match = options.last_match;
		} else
			options.first_match = lprefix;
		options.last_match = prefixend;
	}

	if (lsubstr) {
		if (options.first_match) {
			options.last_match->next_match = lsubstr;
			lsubstr->previous_match = options.last_match;
		} else
			options.first_match = lsubstr;
		options.last_match = substrend;
	}
	options.selection = options.first_match;
	calculate_scroll();
}


void
draw_dmenu(Monitor *monitor)
{
	Option *match = options.current;
	unsigned x = 0;
	unsigned y = 0;

	draw_rect(monitor->canvas, 0, 0, monitor->width, height, highlight);
	draw_rect(monitor->canvas, borderwidth, borderwidth,
	          monitor->width - 2 * borderwidth, height - 2 * borderwidth, background);
	draw_rect(monitor->canvas, 0, fontsize + borderwidth + 2 * padding,
	          monitor->width, borderwidth, highlight);

	if (prompt) {
		x = font_width(monitor->font, prompt) + padding + borderwidth;
		draw_rect(monitor->canvas, 0, 0, x + padding, fontsize + 2 * padding + borderwidth, highlight);
		draw_font(monitor->canvas, monitor->font, prompt, padding + borderwidth,
				borderwidth + fontsize + padding, foreground);
	}

	x = draw_font(monitor->canvas, monitor->font, input_field,
	              x + padding + borderwidth, borderwidth + fontsize + padding, foreground);
	draw_rect(monitor->canvas, x + cursor_horizontal_offset,
	          borderwidth + padding - cursor_vertical_offset, cursor_width,
	          fontsize + 2 * cursor_vertical_offset, foreground);

	y = 2 * borderwidth + fontsize + 2 * padding;

	for (; match && y < height; match = match->next_match) {
		if (match == options.selection)
			draw_rect(monitor->canvas, 0, y, monitor->width, fontsize + 2 * padding, highlight);
		y += fontsize + padding;
		draw_font(monitor->canvas, monitor->font, match->name, padding, y, foreground);
		y += padding;
	}


	wl_surface_attach(monitor->surface, (struct wl_buffer*)monitor->canvas->buffer, 0, 0);
	wl_surface_damage_buffer(monitor->surface, 0, 0, monitor->width, height);
	wl_surface_commit(monitor->surface);
}


Monitor *
add_monitor()
{
	Monitor *last = monitors;
	Monitor *mon = malloc(sizeof(Monitor));
	memset(mon, 0, sizeof(Monitor));

	if (monitors == 0) {
		monitors = mon;
	} else {
		for (; last->next; last = last->next);
		last->next = mon;
	}

	return mon;
}

void
registry_global(void *data, struct wl_registry *wl_registry, uint32_t name, const char *interface, uint32_t version)
{
	struct wl_output *output;
	Monitor *monitor;

	match_then_bind(client.shm, wl_shm_interface, 1)
	or_match match_then_bind(client.compositor, wl_compositor_interface, 4)
	or_match match_then_bind(client.layer, zwlr_layer_shell_v1_interface, 1)
	or_match match_then_bind(client.dwl_manager, zdwl_manager_v1_interface, 1)
	or_match match_then_bind(client.seat, wl_seat_interface, 7)
		wl_seat_add_listener(client.seat, &seat_listener, &client);
	or_match match_then_bind(output, wl_output_interface, 1)
		monitor = add_monitor();
		monitor->output = output;
	end_match
}


void
seat_capabilities(void *data, struct wl_seat *seat, uint32_t capabilities)
{
	int has_keyboard = capabilities & WL_SEAT_CAPABILITY_KEYBOARD;
	
	if (has_keyboard && client.keyboard == NULL) {
		client.keyboard = wl_seat_get_keyboard(client.seat);
		wl_keyboard_add_listener(client.keyboard, &keyboard_listener, &client);
	} else {
		wl_keyboard_release(client.keyboard);
		client.keyboard = NULL;
	}
}


void
dwl_output_active(void *data, struct zdwl_output_v1 *output, uint32_t active)
{
	Monitor *monitor = data;

	if (active && active_monitor == 0) {
		active_monitor = data;
	}
}


void
wlr_layer_surface_configure(void *data, struct zwlr_layer_surface_v1 *surface, uint32_t serial, uint32_t w, uint32_t h)
{
	Monitor *mon = data;
	mon->width = w;

	zwlr_layer_surface_v1_ack_configure(surface, serial);
	
	if (mon->canvas == 0)
		mon->canvas = create_drw(client.shm, mon->width, height);
	draw_dmenu(mon);
}


void
keyboard_keymap(void *data, struct wl_keyboard *keyboard, uint32_t format, int32_t fd, uint32_t size)
{
	char *key_shm;
	struct xkb_keymap *keymap;
	struct xkb_state *state;

	if (format != WL_KEYBOARD_KEYMAP_FORMAT_XKB_V1)
		return;

	key_shm = mmap(NULL, size, PROT_READ, MAP_SHARED, fd, 0);

	if (key_shm == MAP_FAILED)
		return;

	keymap = xkb_keymap_new_from_string(
		client.xkb_context, key_shm,
		XKB_KEYMAP_FORMAT_TEXT_V1, XKB_KEYMAP_COMPILE_NO_FLAGS
	);

	munmap(key_shm, size);
	close(fd);

	state = xkb_state_new(keymap);
	xkb_keymap_unref(client.xkb_keymap);
	xkb_state_unref(client.xkb_state);

	client.xkb_state = state;
	client.xkb_keymap = keymap;
}


void
keyboard_key(void *data, struct wl_keyboard *keyboard, uint32_t serial, uint32_t time, uint32_t key, uint32_t state)
{
	struct itimerspec spec = { 0 };
	uint32_t keycode = key + 8;
	xkb_keysym_t sym = xkb_state_key_get_one_sym(client.xkb_state, keycode);
	client.keys.sym = sym;
	client.keys.state = state;
	handle_keyboard_event();


	if (client.keys.state == WL_KEYBOARD_KEY_STATE_PRESSED && client.repeat.period >= 0) {
		spec.it_value.tv_sec = client.repeat.delay / 1000;
		spec.it_value.tv_nsec = (client.repeat.delay % 1000) * 1000000;
		timerfd_settime(client.repeat.timer, 0, &spec, 0);
	} else if (client.keys.state == WL_KEYBOARD_KEY_STATE_RELEASED) {
		timerfd_settime(client.repeat.timer, 0, &spec, 0);
	}
}


void
keyboard_modifiers(void *data, struct wl_keyboard *keyboard, uint32_t serial, uint32_t depressed, uint32_t latched, uint32_t locked, uint32_t group)
{
	xkb_state_update_mask(client.xkb_state, depressed, latched, locked, 0, 0, group);
	client.keys.alt = xkb_state_mod_name_is_active(client.xkb_state, XKB_MOD_NAME_ALT, XKB_STATE_MODS_DEPRESSED | XKB_STATE_MODS_LATCHED);
	client.keys.control = xkb_state_mod_name_is_active(client.xkb_state, XKB_MOD_NAME_CTRL, XKB_STATE_MODS_DEPRESSED | XKB_STATE_MODS_LATCHED);
	client.keys.shift = xkb_state_mod_name_is_active(client.xkb_state, XKB_MOD_NAME_SHIFT, XKB_STATE_MODS_DEPRESSED | XKB_STATE_MODS_LATCHED);
}


void
keyboard_repeat_info(void *data, struct wl_keyboard *keyboard, int32_t rate, int32_t delay)
{
	client.repeat.delay = delay;

	if (rate > 0)
		client.repeat.period = 1000 / rate;
	else
		client.repeat.period = -1;
}


void
handle_keyboard_event()
{
	char buf[8];

	if (client.keys.state == WL_KEYBOARD_KEY_STATE_PRESSED) {
		if (client.keys.alt) {
			switch (client.keys.sym) {
				case XKB_KEY_j: 
					if (options.selection->next_match)
						options.selection = options.selection->next_match;
					calculate_scroll();
				break;
				case XKB_KEY_k: 
					if (options.selection->previous_match)
						options.selection = options.selection->previous_match;
					calculate_scroll();
				break;
				case XKB_KEY_g:
					options.selection = options.first_match;
					calculate_scroll();
					break;
				case XKB_KEY_G:
					options.selection = options.last_match;
					calculate_scroll();
					break;
				default: break;
			}
		} else {
			switch (client.keys.sym) {
				case XKB_KEY_Down:
					if (options.selection->next_match)
						options.selection = options.selection->next_match;
					calculate_scroll();
					break;
				case XKB_KEY_Up:
					if (options.selection->previous_match)
						options.selection = options.selection->previous_match;
					calculate_scroll();
					break;
				case XKB_KEY_Left:
					break;
				case XKB_KEY_Right:
					break;
				case XKB_KEY_KP_Enter: /* fallthrough */
				case XKB_KEY_Return:
					fputs(options.selection ? options.selection->name: input_field, stdout);
					running = 0;
					break;
				case XKB_KEY_Escape:
					running = 0;
					break;
				case XKB_KEY_Tab:
					if (options.selection)
						strncpy(input_field, options.selection->name, MAX_LINE_LENGTH);
					break;
				case XKB_KEY_BackSpace:
					if (strlen(input_field))
						input_field[strlen(input_field) - 1] = 0;
					generate_matches();
					break;
				default:
					if (xkb_keysym_to_utf8(client.keys.sym, buf, sizeof(buf))) {
						strncpy(input_field + strlen(input_field), buf, MAX_LINE_LENGTH - strlen(input_field));
						generate_matches();
					}
					break;
			}
		}
		
		if (running) {
			draw_dmenu(active_monitor);
		}
	}
}


void
readstdin()
{
	unsigned s;
	char *buf = malloc(MAX_LINE_LENGTH);
	Option *option;

	for (;fgets(buf, MAX_LINE_LENGTH, stdin) != 0; ++numitems) {
		option = malloc(sizeof(Option));
		option->name = buf;
		s = strlen(option->name);
		if (option->name[s - 1] == '\n')
			option->name[s - 1] = 0;
		option->next = 0;
		option->next_match = 0;
		option->previous_match = 0;

		if (options.first == 0) {
			options.first = option;
			options.last = option;
			option->previous = 0;
		} else {
			options.last->next = option;
			option->previous = options.last;
			options.last = option;
		}

		buf = malloc(MAX_LINE_LENGTH);
	}

	free(buf);
}


void
setup()
{
	unsigned count = 0;
	Monitor *mon;
	char *compositor;
	char namespace[] = "dmenu-wl";
	uint32_t layer = ZWLR_LAYER_SHELL_V1_LAYER_TOP;
	uint32_t anchor = ZWLR_LAYER_SURFACE_V1_ANCHOR_TOP | ZWLR_LAYER_SURFACE_V1_ANCHOR_LEFT | ZWLR_LAYER_SURFACE_V1_ANCHOR_RIGHT;

	if (numitems < lines)
		lines = numitems;

	height = (lines + 1) * (fontsize + padding * 2) + 3 * borderwidth;
	generate_matches();
	calculate_scroll();

	client.repeat.timer = timerfd_create(CLOCK_MONOTONIC, 0);
	client.xkb_context = xkb_context_new(XKB_CONTEXT_NO_FLAGS);
	client.display = wl_display_connect(0);

	if (client.display == 0)
		die("could not open display:");

	client.registry = wl_display_get_registry(client.display);

	if (client.registry == 0)
		die("could not open display:");

	wl_registry_add_listener(client.registry, &registry_listnener, &client);
	wl_display_roundtrip(client.display);

	compositor = process_name_from_fd(client.display);
	if (strcmp(compositor, "dwl") == 0) {
		for (mon = monitors; mon; mon = mon->next) {
			mon->dwl_output = zdwl_manager_v1_get_output(client.dwl_manager, mon->output);
			zdwl_output_v1_add_listener(mon->dwl_output, &dwl_output_listener, mon);
		}
		wl_display_roundtrip(client.display);
	} else if (strcmp(compositor, "Hyprland") == 0) {
		char nums[3] = "";
		FILE *nfd = popen("hyprctl monitors | grep '\\(Monitor\\|focused\\)' | paste -d ' ' - - | grep -n yes | grep -o '^[0-9]*'", "r");

		fgets(nums, 3, nfd);

		if (nums[strlen(nums) - 1] == '\n') {
			nums[strlen(nums) - 1] = 0;
		}

		monitor_offset += atoi(nums) - 1;

		pclose(nfd);
	}

	for (mon = monitors; mon && count < monitor_offset;) {
		mon = mon->next;
		count += 1;
	}

	if (mon) {
		active_monitor = mon;
	} else {
		active_monitor = monitors;
	}

	active_monitor->font = create_font(fontpath, fontsize * 1.5);

	active_monitor->surface = wl_compositor_create_surface(client.compositor);
	active_monitor->wlr_surface = zwlr_layer_shell_v1_get_layer_surface(client.layer, active_monitor->surface, active_monitor->output, layer, namespace);

	zwlr_layer_surface_v1_set_exclusive_zone(active_monitor->wlr_surface, -1);
	zwlr_layer_surface_v1_set_keyboard_interactivity(active_monitor->wlr_surface, 1);

	zwlr_layer_surface_v1_add_listener(active_monitor->wlr_surface, &wlr_layer_surface_listener, active_monitor);

	zwlr_layer_surface_v1_set_size(active_monitor->wlr_surface, 0, height);
	zwlr_layer_surface_v1_set_anchor(active_monitor->wlr_surface, anchor);

	wl_surface_commit(active_monitor->surface);
	wl_display_roundtrip(client.display);
}


int
main(int argc, char *argv[])
{
	struct itimerspec spec = { 0 };
	struct pollfd fds[2];
	char **arg;

	for (arg = argv; *arg; ++arg) {
		if (strcmp(*arg, "-p") == 0) {
			prompt = *(++arg);
		} else if (strcmp(*arg, "-h") == 0) {
			fprintf(stderr, "dmenu-wl [-p PROMPT]\n");
		}
	}

	readstdin();
	setup();

	fds[0].fd = wl_display_get_fd(client.display);
	fds[0].events = POLLIN;
	fds[1].fd = client.repeat.timer;
	fds[1].events = POLLIN;

	while (running) {
		if (wl_display_flush(client.display) < 0) {
			if (errno == EAGAIN)
				continue;
			break;
		}

		if (poll(fds, sizeof(fds) / sizeof(*fds), -1) < 0) {
			if (errno == EAGAIN)
				continue;
			break;
		}

		if (fds[0].revents & POLLIN) {
			if (wl_display_dispatch(client.display) < 0) {
				running = 0;
			}
		}

		if (fds[1].revents & POLLIN) {
			handle_keyboard_event();
			spec.it_value.tv_sec = client.repeat.period / 1000;
			spec.it_value.tv_nsec = (client.repeat.period % 1000) * 1000000;
			timerfd_settime(client.repeat.timer, 0, &spec, 0);
		}
	}

	wl_display_disconnect(client.display);
}