#include "gamepad.h"
#include "common/io.h"
#include "common/mallocHelper.h"
#include "common/windows/unicode.h"

#include <windows.h>
#include <hidsdi.h>

static const char* detectKnownDeviceName(uint32_t vendorId, uint32_t productId) {
    switch (vendorId) {
        // Nintendo
        case 0x057E: {
            switch (productId) {
                case 0x2006:
                    return "Nintendo Switch Joycon L";
                case 0x2007:
                    return "Nintendo Switch Joycon R";
                case 0x2009:
                    return "Nintendo Switch Pro Controller";
                case 0x200E:
                    return "Nintendo Switch Charging Grip";
                case 0x2017:
                    return "Nintendo Switch SNES Controller";

                default:
                    return nullptr;
            }
        }

        // Sony
        case 0x054C: {
            switch (productId) {
                case 0x0268:
                    return "Sony DualShock 3 / Six Axis";

                case 0x05C4:
                    return "Sony DualShock 4 Gen1";
                case 0x09CC:
                    return "Sony DualShock 4 Gen2";
                case 0x0BA0:
                    return "Sony DualShock 4 USB receiver";

                case 0x0CE6:
                    return "Sony DualSense";
                case 0x0DF2:
                    return "Sony DualSense Edge";

                default:
                    return nullptr;
            }
        }

        // Logitech
        case 0x046D: {
            switch (productId) {
                case 0xC216:
                    return "Logitech F310, DirectInput";
                case 0xC218:
                    return "Logitech F510, DirectInput";
                case 0xC219:
                    return "Logitech F710, DirectInput";
                case 0xC21D:
                    return "Logitech F310";
                case 0xC21E:
                    return "Logitech F510";
                case 0xC21F:
                    return "Logitech F710";

                default:
                    return nullptr;
            }
        }

        case 0x045E: // Microsoft Xbox compatible controllers should be handled by Windows without problems
        default:
            return nullptr;
    }
}

