#include "DRI/Dri_NkroRaw.h"

#include <QtCore/QVector>
#include <Windows.h>

namespace
{

// Device filter and scancode-to-usage mapping.

QString Dri_NkroRaw_GetDevicePath(HANDLE DeviceHandle)
{
    if (DeviceHandle == nullptr)
    {
        return QString();
    }

    UINT NeedChars = 0;
    GetRawInputDeviceInfoW(DeviceHandle, RIDI_DEVICENAME, nullptr, &NeedChars);
    if (NeedChars == 0)
    {
        return QString();
    }

    QVector<wchar_t> Buffer(static_cast<int>(NeedChars) + 1, 0);
    if (GetRawInputDeviceInfoW(DeviceHandle, RIDI_DEVICENAME, Buffer.data(), &NeedChars) == static_cast<UINT>(-1))
    {
        return QString();
    }

    return QString::fromWCharArray(Buffer.constData()).trimmed();
}

bool Dri_NkroRaw_IsTargetDevice(const QString& DevicePath, const Mid_Struct_DeviceConfig& DeviceConfig)
{
    if (DevicePath.isEmpty())
    {
        return false;
    }

    const QString UpperPath = DevicePath.toUpper();
    return UpperPath.contains(QStringLiteral("VID_%1").arg(DeviceConfig.VendorId, 4, 16, QLatin1Char('0')).toUpper()) &&
        UpperPath.contains(QStringLiteral("PID_%1").arg(DeviceConfig.ProductId, 4, 16, QLatin1Char('0')).toUpper());
}

quint16 Dri_NkroRaw_GetUsage(const RAWKEYBOARD& Keyboard)
{
    const bool IsE0 = (Keyboard.Flags & RI_KEY_E0) != 0;
    const bool IsE1 = (Keyboard.Flags & RI_KEY_E1) != 0;
    const USHORT ScanCode = Keyboard.MakeCode;

    if (IsE1)
    {
        return 0;
    }

    if (IsE0)
    {
        switch (ScanCode)
        {
        case 0x35: return 0x0054;
        case 0x1C: return 0x0058;
        case 0x1D: return 0x00E4;
        case 0x38: return 0x00E6;
        case 0x5B: return 0x00E3;
        case 0x5C: return 0x00E7;
        default:
            return 0;
        }
    }

    switch (ScanCode)
    {
    case 0x45: return 0x0053;
    case 0x37: return 0x0055;
    case 0x4A: return 0x0056;
    case 0x4E: return 0x0057;
    case 0x47: return 0x005F;
    case 0x48: return 0x0060;
    case 0x49: return 0x0061;
    case 0x4B: return 0x005C;
    case 0x4C: return 0x005D;
    case 0x4D: return 0x005E;
    case 0x4F: return 0x0059;
    case 0x50: return 0x005A;
    case 0x51: return 0x005B;
    case 0x52: return 0x0062;
    case 0x53: return 0x0063;
    case 0x1D: return 0x00E0;
    case 0x2A: return 0x00E1;
    case 0x36: return 0x00E5;
    case 0x38: return 0x00E2;
    default:
        return 0;
    }
}

} // namespace

void Dri_NkroRaw_Close(Dri_NkroRaw_Struct_Port* p_Port)
{
    *p_Port = Dri_NkroRaw_Struct_Port();
}

bool Dri_NkroRaw_Init(
    Dri_NkroRaw_Struct_Port* p_Port,
    const Mid_Struct_DeviceConfig& DeviceConfig,
    void* WindowHandle,
    QString* p_TextStatus)
{
    Dri_NkroRaw_Close(p_Port);

    if (WindowHandle == nullptr)
    {
        if (p_TextStatus != nullptr)
        {
            *p_TextStatus = QStringLiteral("NKRO 原生输入链路打开失败：窗口句柄为空。");
        }
        return false;
    }

    RAWINPUTDEVICE Device = {};
    Device.usUsagePage = MID_CONST_USAGE_PAGE_NKRO;
    Device.usUsage = MID_CONST_USAGE_NKRO;
    Device.dwFlags = RIDEV_INPUTSINK;
    Device.hwndTarget = reinterpret_cast<HWND>(WindowHandle);

    if (!RegisterRawInputDevices(&Device, 1, sizeof(Device)))
    {
        if (p_TextStatus != nullptr)
        {
            *p_TextStatus = QStringLiteral("NKRO 原生输入链路注册失败：%1").arg(GetLastError());
        }
        return false;
    }

    p_Port->IsOpened = true;
    p_Port->WindowHandle = WindowHandle;
    p_Port->DeviceConfig = DeviceConfig;

    if (p_TextStatus != nullptr)
    {
        *p_TextStatus = QStringLiteral("已启用 NKRO 原生输入链路。");
    }
    return true;
}

