#include <iostream>
#include <memory>
#include <vector>
#include <unordered_map>
#include <variant>
#include <stdexcept>
#include <functional>
#include <cassert>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/XKBlib.h>


class X_keycode {
public:
	X_keycode() : code{0} {}
	X_keycode(unsigned int code) : code{code} {}
	explicit operator unsigned int() const { return code; }
	bool operator==(X_keycode other) const { return code == other.code; }
private: unsigned int code;
};

template <>
struct std::hash<X_keycode> {
	size_t operator()(X_keycode code) const {
		return std::hash<unsigned int>{}((unsigned int)code);
	}
};

class X_keysym {
public:
	X_keysym() : sym{0} {}
	X_keysym(unsigned int sym) : sym{sym} {}
	explicit operator unsigned int() const { return sym; }
	X_keycode toCode(Display *dpy) const { return XKeysymToKeycode(dpy, sym); }
	bool operator==(X_keysym other) const { return sym == other.sym; }
private: unsigned int sym;
};

template <>
struct std::hash<X_keysym> {
	size_t operator()(X_keysym sym) const {
		return std::hash<unsigned int>{}((unsigned int)sym);
	}
};

void bel(Display *dpy) {
	XkbBell(dpy, None, 100, None);
}

template <typename Cleanup>
class UponExit {
public:
	UponExit(Cleanup cleanup) : cleanup{cleanup} {}
	~UponExit() {
		if (cleanup) (*cleanup)();
	}
	UponExit(const UponExit&) = delete;
	UponExit(UponExit &&other) : cleanup{move(other.cleanup)} {
		other.cleanup.reset();
	}
	UponExit& operator=(const UponExit&) = delete;
	UponExit& operator=(UponExit &&other) {
		cleanup = move(other.cleanup);
		other.cleanup.reset();
	}

private:
	std::optional<Cleanup> cleanup;
};

auto XGrabKeyRAII(Display *dpy, X_keycode code, int modifier, Window win) {
	XGrabKey(dpy, (unsigned int)code, modifier, win, False, GrabModeAsync, GrabModeAsync);
	return UponExit{[dpy, code, modifier, win]() {
		XUngrabKey(dpy, (unsigned int)code, modifier, win);
	}};
}

auto XGrabKeyboardRAII(Display *dpy, Window win) {
	int ret = XGrabKeyboard(dpy, win, False, GrabModeAsync, GrabModeAsync, CurrentTime);
	if (ret == AlreadyGrabbed) {
		XUngrabKeyboard(dpy, CurrentTime);
		throw std::runtime_error("Cannot grab keyboard: already grabbed");
	}
	return UponExit{[dpy]() {
		XUngrabKeyboard(dpy, CurrentTime);
	}};
}

auto XOpenDisplayRAII(const char *name) {
	Display *dpy = XOpenDisplay(name);
	if (dpy == nullptr) {
		std::cerr << "Cannot open X display" << std::endl;
		exit(1);
	}
	return std::make_pair(dpy, UponExit{[dpy]() {
		XCloseDisplay(dpy);
	}});
}

template <typename F>  // return true to stop watch and loop
void globalKeyWatch(Display *dpy, X_keysym headerSym, F callback) {
	const Window root = DefaultRootWindow(dpy);
	const X_keycode headerCode = headerSym.toCode(dpy);

	auto guard = XGrabKeyRAII(dpy, headerCode, AnyModifier, root);

	XSelectInput(dpy, root, KeyPressMask);
	while (true) {
		XEvent ev;
		XNextEvent(dpy, &ev);
		if (ev.type == KeyPress && ev.xkey.keycode == (unsigned int)headerCode) {
			if (callback(ev.xkey)) return;
		}
	}
}

template <typename F>  // return true to lose grab and loop
void globalKeyboardGrab(Display *dpy, F callback) {
	const Window root = DefaultRootWindow(dpy);

	try {
		auto guard = XGrabKeyboardRAII(dpy, root);

		while (true) {
			XEvent ev;
			XNextEvent(dpy, &ev);
			if (ev.type == KeyPress) {
				if (callback(ev.xkey)) return;
			}
		}
	} catch (std::exception &e) {
		std::cerr << e.what() << std::endl;
	}
}

class SeqMatcher {
public:
	using Callback = std::function<void()>;

	struct SymSequence {
		std::vector<X_keysym> syms;
		Callback callback;
	};

	SeqMatcher(Display *dpy) : dpy{dpy} {}

	SeqMatcher(Display *dpy, std::vector<SymSequence> seqs)
			: dpy{dpy} {
		for (const auto &seq : seqs) addSequence(seq.syms, seq.callback);
	}

	void addSequence(const std::vector<X_keysym> &syms, Callback callback) {
		if (syms.empty()) {
			throw std::logic_error("Cannot register empty key sequence");
		}

		Node *current = &rootNode;
		for (X_keysym sym : syms) {
			if (std::holds_alternative<Callback>(current->v)) {
				throw std::logic_error("Overlapping key sequences (second is longer)");
			} else {
				if (!std::holds_alternative<NodeMap>(current->v)) {
					current->v.emplace<NodeMap>();
				}
				NodeMap &map = std::get<NodeMap>(current->v);
				X_keycode code = sym.toCode(dpy);
				auto it = map.find(code);
				if (it != map.end()) {
					current = it->second.get();
				} else {
					current = map.emplace(sym.toCode(dpy), std::make_unique<Node>()).first->second.get();
				}
			}
		}

		if (auto *map = std::get_if<NodeMap>(&current->v)) {
			if (!map->empty()) {
				throw std::logic_error("Overlapping key sequences (second is shorter)");
			}
		}
		if (std::holds_alternative<Callback>(current->v)) {
			throw std::logic_error("Overlapping key sequences (equally long)");
		}
		current->v.emplace<Callback>(callback);
	}

	// Returns a sequence was completed, either successfully or erroneously.
	// Needs more events to complete a sequence iff it returns false.
	bool observe(const XKeyEvent &ev) {
		auto *map = std::get_if<NodeMap>(&curNode->v);
		assert(map);

		auto it = map->find(X_keycode{ev.keycode});
		if (it == map->end()) {
			// Sequence not found
			bel(dpy);
			reset();
			return true;
		}

		curNode = it->second.get();

		if (auto *cb = std::get_if<Callback>(&curNode->v)) {
			(*cb)();
			reset();
			return true;
		}

		// Need more keys
		return false;
	}

	void reset() {
		curNode = &rootNode;
	}

private:
	struct Node;
	using NodeMap = std::unordered_map<X_keycode, std::unique_ptr<Node>>;
	struct Node {
		std::variant<NodeMap, Callback> v;
	};

	Display *const dpy;
	Node rootNode;
	Node *curNode = &rootNode;
};

int main() {
	auto dpy_pair = XOpenDisplayRAII(nullptr);
	Display *dpy = dpy_pair.first;

	bool quitRequested = false;

	SeqMatcher matcher{dpy};
	matcher.addSequence({XK_A, XK_B}, []() {
		std::cout << "key ab" << std::endl;
	});
	matcher.addSequence({XK_B}, []() {
		std::cout << "key b" << std::endl;
	});
	matcher.addSequence({XK_Q}, [&quitRequested]() {
		quitRequested = true;
	});

	globalKeyWatch(
		dpy, XK_Break,
		[dpy, &matcher, &quitRequested](const XKeyEvent&) -> bool {
			matcher.reset();
			globalKeyboardGrab(dpy, [&matcher](const XKeyEvent &ev) -> bool {
				return matcher.observe(ev);
			});
			return quitRequested;
		}
	);
}