#include <windows.h>

#include <stdint.h>
#include <stdlib.h>
#include <stdbool.h>
#include <process.h>

#include "hook/iohook.h"
#include "util/dprintf.h"

#include "ri/ri.h"

static HWND ri_hwnd;
static WNDCLASSEX ri_wndclass;
static uintptr_t ri_thread;

static struct RI_DEVICE selected_ri_dev = 
{
    .hDevice = INVALID_HANDLE_VALUE,
    .capeabilities = NULL,
    .preparsed_data = NULL,
    .prod_name = L"",
    .dev_name = L""
};

static ri_mouse_cb_t mouse_cb = NULL;
static ri_kb_cb_t kb_cb = NULL;
static ri_hid_cb_t hid_cb = NULL;

LRESULT CALLBACK ri_win_proc(HWND hWnd, UINT msg, WPARAM wparam, LPARAM lParam);
static unsigned int __stdcall ri_thread_proc(void *ctx);

HRESULT ri_init(uint8_t dev_num, ri_mouse_cb_t in_mouse_cb, ri_kb_cb_t in_kb_cb, ri_hid_cb_t in_hid_cb)
{
    wchar_t ri_dev_name[MAX_PATH] = L"\0";
    wchar_t hid_prod_name[MAX_PATH] = L"\0";
    PHIDP_PREPARSED_DATA ppdata = NULL;
    PHIDP_CAPS caps = NULL;
    UINT num_ri_devs = 0;
    UINT len_ri_dev_name = 0;
    UINT len_ri_preparsed_data = 0;
    DWORD err = 0;
    PRAWINPUTDEVICELIST ri_devs = NULL;
    uint32_t num_final_devs = 0;
    PRI_DEVICE final_dev_list = NULL;

    mouse_cb = in_mouse_cb;
    kb_cb = in_kb_cb;
    hid_cb = in_hid_cb;

    if (GetRawInputDeviceList(NULL, &num_ri_devs, sizeof(RAWINPUTDEVICELIST)) != 0) {
        err = GetLastError();
        dprintf("RawInput: GetRawInputDeviceList failed to get number of devices %08lX\n", err);
        return E_FAIL;
    }

    ri_devs = malloc(sizeof(RAWINPUTDEVICELIST) * num_ri_devs);
    GetRawInputDeviceList(ri_devs, &num_ri_devs, sizeof(RAWINPUTDEVICELIST));

    if (num_ri_devs <= 0) {
        err = GetLastError();
        dprintf("RawInput: GetRawInputDeviceList failed to get devices info %08lX\n", err);
        free(ri_devs);
        return E_FAIL;
    }

    dprintf("RawInput: Raw Input device list (%d devices):\n", num_ri_devs);
    for (int i = 0; i < num_ri_devs; i++) {
        len_ri_dev_name = 0;
        len_ri_preparsed_data = 0;
        if (GetRawInputDeviceInfo(ri_devs[i].hDevice, RIDI_PREPARSEDDATA, NULL, &len_ri_preparsed_data) != 0) {
            dprintf("RawInput: GetRawInputDeviceInfo #%d failed to get preprocessed data size\n", i);
            break;
        }
        
        if (len_ri_preparsed_data <= 0) { 
            dprintf("RawInput: GetRawInputDeviceInfoW #%d len_ri_preparsed_data <= 0\n", i);
            break;
        }
        
        if (GetRawInputDeviceInfo(ri_devs[i].hDevice, RIDI_PREPARSEDDATA, ppdata, &len_ri_preparsed_data) != 0) {
            dprintf("RawInput: GetRawInputDeviceInfo #%d failed to get preprocessed data\n", i);
            break;
        }

        if (HidP_GetCaps(ppdata, caps) != HIDP_STATUS_SUCCESS) {
            dprintf("RawInput: HidP_GetCaps #%d failed to get capeabilities data size\n", i);
            break;
        }

        if ((caps->UsagePage >> 8) == 0xFF) {
            dprintf("RawInput: Skipping vendor-specific device #%d\n", i);
            continue;
        }

        if (GetRawInputDeviceInfoW(ri_devs[i].hDevice, RIDI_DEVICENAME, NULL, &len_ri_dev_name) != 0) {
            dprintf("RawInput: GetRawInputDeviceInfoW #%d failed to get name size\n", i);
            break;
        }
        
        if (len_ri_dev_name <= 0) { 
            dprintf("RawInput: GetRawInputDeviceInfoW #%d len_ri_dev_name <= 0\n", i);
            break;
        }

        if (len_ri_dev_name >= MAX_PATH) {
            dprintf("\tDev #%d - (name too long, %d characters)\n", i, len_ri_dev_name);
            continue;
        }

        if (GetRawInputDeviceInfoW(ri_devs[i].hDevice, RIDI_DEVICENAME, ri_dev_name, &len_ri_dev_name) <= 0) {
            dprintf("RawInput: GetRawInputDeviceInfoW #%d failed to get name\n", i);
            break;
        }

        HANDLE hiddev = CreateFileW(ri_dev_name, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

        if (hiddev == NULL) {
            dprintf("RawInput: CreateFileW failed on dev #%d for device %ls\n", i, ri_dev_name);
            continue;
        }

        if (!HidD_GetProductString(hiddev, hid_prod_name, MAX_PATH)) {
            continue;
        }

        PRI_DEVICE new_dev;
        PRI_DEVICE new_mem = realloc(final_dev_list, (num_final_devs + 1) * sizeof(struct RI_DEVICE));
        
        new_dev = &new_mem[num_final_devs];
        new_dev->hDevice = ri_devs[i].hDevice;
        // does this need memcpy?
        new_dev->capeabilities = caps;
        new_dev->preparsed_data = ppdata;
        new_dev->dev_name = ri_dev_name;
        new_dev->prod_name = hid_prod_name;

        final_dev_list = new_mem;
        num_final_devs++;
    }

    if (dev_num >= num_final_devs) {
        dprintf("RawInput: No RawInput device set, please select from the list below:\n");
        for (int i = 0; i < num_final_devs; i++) {
            dprintf("\tDev #%d - %ls\n", i, hid_prod_name);
        }
        
        free(ri_devs);
        free(final_dev_list);
        return S_FALSE;
    }
    
    selected_ri_dev = final_dev_list[dev_num];
    dprintf("RawInput: Using device %ls\n", selected_ri_dev.prod_name);

    ri_wndclass.cbSize = sizeof(WNDCLASSEX);
    ri_wndclass.hInstance = GetModuleHandle(NULL);
    ri_wndclass.lpfnWndProc = ri_win_proc;
    ri_wndclass.lpszClassName = "Taitools Input";

    RegisterClassEx(&ri_wndclass);

    ri_thread = _beginthreadex(NULL, 0, ri_thread_proc, NULL, 0, NULL);
    
    free(ri_devs);
    free(final_dev_list);
    return S_OK;
}

void ri_exit()
{
    if (ri_hwnd) {
        PostMessage(ri_hwnd, WM_CLOSE, 0, 0);
    }
}

static unsigned int __stdcall ri_thread_proc(void *ctx)
{
    if (selected_ri_dev.hDevice == INVALID_HANDLE_VALUE) {
        dprintf("RawInput: ri_thread_proc cannont begin, no device selected\n");
        return (LRESULT)0;
    }

    dprintf("RawInput: ri_thread_proc begin\n");
    SetThreadPriority(GetCurrentThread(), THREAD_PRIORITY_TIME_CRITICAL);

    ri_hwnd = CreateWindowExA(0, ri_wndclass.lpszClassName, "Taitools Input", 0, 0, 0, 0, 0, NULL, NULL, ri_wndclass.hInstance, NULL);   
 
    RAWINPUTDEVICE Rid[4];
    Rid[0].usUsagePage = HID_USAGE_PAGE_GENERIC;
    Rid[0].usUsage     = HID_USAGE_GENERIC_MOUSE;
    Rid[0].dwFlags     = RIDEV_NOLEGACY;    // adds mouse and also ignores legacy mouse messages
    Rid[0].hwndTarget  = ri_hwnd;

    Rid[1].usUsagePage = HID_USAGE_PAGE_GENERIC;
    Rid[1].usUsage     = HID_USAGE_GENERIC_KEYBOARD;
    Rid[1].dwFlags     = RIDEV_NOLEGACY;    // adds keyboard and also ignores legacy keyboard messages
    Rid[1].hwndTarget  = ri_hwnd;

    Rid[2].usUsagePage = HID_USAGE_PAGE_GENERIC;
    Rid[2].usUsage     = HID_USAGE_GENERIC_GAMEPAD;
    Rid[2].dwFlags     = 0;                 // adds game pad
    Rid[2].hwndTarget  = ri_hwnd;

    Rid[3].usUsagePage = HID_USAGE_PAGE_GENERIC;
    Rid[3].usUsage     = HID_USAGE_GENERIC_JOYSTICK;
    Rid[3].dwFlags     = 0;                 // adds joystick
    Rid[3].hwndTarget  = ri_hwnd;
    
    if (!RegisterRawInputDevices(Rid, 4, sizeof(Rid[0]))) {
        dprintf("RawInput: Failed to register RI devices\n");
        return E_FAIL;
    }

    MSG msg;
    while (GetMessage(&msg, ri_hwnd, 0, 0) > 0) {
        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }

    DestroyWindow(ri_hwnd);
    ri_hwnd = NULL;
    UnregisterClass(ri_wndclass.lpszClassName, ri_wndclass.hInstance);
    dprintf("RawInput: ri_thread_proc end\n");
    return S_OK;
}

LRESULT CALLBACK ri_win_proc(HWND hWnd, UINT msg, WPARAM wparam, LPARAM lParam)
{
    switch (msg) {        
        case WM_INPUT:
            //dprintf("RawInput: Got input\n");
            UINT pdata_size = 0;
            if (GetRawInputData((HRAWINPUT)lParam, RID_INPUT, NULL, &pdata_size, sizeof(RAWINPUTHEADER)) != -1 && pdata_size > 0) {
                PRAWINPUT ridata = (PRAWINPUT)malloc(pdata_size);
                GetRawInputData((HRAWINPUT)lParam, RID_INPUT, ridata, &pdata_size, sizeof(RAWINPUTHEADER));

                if (ridata->header.hDevice != selected_ri_dev.hDevice) {
                    free(ridata);
                    return (LRESULT)0;
                }

                switch (ridata->header.dwType) {
                    case RIM_TYPEMOUSE:
                        if (mouse_cb != NULL) {
                            mouse_cb(&selected_ri_dev, &ridata->data.mouse);
                        }
                        break;
                    case RIM_TYPEKEYBOARD:
                        if (kb_cb != NULL) {
                            kb_cb(&selected_ri_dev, &ridata->data.keyboard);
                        }
                        break;
                    case RIM_TYPEHID:
                        if (hid_cb != NULL) {
                            hid_cb(&selected_ri_dev, &ridata->data.hid);
                        }
                        break;
                }

                free(ridata);
            }
            return (LRESULT)0;
        case WM_QUIT:
            dprintf("RI: Exit requested\n");
            ExitProcess(0);
            return (LRESULT)0;
    }

    return DefWindowProc(hWnd, msg, wparam, lParam);
}