bool Dri_NkroRaw_HandleMessage(
    Dri_NkroRaw_Struct_Port* p_Port,
    void* p_Message,
    QString* p_TextStatus)
{
    if (!p_Port->IsOpened || (p_Message == nullptr))
    {
        return false;
    }

    MSG* p_Msg = reinterpret_cast<MSG*>(p_Message);
    if (p_Msg->message != WM_INPUT)
    {
        return false;
    }

    UINT NeedSize = 0;
    GetRawInputData(reinterpret_cast<HRAWINPUT>(p_Msg->lParam), RID_INPUT, nullptr, &NeedSize, sizeof(RAWINPUTHEADER));
    if (NeedSize == 0)
    {
        return false;
    }

    QByteArray Buffer(static_cast<int>(NeedSize), 0);
    if (GetRawInputData(
            reinterpret_cast<HRAWINPUT>(p_Msg->lParam),
            RID_INPUT,
            Buffer.data(),
            &NeedSize,
            sizeof(RAWINPUTHEADER)) == static_cast<UINT>(-1))
    {
        return false;
    }

    RAWINPUT* p_Input = reinterpret_cast<RAWINPUT*>(Buffer.data());
    if (p_Input->header.dwType != RIM_TYPEKEYBOARD)
    {
        return false;
    }

    const QString DevicePath = Dri_NkroRaw_GetDevicePath(p_Input->header.hDevice);
    if (!Dri_NkroRaw_IsTargetDevice(DevicePath, p_Port->DeviceConfig))
    {
        return false;
    }

    const quint16 Usage = Dri_NkroRaw_GetUsage(p_Input->data.keyboard);
    if (Usage == 0)
    {
        return false;
    }

    const bool IsPressed = (p_Input->data.keyboard.Flags & RI_KEY_BREAK) == 0;
    bool IsChanged = false;

    if ((Usage >= 0x00E0) && (Usage <= 0x00E7))
    {
        const quint8 BitMask = static_cast<quint8>(1U << (Usage - 0x00E0));
        const quint8 OldModifier = p_Port->Modifier;
        if (IsPressed)
        {
            p_Port->Modifier = static_cast<quint8>(p_Port->Modifier | BitMask);
        }
        else
        {
            p_Port->Modifier = static_cast<quint8>(p_Port->Modifier & static_cast<quint8>(~BitMask));
        }
        IsChanged = (OldModifier != p_Port->Modifier);
    }
    else
    {
        const int ByteIndex = Usage / 8;
        const quint8 BitMask = static_cast<quint8>(1U << (Usage % 8));
        quint8 Value = static_cast<quint8>(p_Port->UsageBitmap.at(ByteIndex));
        const bool OldPressed = (Value & BitMask) != 0;
        if (OldPressed == IsPressed)
        {
            return false;
        }
        Value = IsPressed ? static_cast<quint8>(Value | BitMask) : static_cast<quint8>(Value & static_cast<quint8>(~BitMask));
        p_Port->UsageBitmap[ByteIndex] = static_cast<char>(Value);
        IsChanged = true;
    }

    if (!IsChanged)
    {
        return false;
    }

    if (p_Port->DevicePath != DevicePath)
    {
        p_Port->DevicePath = DevicePath;
        if (p_TextStatus != nullptr)
        {
            *p_TextStatus = QStringLiteral("原生输入已命中目标设备：%1").arg(DevicePath);
        }
    }

    Mid_Struct_RawPacket Packet;
    Packet.IsValid = true;
    Packet.Source = Mid_Enum_RawPacketSource_UsbNkroRaw;
    Packet.PortName = QStringLiteral("NKRO(原生输入)");
    Packet.ByteArray = QByteArray(MID_CONST_PACKET_SIZE_NKRO, 0);
    Packet.ByteArray[0] = static_cast<char>(Mid_Enum_ReportId_Nkro);
    Packet.ByteArray[1] = static_cast<char>(p_Port->Modifier);

    for (int Index = 0; Index < MID_CONST_USAGE_BITMAP_SIZE; ++Index)
    {
        Packet.ByteArray[2 + Index] = p_Port->UsageBitmap.at(Index);
    }

    p_Port->PacketQueue.append(Packet);
    return true;
}

bool Dri_NkroRaw_Read(
    Dri_NkroRaw_Struct_Port* p_Port,
    Mid_Struct_RawPacket* p_Packet,
    QString*)
{
    p_Packet->IsValid = false;
    p_Packet->Source = Mid_Enum_RawPacketSource_UsbNkroRaw;
    p_Packet->ByteArray.clear();
    p_Packet->PortName = QStringLiteral("NKRO(原生输入)");

    if (!p_Port->IsOpened || p_Port->PacketQueue.isEmpty())
    {
        return false;
    }

    *p_Packet = p_Port->PacketQueue.takeFirst();
    return p_Packet->IsValid;
}