#include <errno.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <X11/Xlib.h>
#include <X11/keysym.h>

struct named_mask {
	const char *name;
	unsigned int mask;
};

static const struct named_mask NAMED_MASKS[] = {
	{ "Shift", ShiftMask },
	{ "Control", ControlMask },
	{ "Ctrl", ControlMask },
	{ "Alt", Mod1Mask },
	{ "Super", Mod4Mask },
	{ "Mod1", Mod1Mask },
	{ "Mod2", Mod2Mask },
	{ "Mod3", Mod3Mask },
	{ "Mod4", Mod4Mask },
	{ "Mod5", Mod5Mask }
};
#define NUM_NAMED_MASKS (sizeof(NAMED_MASKS) / sizeof(*NAMED_MASKS))

static unsigned int name_to_mask(const char *s)
{
	for (const struct named_mask *m = NAMED_MASKS;
			m < NAMED_MASKS + NUM_NAMED_MASKS; ++m) {
		if (!strcmp(s, m->name))
			return m->mask;
	}

	return 0;
}

static void die(const char *fmt, ...)
{
	va_list ap;
	va_start(ap, fmt);

	char *s;
	if (vasprintf(&s, fmt, ap) != -1)
		fprintf(stderr, "xkbdlock: %s\n", s);
	else
		fprintf(stderr, "xkbdlock: %s\n", strerror(errno));

	va_end(ap);
	exit(1);
}

static bool grab_kbd(Display *d, int screen)
{
	Window w = RootWindow(d, screen);
	int res;

	// attempt to grab keyboard for 500ms if already grabbed.
	// thanks to slock code for the pointer.
	for (int i = 0;;) {
		res = XGrabKeyboard(d, w, True, GrabModeAsync, GrabModeAsync,
			CurrentTime);
		if (res != AlreadyGrabbed || ++i > 5)
			break;

		usleep(100000);
	}

	return (res == GrabSuccess);
}

static void wait_for_release(Display *d, KeySym key, unsigned int mask)
{
	XEvent ev;
	while (!XNextEvent(d, &ev)) {
		if (ev.type == KeyPress &&
				XLookupKeysym(&ev.xkey, 0) == key &&
				(mask & ev.xkey.state) == mask)
			break;
	}
}

int main(int argc, char **argv)
{
	unsigned int mask = 0, arg_mask;
	KeySym key = NoSymbol, arg_key;

	for (int i = 1; i < argc; ++i) {
		if ((arg_mask = name_to_mask(argv[i]))) {
			mask = mask | arg_mask;
		} else if (key != NoSymbol) {
			die("more than one non-modifier key provided");
		} else {
			if ((arg_key = XStringToKeysym(argv[i])) == NoSymbol)
				die("unknown key '%s'", argv[i]);
			if (arg_key >= XK_A && arg_key <= XK_Z)
				die("use Shift modifier for uppercase key '%s'",
					argv[i]);
			key = arg_key;
		}
	}

	if (key == NoSymbol) {
		if (mask)
			die("no non-modifier key provided");
		fprintf(stderr,
			"xkbdlock: no release key provided, using Escape\n");
		key = XK_Escape;
	}

	Display *d = XOpenDisplay(NULL);
	if (!d)
		die("cannot open display '%s'", getenv("DISPLAY"));

	for (int i = 0; i < ScreenCount(d); ++i) {
		if (!grab_kbd(d, i))
			die("failed to grab keyboard on screen %d", i);
	}

	wait_for_release(d, key, mask);
	return 0;
}