#include "RawInputDetector.h"

#include <QCoreApplication>
#include <QDebug>
#include <QGuiApplication>

#ifdef _WIN32
#include <windows.h>
#include <hidusage.h>

namespace {
constexpr UINT kRidInput = 0x10000003;
constexpr USHORT kE0Flag = 0x02;
constexpr USHORT kMouseMoveAbsolute = 0x0001;
constexpr USHORT kMouseLeftDown = 0x0001;
constexpr USHORT kMouseLeftUp = 0x0002;
constexpr USHORT kMouseRightDown = 0x0004;
constexpr USHORT kMouseRightUp = 0x0008;
constexpr USHORT kMouseMiddleDown = 0x0010;
constexpr USHORT kMouseMiddleUp = 0x0020;
constexpr USHORT kMouseX1Down = 0x0040;
constexpr USHORT kMouseX1Up = 0x0080;
constexpr USHORT kMouseX2Down = 0x0100;
constexpr USHORT kMouseX2Up = 0x0200;
constexpr USHORT kMouseWheel = 0x0400;
}
#endif

namespace FancyInput {

RawInputDetector::RawInputDetector(QObject *parent)
    : QObject(parent)
{
    m_rawBuffer.resize(64);
}

RawInputDetector::~RawInputDetector()
{
    stop();
}

bool RawInputDetector::start(QWindow *targetWindow)
{
#ifdef _WIN32
    if (m_running || targetWindow == nullptr || targetWindow->winId() == 0)
        return false;

    if (!registerDevices(reinterpret_cast<HWND>(targetWindow->winId())))
        return false;

    m_targetWindow = targetWindow;
    m_pressedKeys.clear();
    QCoreApplication::instance()->installNativeEventFilter(this);
    m_running = true;
    return true;
#else
    Q_UNUSED(targetWindow)
    return false;
#endif
}

void RawInputDetector::stop()
{
#ifdef _WIN32
    if (!m_running)
        return;

    QCoreApplication::instance()->removeNativeEventFilter(this);
    unregisterDevices();
    m_pressedKeys.clear();
    m_targetWindow = nullptr;
    m_running = false;
#endif
}

bool RawInputDetector::nativeEventFilter(const QByteArray &eventType,
                                         void *message,
                                         qintptr *result)
{
    Q_UNUSED(result)
#ifdef _WIN32
    if (!m_running || eventType != "windows_generic_MSG" || message == nullptr)
        return false;

    const MSG *msg = static_cast<const MSG *>(message);
    if (msg->message == WM_INPUT)
        handleRawInput(reinterpret_cast<HRAWINPUT>(msg->lParam));
#else
    Q_UNUSED(eventType)
    Q_UNUSED(message)
#endif
    return false;
}

#ifdef _WIN32
bool RawInputDetector::registerDevices(HWND targetHandle)
{
    RAWINPUTDEVICE devices[2] = {};
    devices[0].usUsagePage = HID_USAGE_PAGE_GENERIC;
    devices[0].usUsage = HID_USAGE_GENERIC_KEYBOARD;
    devices[0].dwFlags = RIDEV_INPUTSINK;
    devices[0].hwndTarget = targetHandle;
    devices[1].usUsagePage = HID_USAGE_PAGE_GENERIC;
    devices[1].usUsage = HID_USAGE_GENERIC_MOUSE;
    devices[1].dwFlags = RIDEV_INPUTSINK;
    devices[1].hwndTarget = targetHandle;

    return RegisterRawInputDevices(devices, 2, sizeof(RAWINPUTDEVICE)) != FALSE;
}

void RawInputDetector::unregisterDevices()
{
    RAWINPUTDEVICE devices[2] = {};
    devices[0].usUsagePage = HID_USAGE_PAGE_GENERIC;
    devices[0].usUsage = HID_USAGE_GENERIC_KEYBOARD;
    devices[0].dwFlags = RIDEV_REMOVE;
    devices[1].usUsagePage = HID_USAGE_PAGE_GENERIC;
    devices[1].usUsage = HID_USAGE_GENERIC_MOUSE;
    devices[1].dwFlags = RIDEV_REMOVE;
    RegisterRawInputDevices(devices, 2, sizeof(RAWINPUTDEVICE));
}

void RawInputDetector::handleRawInput(HRAWINPUT rawInput)
{
    UINT size = 0;
    if (GetRawInputData(rawInput, RID_INPUT, nullptr, &size,
                        sizeof(RAWINPUTHEADER)) == static_cast<UINT>(-1) || size == 0)
        return;

    if (m_rawBuffer.size() < static_cast<qsizetype>(size))
        m_rawBuffer.resize(static_cast<qsizetype>(size));
    if (GetRawInputData(rawInput, RID_INPUT, m_rawBuffer.data(), &size,
                        sizeof(RAWINPUTHEADER)) == static_cast<UINT>(-1))
        return;

    const auto *input = reinterpret_cast<const RAWINPUT *>(m_rawBuffer.constData());
    if (input->header.dwType == RIM_TYPEKEYBOARD)
        handleKeyboard(input->data.keyboard);
    else if (input->header.dwType == RIM_TYPEMOUSE)
        handleMouse(input->data.mouse);
}

void RawInputDetector::handleKeyboard(const RAWKEYBOARD &keyboard)
{
    const bool keyDown = keyboard.Message == WM_KEYDOWN || keyboard.Message == WM_SYSKEYDOWN;
    const bool keyUp = keyboard.Message == WM_KEYUP || keyboard.Message == WM_SYSKEYUP;
    if (!keyDown && !keyUp)
        return;

    const bool extended = (keyboard.Flags & kE0Flag) != 0;
    const quint16 virtualKey = keyboard.VKey;
    const FipKeys key = mapKey(virtualKey, keyboard.MakeCode, extended);
    const FipRawKeys rawKey = mapRawKey(virtualKey, keyboard.MakeCode, extended);
    if (keyDown) {
        if (m_pressedKeys.contains(virtualKey))
            return;
        m_pressedKeys.insert(virtualKey);
    }
    if (keyUp)
        m_pressedKeys.remove(virtualKey);

    const bool ctrl = (GetAsyncKeyState(VK_LCONTROL) & 0x8000) || (GetAsyncKeyState(VK_RCONTROL) & 0x8000);
    const bool shift = (GetAsyncKeyState(VK_LSHIFT) & 0x8000) || (GetAsyncKeyState(VK_RSHIFT) & 0x8000);
    const bool alt = (GetAsyncKeyState(VK_LMENU) & 0x8000) || (GetAsyncKeyState(VK_RMENU) & 0x8000);
    const bool win = (GetAsyncKeyState(VK_LWIN) & 0x8000) || (GetAsyncKeyState(VK_RWIN) & 0x8000);
    const bool numLock = (GetKeyState(VK_NUMLOCK) & 1) != 0;
    const bool capsLock = (GetKeyState(VK_CAPITAL) & 1) != 0;
    const bool scrollLock = (GetKeyState(VK_SCROLL) & 1) != 0;

    if (keyDown)
        emit keyPressed(key, rawKey, virtualKey, keyboard.MakeCode, extended, ctrl, shift, alt, win, numLock, capsLock, scrollLock);
    else
        emit keyReleased(key, rawKey, virtualKey, keyboard.MakeCode, extended, ctrl, shift, alt, win, numLock, capsLock, scrollLock);
}

void RawInputDetector::handleMouse(const RAWMOUSE &mouse)
{
    const USHORT flags = mouse.usButtonFlags;
    if (flags & kMouseLeftDown) emit mouseButtonChanged(MouseCodeType::Left, true);
    if (flags & kMouseLeftUp) emit mouseButtonChanged(MouseCodeType::Left, false);
    if (flags & kMouseRightDown) emit mouseButtonChanged(MouseCodeType::Right, true);
    if (flags & kMouseRightUp) emit mouseButtonChanged(MouseCodeType::Right, false);
    if (flags & kMouseMiddleDown) emit mouseButtonChanged(MouseCodeType::Middle, true);
    if (flags & kMouseMiddleUp) emit mouseButtonChanged(MouseCodeType::Middle, false);
    if (flags & kMouseX1Down) emit mouseButtonChanged(MouseCodeType::XButton1, true);
    if (flags & kMouseX1Up) emit mouseButtonChanged(MouseCodeType::XButton1, false);
    if (flags & kMouseX2Down) emit mouseButtonChanged(MouseCodeType::XButton2, true);
    if (flags & kMouseX2Up) emit mouseButtonChanged(MouseCodeType::XButton2, false);
    if (flags & kMouseWheel)
        emit mouseWheelMoved(static_cast<qint16>(mouse.usButtonData));
    if ((mouse.usFlags & kMouseMoveAbsolute) == 0 && (mouse.lLastX != 0 || mouse.lLastY != 0))
        emit mouseMoved(mouse.lLastX, mouse.lLastY);
}
#endif

FipKeys RawInputDetector::mapKey(quint16 virtualKey, quint16 scanCode, bool extended)
{
    if (virtualKey == VK_PAUSE) return FipKeys::Pause;
    if (virtualKey == VK_SNAPSHOT) return FipKeys::PrintScreen;
    if (virtualKey == VK_RETURN && scanCode == 0x1C) return extended ? FipKeys::NumpadEnter : FipKeys::Enter;
    if (virtualKey == VK_CONTROL && scanCode == 0x1D) return extended ? FipKeys::RControlKey : FipKeys::LControlKey;
    if (virtualKey == VK_MENU && scanCode == 0x38) return extended ? FipKeys::RMenu : FipKeys::LMenu;
    if (virtualKey == VK_SHIFT) {
        if (scanCode == 0x2A) return FipKeys::LShiftKey;
        if (scanCode == 0x36) return FipKeys::RShiftKey;
    }
    if (virtualKey == VK_LWIN) return FipKeys::LWin;
    if (virtualKey == VK_RWIN) return FipKeys::RWin;
    if (virtualKey == VK_APPS) return FipKeys::Apps;
    if (extended) {
        switch (scanCode) {
        case 0x48: return FipKeys::Up;
        case 0x4B: return FipKeys::Left;
        case 0x4D: return FipKeys::Right;
        case 0x50: return FipKeys::Down;
        case 0x47: return FipKeys::Home;
        case 0x4F: return FipKeys::End;
        case 0x49: return FipKeys::PageUp;
        case 0x51: return FipKeys::PageDown;
        case 0x52: return FipKeys::Insert;
        case 0x53: return FipKeys::Delete;
        case 0x35: return FipKeys::Divide;
        default: break;
        }
    }
    return static_cast<FipKeys>(virtualKey);
}

FipRawKeys RawInputDetector::mapRawKey(quint16 virtualKey, quint16 scanCode, bool extended)
{
    if (virtualKey == VK_PAUSE) return FipRawKeys::Pause;
    if (virtualKey == VK_SNAPSHOT) return FipRawKeys::PrintScreen;
    if (virtualKey == VK_CONTROL && scanCode == 0x1D) return extended ? FipRawKeys::RControlKey : FipRawKeys::LControlKey;
    if (virtualKey == VK_MENU && scanCode == 0x38) return extended ? FipRawKeys::RMenu : FipRawKeys::LMenu;
    if (virtualKey == VK_RETURN && scanCode == 0x1C) return extended ? FipRawKeys::NumpadEnter : FipRawKeys::Enter;
    if (virtualKey == VK_LWIN) return FipRawKeys::LWin;
    if (virtualKey == VK_RWIN) return FipRawKeys::RWin;
    if (virtualKey == VK_APPS) return FipRawKeys::Apps;
    const quint32 code = extended ? (0xE000u | scanCode) : scanCode;
    return static_cast<FipRawKeys>(code);
}

} // namespace FancyInput