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

using namespace std;

typedef KBDLLHOOKSTRUCT KeyboardHookStruct;
typedef MSLLHOOKSTRUCT MouseHookStruct;
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;
}

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

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

void KeyboardHook::poll_events() {
	hook_handles[this_thread::get_id()] = this;
	this->keyboad_hook_id = SetWindowsHookEx(WH_KEYBOARD_LL, _keyboard_hook_callback, GetModuleHandle(nullptr), 0);
	if(!this->keyboad_hook_id) {
		cerr << "Failed to register keyboard hook" << endl;
		return;
	}
	
    this->mouse_hook_id = SetWindowsHookEx(WH_MOUSE_LL, _mouse_hook_callback, GetModuleHandle(nullptr), 0);
    if(!this->keyboad_hook_id) {
        UnhookWindowsHookEx(this->keyboad_hook_id);
        cerr << "Failed to register mouse hook" << endl;
        return;
    }

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

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

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::keyboard_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(!typed)
            this->trigger_key_event(KeyEvent::PRESS, type == KeyType::KEY_NORMAL ? key_code(keyboard) : key_code(this->map_special[KeyType::KEY_NORMAL], false));
        this->trigger_key_event(KeyEvent::TYPE, type == KeyType::KEY_NORMAL ? key_code(keyboard) : key_code(this->map_special[KeyType::KEY_NORMAL], false));
	} 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] = 0;

        this->trigger_key_event(KeyEvent::RELEASE, type == KeyType::KEY_NORMAL ? key_code(keyboard) : key_code(this->map_special[KeyType::KEY_NORMAL], false));
	}

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

bool KeyboardHook::mouse_hook_callback(int nCode, WPARAM event, LPARAM ptr_mouse) { 
    auto mouse = (MouseHookStruct*) ptr_mouse;
    if(event == WM_RBUTTONDOWN)
        this->trigger_key_event(KeyEvent::PRESS, "MOUSE1");
    else if(event == WM_RBUTTONUP)
        this->trigger_key_event(KeyEvent::RELEASE, "MOUSE1");
    if(event == WM_LBUTTONDOWN)
        this->trigger_key_event(KeyEvent::PRESS, "MOUSE2");
    else if(event == WM_LBUTTONUP)
        this->trigger_key_event(KeyEvent::RELEASE, "MOUSE2");
    if(event == WM_MBUTTONDOWN)
        this->trigger_key_event(KeyEvent::PRESS, "MOUSE3");
    else if(event == WM_MBUTTONUP) 
        this->trigger_key_event(KeyEvent::RELEASE, "MOUSE3");
    else if(event == WM_XBUTTONDOWN || event == WM_XBUTTONUP) { 
        auto x_index = GET_XBUTTON_WPARAM(mouse->mouseData);
        this->trigger_key_event(event == WM_XBUTTONDOWN ? KeyEvent::PRESS : KeyEvent::RELEASE, "MOUSEX" + std::to_string(x_index));
    }
    return false;
}