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

using namespace std;

typedef KBDLLHOOKSTRUCT KeyboardHookStruct;
typedef MSLLHOOKSTRUCT MouseHookStruct;
thread_local KeyboardHook* thread_hook{nullptr};

struct MouseButtonEventEntry {
    MouseButtonEventEntry* next;
    KeyboardHook* hook;

    enum KeyboardHook::KeyEvent::type type;
    std::string key;
};

inline MouseButtonEventEntry* allocate_mb_event() {
    return new MouseButtonEventEntry{};
}

inline void delete_mb_event(MouseButtonEventEntry* event) {
    delete event;
}

struct MouseButtonEventDispatcher {
    bool active{true};

    std::thread dispatcher{};

    CRITICAL_SECTION mutex;
    CONDITION_VARIABLE cv_flushed;
    CONDITION_VARIABLE cv_work;

    MouseButtonEventEntry* event_head{nullptr};
    MouseButtonEventEntry** event_tail{&event_head};
};

MouseButtonEventDispatcher* global_event_dispatcher{};
size_t global_ed_ref_count{0};

void init_global_ed() {
    if(global_event_dispatcher) {
        global_ed_ref_count++;
        return;
    }

    global_event_dispatcher = new MouseButtonEventDispatcher{};
    InitializeCriticalSection(&global_event_dispatcher->mutex);
    InitializeConditionVariable(&global_event_dispatcher->cv_flushed);
    InitializeConditionVariable(&global_event_dispatcher->cv_work);

    global_event_dispatcher->dispatcher = std::thread([]{
        auto ed = global_event_dispatcher;

        while(ed->active) {
            MouseButtonEventEntry* entry{nullptr};
            {
                EnterCriticalSection(&ed->mutex);
                while(!global_event_dispatcher->event_head && ed->active) {
                    WakeAllConditionVariable(&ed->cv_flushed);
                    SleepConditionVariableCS(&ed->cv_work, &ed->mutex, INFINITE);
                }

                entry = global_event_dispatcher->event_head;
                global_event_dispatcher->event_head = nullptr;
                global_event_dispatcher->event_tail = &global_event_dispatcher->event_head;
                LeaveCriticalSection(&ed->mutex);
            }

            while(entry) {
                entry->hook->trigger_key_event(entry->type, std::string{entry->key});

                auto next = entry->next;
                delete_mb_event(entry);
                entry = next;
            }
        }
    });
}

void shutdown_global_ed() {
    if(--global_ed_ref_count > 0) return;

    auto ed = std::exchange(global_event_dispatcher, nullptr);
    ed->active = false;
    WakeAllConditionVariable(&ed->cv_work);

    if(ed->dispatcher.joinable())
        ed->dispatcher.join();

    DeleteCriticalSection(&ed->mutex);
    delete ed;
}

KeyboardHook::KeyboardHook() = default;

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

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

    init_global_ed();
	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();

	/* all events flushed */
    EnterCriticalSection(&global_event_dispatcher->mutex);
    WakeAllConditionVariable(&global_event_dispatcher->cv_work);
    SleepConditionVariableCS(&global_event_dispatcher->cv_flushed, &global_event_dispatcher->mutex, INFINITE);
    LeaveCriticalSection(&global_event_dispatcher->mutex);
    shutdown_global_ed();

	this->_attached = false;
}

LRESULT _keyboard_hook_callback(int nCode, WPARAM event, LPARAM ptr_keyboard) {
    assert(thread_hook);
    auto consume = thread_hook->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) {
    assert(thread_hook);
    auto consume = thread_hook->mouse_hook_callback(nCode, event, ptr_keyboard);
    auto end = std::chrono::high_resolution_clock::now();
    if(consume)
        return 1;
    return CallNextHookEx(nullptr, nCode, event, ptr_keyboard);
}

void KeyboardHook::poll_events() {
    thread_hook = this;

    if(!SetPriorityClass(GetCurrentProcess(), REALTIME_PRIORITY_CLASS))
        std::cerr << "Failed to set priority class to realtime!" << std::endl;

#if 0
    int    cur_priority		  = GetThreadPriority(GetCurrentThread());
    DWORD  cur_priority_class = GetPriorityClass(GetCurrentProcess());
    std::cout << "P: " << cur_priority << " C: " << cur_priority_class << std::endl;
#endif

	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);
	}
    std::this_thread::sleep_for(std::chrono::seconds{1});

	UnhookWindowsHookEx(this->mouse_hook_id);
    UnhookWindowsHookEx(this->keyboad_hook_id);
    thread_hook = 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) {
    MouseButtonEventEntry* mb_event;
    switch (event) {
        case WM_LBUTTONDOWN:
            mb_event = allocate_mb_event();
            mb_event->type = KeyEvent::PRESS;
            mb_event->key = "MOUSE1";
            break;
        case WM_LBUTTONUP:
            mb_event = allocate_mb_event();
            mb_event->type = KeyEvent::RELEASE;
            mb_event->key = "MOUSE1";
            break;

        case WM_RBUTTONDOWN:
            mb_event = allocate_mb_event();
            mb_event->type = KeyEvent::PRESS;
            mb_event->key = "MOUSE3";
            break;
        case WM_RBUTTONUP:
            mb_event = allocate_mb_event();
            mb_event->type = KeyEvent::RELEASE;
            mb_event->key = "MOUSE3";
            break;

        case WM_XBUTTONDOWN: {
            auto mouse = (MouseHookStruct*) ptr_mouse;
            auto x_index = GET_XBUTTON_WPARAM(mouse->mouseData);

            mb_event = allocate_mb_event();
            mb_event->type = KeyEvent::PRESS;
            mb_event->key = "MOUSEX" + std::to_string(x_index);
            break;
        }
        case WM_XBUTTONUP: {
            auto mouse = (MouseHookStruct*) ptr_mouse;
            auto x_index = GET_XBUTTON_WPARAM(mouse->mouseData);

            mb_event = allocate_mb_event();
            mb_event->type = KeyEvent::RELEASE;
            mb_event->key = "MOUSEX" + std::to_string(x_index);
            break;
        }
        default:
            return false;
    }

    mb_event->next = nullptr;
    mb_event->hook = thread_hook;

    EnterCriticalSection(&global_event_dispatcher->mutex);
    *global_event_dispatcher->event_tail = mb_event;
    global_event_dispatcher->event_tail = &mb_event->next;
    WakeAllConditionVariable(&global_event_dispatcher->cv_work);
    LeaveCriticalSection(&global_event_dispatcher->mutex);
    return false;
}