#include <stdio.h>
#include <string.h>
#include <sys/mman.h>
#include <stdlib.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 <sys/time.h>
#include <errno.h>

#include <pwd.h>
#include <grp.h>
#include <errno.h>
#include <stdio.h>
#include <string.h>

#include "util.h"
#include "drw.h"
#include "config.h"
#include "oom.h"
#include "hash.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 sym;
		uint32_t state;
	} keys;
	struct {
		int timer;
		int delay;
		int period;
	} repeat;
} Client;

typedef struct Monitor Monitor;
struct Monitor {
	uint32_t width;
	uint32_t height;
	uint32_t registry_name;
	Canvas *canvas;
	struct wl_output *output;
	struct wl_surface *surface;
	struct zwlr_layer_surface_v1 *wlr_surface;
	Monitor *next;
};

/* function definition */
static void draw_to_monitor(Monitor *mon);
static void update_time();
static void setup();
static Monitor *add_monitor();
static void remove_monitor(Monitor *mon);
static void registry_global(void *data, struct wl_registry *wl_registry, uint32_t name, const char *interface, uint32_t version);
static void remove_global(void *data, struct wl_registry *wl_registry, uint32_t name);
static void seat_capabilities(void *data, struct wl_seat *seat, uint32_t capabilities);
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_repeat_info(void *data, struct wl_keyboard *keyboard, int32_t rate, int32_t delay);
static void handle_keyboard_event();
static void register_monitor(Monitor *mon);
static void dummy() {}

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

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 = dummy,
	.repeat_info = keyboard_repeat_info,
};

Client client = { 0 };
Monitor *monitors = 0;
char input_field[MAX_LINE_LENGTH] = "";
int running = 1;
const char *hash;
const char *username;
char timestr[64];
Font *username_font;
Font *clock_font;
int ready = 0;

/* function implementations */
void
draw_to_monitor(Monitor *mon)
{
	int inputlen = strlen(input_field);
	int x;

	draw_rect(mon->canvas, 0, 0, mon->width, mon->height, background);
	draw_rect(mon->canvas, (mon->width - inputlen * 20) / 2, mon->height / 2, inputlen * 20, 10, inputbar);

	x = font_width(username_font, username);
	draw_font(mon->canvas, username_font, username, (mon->width - x) / 2, mon->height / 2 - 2 * clock_fontsize - 20, foreground);

	x = font_width(clock_font, timestr);
	draw_font(mon->canvas, clock_font, timestr, (mon->width - x) / 2, mon->height / 2 - clock_fontsize , foreground);

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


void
update_time()
{
	Monitor *mon;
	time_t rawtime;
	struct tm *timeinfo;

	time(&rawtime);
	timeinfo = localtime(&rawtime);
	sprintf(timestr, "%02d:%02d", timeinfo->tm_hour, timeinfo->tm_min);

	for (mon = monitors; mon; mon = mon->next)
		draw_to_monitor(mon);
}


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
remove_monitor(Monitor *mon)
{
	wl_surface_destroy(mon->surface);
	zwlr_layer_surface_v1_destroy(mon->wlr_surface);
	wl_output_destroy(mon->output);
	free_drw(mon->canvas);
	free(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.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;
		monitor->registry_name = name;
		if (ready) {
			register_monitor(monitor);
		}
	end_match
}


void
remove_global(void *data, struct wl_registry *wl_registry, uint32_t name)
{
	Monitor *mon = monitors;
	Monitor *prev = 0;
	for (; mon; mon = mon->next) {
		if (mon->registry_name == name) {
			if (prev) {
				prev->next = mon->next;
			} else {
				monitors = mon->next;
			}

			remove_monitor(mon);
			break;
		}

		prev = mon;
	}
}


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
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;
	mon->height = h;

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

		if (mon->canvas == 0) {
			die("wlock: could not create canvas:", strerror(errno));
		}
	}

	draw_to_monitor(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_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];
	char *inputhash;
	Monitor *mon;

	if (client.keys.state == WL_KEYBOARD_KEY_STATE_PRESSED) {
		switch (client.keys.sym) {
			case XKB_KEY_Escape:
				memset(input_field, 0, MAX_LINE_LENGTH);
				break;
			case XKB_KEY_KP_Enter: /* fallthrough */
			case XKB_KEY_Return:
				if (!(inputhash = crypt(input_field, hash)))
					fprintf(stderr, "slock: crypt: %s\n", strerror(errno));
				else
					running = !!strcmp(inputhash, hash);
				memset(input_field, 0, MAX_LINE_LENGTH);
				break;
			case XKB_KEY_BackSpace:
				if (strlen(input_field))
					input_field[strlen(input_field) - 1] = 0;
				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));
				}
				break;
		}
		
		if (running) {
			for (mon = monitors; mon; mon = mon->next)
				draw_to_monitor(mon);
		}
	}
}


void
register_monitor(Monitor *mon)
{
	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_BOTTOM |
	                  ZWLR_LAYER_SURFACE_V1_ANCHOR_LEFT |
	                  ZWLR_LAYER_SURFACE_V1_ANCHOR_RIGHT;

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

	zwlr_layer_surface_v1_set_exclusive_zone(mon->wlr_surface, -1);
	zwlr_layer_surface_v1_set_keyboard_interactivity(monitors->wlr_surface, 1);

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

	zwlr_layer_surface_v1_set_anchor(mon->wlr_surface, anchor);

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


void
setup()
{
	Monitor *mon;
	username_font = create_font(fontpath, username_fontsize * 1.5);
	clock_font = create_font(fontpath, clock_fontsize * 1.5);
	client.repeat.timer = timerfd_create(CLOCK_MONOTONIC, 0);
	client.xkb_context = xkb_context_new(XKB_CONTEXT_NO_FLAGS);
	client.display = wl_display_connect(0);
	client.registry = wl_display_get_registry(client.display);
	wl_registry_add_listener(client.registry, &registry_listnener, &client);
	wl_display_roundtrip(client.display);

	for (mon = monitors; mon; mon = mon->next) {
		register_monitor(mon);
	}

	ready = 1;
}


int
main(int argc, char *argv[])
{
	struct itimerspec spec = { 0 };
	struct pollfd fds[3];
	struct passwd *pw;
	struct group *grp;

	username = getenv("USER");

	if (!(pw = getpwnam(user)))
		die("slock: getpwname: user entry not found");
	if (!(grp = getgrnam(group)))
		die("slock: getpwname: user entry not found");

	hash = get_hash();
	if (!crypt("", hash))
		die("slock: crypt: %s\n", strerror(errno));

#ifdef __linux__
	disable_oom_killer();
#endif

	setup();


	if (setgroups(0, NULL) < 0)
		die("slock: setgroups: %s\n", strerror(errno));
	if (setgid(grp->gr_gid) < 0)
		die("slock: setgid: %s\n", strerror(errno));
	if (setuid(pw->pw_uid) < 0)
		die("slock: setuid: %s\n", strerror(errno));

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

	spec.it_interval.tv_sec = 1;
	spec.it_value.tv_nsec = 1;
	timerfd_settime(fds[2].fd, 0, &spec, 0);
	spec.it_interval.tv_nsec = 0;

	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);
		}

		if (fds[2].revents & POLLIN) {
			update_time();
		}
	}

	wl_display_dispatch(client.display);
}