#include <iostream>
#include <cassert>
#include "KeyboardHook.h"

using namespace std;

typedef KBDLLHOOKSTRUCT KeyboardHookStruct;
LRESULT CALLBACK keyboard_hook_callback(int, WPARAM, LPARAM);
std::map<thread::id, KeyboardHook*> hook_handles;

KeyboardHook::KeyboardHook() {}
KeyboardHook::~KeyboardHook() {
	if(this->_attached)
		this->detach();
}

bool KeyboardHook::attach() {
	assert(!this->_attached);
	this->active = true;

	this->poll_thread = std::thread(std::bind(&KeyboardHook::poll_events, this));

	this->_attached = true;
	return true;
}

void KeyboardHook::detach() {
	assert(this->_attached);
	this->active = false;
	{
		//TODO trigger no message!
	}
	if(this->poll_thread.joinable())
		this->poll_thread.join();

	this->_attached = false;
}

void KeyboardHook::poll_events() {
	hook_handles[this_thread::get_id()] = this;
	this->hook_id = SetWindowsHookEx(WH_KEYBOARD_LL, keyboard_hook_callback, GetModuleHandle(nullptr), 0);
	if(!this->hook_id) {
		cerr << "Failed to register hook!" << endl;
		return;
	}

	MSG msg;
	while(!GetMessage(&msg, nullptr, 0, 0) && this->active) {
		TranslateMessage(&msg);
		DispatchMessage(&msg);
	}

	UnhookWindowsHookEx(this->hook_id);
	hook_handles[this_thread::get_id()] = nullptr;
}

LRESULT keyboard_hook_callback(int nCode, WPARAM event, LPARAM ptr_keyboard) {
	auto handle = hook_handles[this_thread::get_id()];
	assert(handle);
	auto consume = handle->_hook_callback(nCode, event, ptr_keyboard);
	if(consume)
		return 1;
	return CallNextHookEx(nullptr, nCode, event, ptr_keyboard);
}

inline std::string key_code(DWORD code, bool extended = false) {
	auto scan_code = MapVirtualKey(code, MAPVK_VK_TO_VSC);
	if(extended)
		scan_code |= KF_EXTENDED;

	char key_buffer[255];
	auto length = GetKeyNameTextA(scan_code << 16, key_buffer, 255);
	if(length == 0)
		return "error";
	else
		return string(key_buffer, length);
}

inline std::string key_code(KeyboardHookStruct* keyboard) {
	return key_code(keyboard->vkCode, (keyboard->flags & LLKHF_EXTENDED) > 0);
}

using KeyType = KeyboardHook::KeyType;
//https://docs.microsoft.com/en-us/windows/desktop/inputdev/virtual-key-codes
inline KeyType::value key_type(DWORD vk_code) {
	switch(vk_code) {
		case VK_CONTROL:
		case VK_LCONTROL:
		case VK_RCONTROL:
			return KeyType::KEY_CTRL;
		case VK_MENU:
		case VK_RMENU:
		case VK_LMENU:
			return KeyType::KEY_ALT;
		case VK_SHIFT:
		case VK_RSHIFT:
		case VK_LSHIFT:
			return KeyType::KEY_SHIFT;
		case VK_LWIN:
		case VK_RWIN:
			return KeyType::KEY_WIN;
		default:
			return KeyType::KEY_NORMAL;
	}
}

bool KeyboardHook::_hook_callback(int nCode, WPARAM event, LPARAM ptr_keyboard) {
	auto keyboard = (KeyboardHookStruct*) ptr_keyboard;
	if(event == WM_KEYDOWN || event == WM_SYSKEYDOWN) {
		auto& state = this->map_key[keyboard->vkCode];
		bool typed = state;
		state = true;

		auto type = key_type(keyboard->vkCode);
		if(type != KeyType::KEY_NORMAL)
			this->map_special[type] = true;
		else
			this->map_special[type] = keyboard->vkCode;

		if (this->callback_event) {
			if(!typed) {
				auto e = make_shared<KeyboardHook::KeyEvent>();
				e->type = KeyboardHook::KeyEvent::PRESS;
				if(type == KeyType::KEY_NORMAL) {
					e->code = key_code(keyboard);
				} else {
					e->code = key_code(this->map_special[KeyType::KEY_NORMAL], false);
				}

				e->key_alt = this->map_special[KeyType::KEY_ALT];
				e->key_ctrl = this->map_special[KeyType::KEY_CTRL];
				e->key_windows = this->map_special[KeyType::KEY_WIN];
				e->key_shift = this->map_special[KeyType::KEY_SHIFT];

				this->callback_event(e);
			}
			{
				auto e = make_shared<KeyboardHook::KeyEvent>();
				e->type = KeyboardHook::KeyEvent::TYPE;
				if(type == KeyType::KEY_NORMAL) {
					e->code = key_code(keyboard);
				} else {
					e->code = key_code(this->map_special[KeyType::KEY_NORMAL], false);
				}

				e->key_alt = this->map_special[KeyType::KEY_ALT];
				e->key_ctrl = this->map_special[KeyType::KEY_CTRL];
				e->key_windows = this->map_special[KeyType::KEY_WIN];
				e->key_shift = this->map_special[KeyType::KEY_SHIFT];

				this->callback_event(e);
			}
		}
	} else if(event == WM_KEYUP || event == WM_SYSKEYUP) {
		auto& state = this->map_key[keyboard->vkCode];
		if(!state) return false; //Duplicate
		state = false;

		auto type = key_type(keyboard->vkCode);
		if(type != KeyType::KEY_NORMAL)
			this->map_special[type] = false;
		else if(this->map_special[KeyType::KEY_NORMAL] == keyboard->vkCode)
			this->map_special[KeyType::KEY_NORMAL] = 0xFF;

		if (this->callback_event) {
			auto e = make_shared<KeyboardHook::KeyEvent>();
			e->type = KeyboardHook::KeyEvent::RELEASE;
			if(type == KeyType::KEY_NORMAL) {
				e->code = key_code(keyboard);
			} else {
				e->code = key_code(this->map_special[KeyType::KEY_NORMAL], false);
			}

			e->key_alt = this->map_special[KeyType::KEY_ALT];
			e->key_ctrl = this->map_special[KeyType::KEY_CTRL];
			e->key_windows = this->map_special[KeyType::KEY_WIN];
			e->key_shift = this->map_special[KeyType::KEY_SHIFT];

			this->callback_event(e);
		}
	}

	//Consume the event: return 1
	return false;
}