const char* ffDetectGamepad(FFlist* devices /* List of FFGamepadDevice */) {
    UINT nDevices = 0;
    if (GetRawInputDeviceList(nullptr, &nDevices, sizeof(RAWINPUTDEVICELIST))) {
        return "GetRawInputDeviceList(nullptr) failed";
    }
    if (nDevices == 0) {
        return "No HID devices found";
    }
    FF_AUTO_FREE RAWINPUTDEVICELIST* pRawInputDeviceList = (RAWINPUTDEVICELIST*) malloc(sizeof(RAWINPUTDEVICELIST) * nDevices);
    if ((nDevices = GetRawInputDeviceList(pRawInputDeviceList, &nDevices, sizeof(RAWINPUTDEVICELIST))) == (UINT) -1) {
        return "GetRawInputDeviceList(pRawInputDeviceList) failed";
    }

    for (UINT i = 0; i < nDevices; ++i) {
        if (pRawInputDeviceList[i].dwType != RIM_TYPEHID) {
            continue;
        }

        HANDLE hDevice = pRawInputDeviceList[i].hDevice;

        RID_DEVICE_INFO rdi;
        UINT rdiSize = sizeof(rdi);
        if (GetRawInputDeviceInfoW(hDevice, RIDI_DEVICEINFO, &rdi, &rdiSize) == (UINT) -1) {
            continue;
        }

        if (rdi.hid.usUsagePage != 1 || (rdi.hid.usUsage != 4 /*Joystick*/ && rdi.hid.usUsage != 5 /*Gamepad*/)) {
            continue;
        }

        WCHAR devName[MAX_PATH] = L"";
        UINT nameSize = MAX_PATH;
        if (GetRawInputDeviceInfoW(hDevice, RIDI_DEVICENAME, devName, &nameSize) == (UINT) -1) {
            continue;
        }

        FFGamepadDevice* device = FF_LIST_ADD(FFGamepadDevice, *devices);
        ffStrbufInit(&device->serial);
        ffStrbufInit(&device->name);
        device->battery = 0;

        const char* knownGamepad = detectKnownDeviceName(rdi.hid.dwVendorId, rdi.hid.dwProductId);
        if (knownGamepad) {
            ffStrbufSetS(&device->name, knownGamepad);
        }
        FF_AUTO_CLOSE_FD HANDLE hHidFile = CreateFileW(devName, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr, OPEN_EXISTING, FILE_FLAG_OVERLAPPED, nullptr);
        if (hHidFile == INVALID_HANDLE_VALUE) {
            if (!knownGamepad) {
                ffStrbufSetF(&device->name, "Unknown gamepad %04X-%04X", (unsigned) rdi.hid.dwVendorId, (unsigned) rdi.hid.dwProductId);
            }
            continue;
        }

        if (!knownGamepad) {
            wchar_t displayName[126];
            if (HidD_GetProductString(hHidFile, displayName, sizeof(displayName) /*in bytes*/)) {
                wchar_t manufacturer[126];
                if (HidD_GetManufacturerString(hHidFile, manufacturer, sizeof(manufacturer) /*in bytes*/)) {
                    ffStrbufSetWS(&device->name, manufacturer);
                    FF_STRBUF_AUTO_DESTROY displayNameStr = ffStrbufCreateWS(displayName);
                    ffStrbufAppendC(&device->name, ' ');
                    ffStrbufAppend(&device->name, &displayNameStr);
                } else {
                    ffStrbufSetWS(&device->name, displayName);
                }
            }
        }

        wchar_t serialNumber[127] = L"";
        if (HidD_GetSerialNumberString(hHidFile, serialNumber, sizeof(serialNumber) /*in bytes*/)) {
            ffStrbufSetWS(&device->serial, serialNumber);
        }

        PHIDP_PREPARSED_DATA preparsedData = nullptr;
        if (HidD_GetPreparsedData(hHidFile, &preparsedData)) {
            HIDP_CAPS caps;
            NTSTATUS capsResult = HidP_GetCaps(preparsedData, &caps);
            HidD_FreePreparsedData(preparsedData);
            if (!NT_SUCCESS(capsResult)) {
                continue;
            }

            if (
                (rdi.hid.dwVendorId == 0x054C && (rdi.hid.dwProductId == 0x05C4 || // PS4 Gen1
                                                     rdi.hid.dwProductId == 0x09CC // PS4 Gen2
                                                     )) ||
                (rdi.hid.dwVendorId == 0x057E && (rdi.hid.dwProductId == 0x2009 // NS Pro
                                                     ))) {
                // Controller must be connected by other programs
                FF_AUTO_FREE uint8_t* reportBuffer = malloc(caps.InputReportByteLength);
                OVERLAPPED overlapped = {};
                DWORD nBytes;
                if (ReadFile(hHidFile, reportBuffer, caps.InputReportByteLength, &nBytes, &overlapped) ||
                    GetOverlappedResultEx(hHidFile, &overlapped, &nBytes, FF_IO_TERM_RESP_WAIT_MS, TRUE)) {
                    if (rdi.hid.dwVendorId == 0x054C) {
                        if (nBytes > 31) {
                            uint8_t batteryInfo = reportBuffer[caps.InputReportByteLength == 64 /*USB?*/ ? 30 : 32];
                            device->battery = (uint8_t) ((batteryInfo & 0x0f) * 100 / (batteryInfo & 0x10 /*charging?*/ ? 11 /*BATTERY_MAX_USB*/ : 8 /*BATTERY_MAX*/));
                            if (device->battery > 100) {
                                device->battery = 100;
                            }
                        }
                    } else {
                        if (nBytes > 3 && reportBuffer[0] == 0x30) {
                            uint8_t batteryInfo = reportBuffer[2];
                            device->battery = (uint8_t) (((batteryInfo & 0xE0) >> 4) * 100 / 8);
                            if (device->battery == 0) {
                                device->battery = 1;
                            } else if (device->battery > 100) {
                                device->battery = 100;
                            }
                        }
                    }
                } else {
                    CancelIo(hHidFile);
                }
            }
        }
    }

    return nullptr;
}
