#include "usbd_dinput.h"

#include "usbd_ctlreq.h"
#include <string.h>

#define HID_DESCRIPTOR_TYPE  0x21U
#define HID_REPORT_DESC      0x22U
#define HID_REQ_GET_REPORT   0x01U
#define HID_REQ_GET_IDLE     0x02U
#define HID_REQ_GET_PROTOCOL 0x03U
#define HID_REQ_SET_REPORT   0x09U
#define HID_REQ_SET_IDLE     0x0AU
#define HID_REQ_SET_PROTOCOL 0x0BU

typedef struct
{
  uint8_t altSetting[3];
  uint8_t protocol;
  uint8_t idle;
  uint8_t busy;
} USBD_DINPUT_HandleTypeDef;

static uint8_t DINPUT_Init(USBD_HandleTypeDef *pdev, uint8_t cfgidx);
static uint8_t DINPUT_DeInit(USBD_HandleTypeDef *pdev, uint8_t cfgidx);
static uint8_t DINPUT_Setup(USBD_HandleTypeDef *pdev,
                            USBD_SetupReqTypedef *req);
static uint8_t DINPUT_DataIn(USBD_HandleTypeDef *pdev, uint8_t epnum);
static uint8_t DINPUT_DataOut(USBD_HandleTypeDef *pdev, uint8_t epnum);
static uint8_t *DINPUT_GetCfgDesc(uint16_t *length);
static uint8_t *DINPUT_GetDeviceQualifierDesc(uint16_t *length);
#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
static uint8_t *DINPUT_GetUserString(USBD_HandleTypeDef *pdev,
                                     uint8_t index,
                                     uint16_t *length);
#endif

static uint8_t dinputInBuffer[DINPUT_REPORT_SIZE];
static uint8_t dinputSetupOutBuffer[DINPUT_EP_SIZE];
#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
static uint8_t dinputStringBuffer[32];
#endif

USBD_ClassTypeDef USBD_DINPUT =
{
  DINPUT_Init,
  DINPUT_DeInit,
  DINPUT_Setup,
  NULL,
  NULL,
  DINPUT_DataIn,
  DINPUT_DataOut,
  NULL,
  NULL,
  NULL,
  DINPUT_GetCfgDesc,
  DINPUT_GetCfgDesc,
  DINPUT_GetCfgDesc,
  DINPUT_GetDeviceQualifierDesc,
#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
  DINPUT_GetUserString,
#endif
};

/* Descriptor order and report sizes match the original three-interface unit.
 * IF0 is the only functional interface. IF1/IF2 are inert compatibility
 * shells and their data endpoints are halted in DINPUT_Init(). */
__ALIGN_BEGIN static uint8_t dinputConfigDescriptor[USB_DINPUT_CONFIG_DESC_SIZ] __ALIGN_END =
{
  0x09, USB_DESC_TYPE_CONFIGURATION,
  LOBYTE(USB_DINPUT_CONFIG_DESC_SIZ), HIBYTE(USB_DINPUT_CONFIG_DESC_SIZ),
  0x03, 0x01, 0x00, 0x80, 0x32,

  0x09, USB_DESC_TYPE_INTERFACE,
  0x00, 0x00, 0x01,
  0x03, 0x00, 0x00, 0x05,
  0x09, HID_DESCRIPTOR_TYPE,
  0x11, 0x01, 0x00, 0x01, HID_REPORT_DESC,
  0x51, 0x00,
  0x07, USB_DESC_TYPE_ENDPOINT,
  DINPUT_GAMEPAD_EPIN_ADDR, 0x03,
  LOBYTE(DINPUT_EP_SIZE), HIBYTE(DINPUT_EP_SIZE), 0x01,

  0x09, USB_DESC_TYPE_INTERFACE,
  0x01, 0x00, 0x02,
  0x03, 0x00, 0x00, 0x06,
  0x09, HID_DESCRIPTOR_TYPE,
  0x11, 0x01, 0x00, 0x01, HID_REPORT_DESC,
  0x22, 0x00,
  0x07, USB_DESC_TYPE_ENDPOINT,
  DINPUT_SETUP_EPIN_ADDR, 0x03,
  LOBYTE(DINPUT_EP_SIZE), HIBYTE(DINPUT_EP_SIZE), 0x01,
  0x07, USB_DESC_TYPE_ENDPOINT,
  DINPUT_SETUP_EPOUT_ADDR, 0x03,
  LOBYTE(DINPUT_EP_SIZE), HIBYTE(DINPUT_EP_SIZE), 0x01,

  0x09, USB_DESC_TYPE_INTERFACE,
  0x02, 0x00, 0x01,
  0x03, 0x00, 0x00, 0x07,
  0x09, HID_DESCRIPTOR_TYPE,
  0x11, 0x01, 0x00, 0x01, HID_REPORT_DESC,
  0x15, 0x00,
  0x07, USB_DESC_TYPE_ENDPOINT,
  DINPUT_RAW_EPIN_ADDR, 0x03,
  LOBYTE(DINPUT_EP_SIZE), HIBYTE(DINPUT_EP_SIZE), 0x01,
};

/* 11 buttons, hat and six unsigned 8-bit axes: 9 bytes total. */
__ALIGN_BEGIN static uint8_t gamepadReportDescriptor[] __ALIGN_END =
{
  0x05, 0x01, 0x09, 0x05, 0xA1, 0x01,
  0x05, 0x09, 0x19, 0x01, 0x29, 0x0B,
  0x15, 0x00, 0x25, 0x01, 0x75, 0x01,
  0x95, 0x0B, 0x81, 0x02,
  0x75, 0x01, 0x95, 0x05, 0x81, 0x03,
  0x05, 0x01, 0x09, 0x39, 0x15, 0x00,
  0x25, 0x07, 0x35, 0x00, 0x46, 0x3B, 0x01,
  0x65, 0x14, 0x75, 0x04, 0x95, 0x01,
  0x81, 0x42, 0x75, 0x04, 0x95, 0x01, 0x81, 0x03,
  0x05, 0x01,
  0x09, 0x30, 0x09, 0x31, 0x09, 0x32,
  0x09, 0x33, 0x09, 0x34, 0x09, 0x35,
  0x15, 0x00, 0x26, 0xFF, 0x00,
  0x75, 0x08, 0x95, 0x06, 0x81, 0x02,
  0xC0
};

__ALIGN_BEGIN static uint8_t setupReportDescriptor[] __ALIGN_END =
{
  0x06, 0x00, 0xFF, 0x09, 0x01, 0xA1, 0x01,
  0x09, 0x10, 0x15, 0x00, 0x26, 0xFF, 0x00,
  0x75, 0x08, 0x95, 0x1A, 0x81, 0x02,
  0x09, 0x11, 0x15, 0x00, 0x26, 0xFF, 0x00,
  0x75, 0x08, 0x95, 0x1B, 0x91, 0x02,
  0xC0
};

__ALIGN_BEGIN static uint8_t rawReportDescriptor[] __ALIGN_END =
{
  0x06, 0x00, 0xFF, 0x09, 0x02, 0xA1, 0x01,
  0x09, 0x20, 0x15, 0x00, 0x26, 0xFF, 0x00,
  0x75, 0x08, 0x95, 0x18, 0x81, 0x02,
  0xC0
};

__ALIGN_BEGIN static uint8_t dinputDeviceQualifierDescriptor[USB_LEN_DEV_QUALIFIER_DESC] __ALIGN_END =
{
  USB_LEN_DEV_QUALIFIER_DESC, USB_DESC_TYPE_DEVICE_QUALIFIER,
  0x00, 0x02, 0x00, 0x00, 0x00, USB_MAX_EP0_SIZE, 0x01, 0x00
};

_Static_assert(sizeof(gamepadReportDescriptor) == 81U,
               "DInput gamepad report descriptor length changed");
_Static_assert(sizeof(setupReportDescriptor) == 34U,
               "Setup report descriptor length changed");
_Static_assert(sizeof(rawReportDescriptor) == 21U,
               "Raw report descriptor length changed");

static uint8_t DINPUT_Init(USBD_HandleTypeDef *pdev, uint8_t cfgidx)
{
  USBD_DINPUT_HandleTypeDef *handle;
  UNUSED(cfgidx);

  pdev->pClassData = USBD_malloc(sizeof(USBD_DINPUT_HandleTypeDef));
  if (pdev->pClassData == NULL)
  {
    return USBD_FAIL;
  }
  handle = (USBD_DINPUT_HandleTypeDef *)pdev->pClassData;
  memset(handle, 0, sizeof(*handle));

  USBD_LL_OpenEP(pdev, DINPUT_GAMEPAD_EPIN_ADDR,
                 USBD_EP_TYPE_INTR, DINPUT_EP_SIZE);
  pdev->ep_in[DINPUT_GAMEPAD_EPIN_ADDR & 0x0FU].is_used = 1U;

  USBD_LL_OpenEP(pdev, DINPUT_SETUP_EPIN_ADDR,
                 USBD_EP_TYPE_INTR, DINPUT_EP_SIZE);
  pdev->ep_in[DINPUT_SETUP_EPIN_ADDR & 0x0FU].is_used = 1U;
  USBD_LL_OpenEP(pdev, DINPUT_SETUP_EPOUT_ADDR,
                 USBD_EP_TYPE_INTR, DINPUT_EP_SIZE);
  pdev->ep_out[DINPUT_SETUP_EPOUT_ADDR & 0x0FU].is_used = 1U;
  USBD_LL_OpenEP(pdev, DINPUT_RAW_EPIN_ADDR,
                 USBD_EP_TYPE_INTR, DINPUT_EP_SIZE);
  pdev->ep_in[DINPUT_RAW_EPIN_ADDR & 0x0FU].is_used = 1U;

  /* Stub HID data paths: unused IN endpoints stay at NAK. Setup OUT packets
   * are acknowledged and discarded so the interface can start cleanly. */
  USBD_LL_PrepareReceive(pdev, DINPUT_SETUP_EPOUT_ADDR,
                         dinputSetupOutBuffer, DINPUT_EP_SIZE);
  return USBD_OK;
}

static uint8_t DINPUT_DeInit(USBD_HandleTypeDef *pdev, uint8_t cfgidx)
{
  UNUSED(cfgidx);
  USBD_LL_CloseEP(pdev, DINPUT_GAMEPAD_EPIN_ADDR);
  USBD_LL_CloseEP(pdev, DINPUT_SETUP_EPIN_ADDR);
  USBD_LL_CloseEP(pdev, DINPUT_SETUP_EPOUT_ADDR);
  USBD_LL_CloseEP(pdev, DINPUT_RAW_EPIN_ADDR);
  pdev->ep_in[DINPUT_GAMEPAD_EPIN_ADDR & 0x0FU].is_used = 0U;
  pdev->ep_in[DINPUT_SETUP_EPIN_ADDR & 0x0FU].is_used = 0U;
  pdev->ep_out[DINPUT_SETUP_EPOUT_ADDR & 0x0FU].is_used = 0U;
  pdev->ep_in[DINPUT_RAW_EPIN_ADDR & 0x0FU].is_used = 0U;
  if (pdev->pClassData != NULL)
  {
    USBD_free(pdev->pClassData);
    pdev->pClassData = NULL;
  }
  return USBD_OK;
}

static uint8_t DINPUT_Setup(USBD_HandleTypeDef *pdev,
                            USBD_SetupReqTypedef *req)
{
  USBD_DINPUT_HandleTypeDef *handle =
      (USBD_DINPUT_HandleTypeDef *)pdev->pClassData;
  uint8_t interface = (uint8_t)req->wIndex;
  uint16_t status = 0U;
  uint8_t *descriptor = NULL;
  uint16_t length = 0U;

  if ((handle == NULL) || (interface > 2U))
  {
    USBD_CtlError(pdev, req);
    return USBD_FAIL;
  }

  if ((req->bmRequest & USB_REQ_TYPE_MASK) == USB_REQ_TYPE_CLASS)
  {
    /* The original Setup/Raw protocols are intentionally not implemented. */
    if (interface != 0U)
    {
      USBD_CtlError(pdev, req);
      return USBD_FAIL;
    }

    switch (req->bRequest)
    {
      case HID_REQ_SET_PROTOCOL:
        handle->protocol = (uint8_t)req->wValue;
        return USBD_OK;
      case HID_REQ_GET_PROTOCOL:
        USBD_CtlSendData(pdev, &handle->protocol, 1U);
        return USBD_OK;
      case HID_REQ_SET_IDLE:
        handle->idle = (uint8_t)(req->wValue >> 8);
        return USBD_OK;
      case HID_REQ_GET_IDLE:
        USBD_CtlSendData(pdev, &handle->idle, 1U);
        return USBD_OK;
      case HID_REQ_GET_REPORT:
      case HID_REQ_SET_REPORT:
      default:
        USBD_CtlError(pdev, req);
        return USBD_FAIL;
    }
  }

  if ((req->bmRequest & USB_REQ_TYPE_MASK) == USB_REQ_TYPE_STANDARD)
  {
    switch (req->bRequest)
    {
      case USB_REQ_GET_STATUS:
        USBD_CtlSendData(pdev, (uint8_t *)&status, sizeof(status));
        return USBD_OK;
      case USB_REQ_GET_INTERFACE:
        USBD_CtlSendData(pdev, &handle->altSetting[interface], 1U);
        return USBD_OK;
      case USB_REQ_SET_INTERFACE:
        if (req->wValue == 0U)
        {
          handle->altSetting[interface] = 0U;
          return USBD_OK;
        }
        break;
      case USB_REQ_GET_DESCRIPTOR:
        if ((req->wValue >> 8) == HID_REPORT_DESC)
        {
          if (interface == 0U)
          {
            descriptor = gamepadReportDescriptor;
            length = sizeof(gamepadReportDescriptor);
          }
          else if (interface == 1U)
          {
            descriptor = setupReportDescriptor;
            length = sizeof(setupReportDescriptor);
          }
          else
          {
            descriptor = rawReportDescriptor;
            length = sizeof(rawReportDescriptor);
          }
        }
        else if ((req->wValue >> 8) == HID_DESCRIPTOR_TYPE)
        {
          static const uint8_t hidOffsets[3] = {18U, 43U, 75U};
          descriptor = &dinputConfigDescriptor[hidOffsets[interface]];
          length = 9U;
        }
        if (descriptor != NULL)
        {
          USBD_CtlSendData(pdev, descriptor, MIN(length, req->wLength));
          return USBD_OK;
        }
        break;
      default:
        break;
    }
  }

  USBD_CtlError(pdev, req);
  return USBD_FAIL;
}

uint8_t USBD_DINPUT_SendReport(USBD_HandleTypeDef *pdev,
                               uint8_t *report,
                               uint16_t len)
{
  USBD_DINPUT_HandleTypeDef *handle =
      (USBD_DINPUT_HandleTypeDef *)pdev->pClassData;
  if ((pdev->dev_state == USBD_STATE_CONFIGURED) &&
      (handle != NULL) && (handle->busy == 0U) &&
      (len == DINPUT_REPORT_SIZE))
  {
    handle->busy = 1U;
    memcpy(dinputInBuffer, report, DINPUT_REPORT_SIZE);
    return USBD_LL_Transmit(pdev, DINPUT_GAMEPAD_EPIN_ADDR,
                            dinputInBuffer, len);
  }
  return USBD_OK;
}

static uint8_t DINPUT_DataIn(USBD_HandleTypeDef *pdev, uint8_t epnum)
{
  if ((epnum == (DINPUT_GAMEPAD_EPIN_ADDR & 0x0FU)) &&
      (pdev->pClassData != NULL))
  {
    ((USBD_DINPUT_HandleTypeDef *)pdev->pClassData)->busy = 0U;
    return USBD_OK;
  }
  return USBD_FAIL;
}

static uint8_t DINPUT_DataOut(USBD_HandleTypeDef *pdev, uint8_t epnum)
{
  if (epnum == (DINPUT_SETUP_EPOUT_ADDR & 0x0FU))
  {
    USBD_LL_PrepareReceive(pdev, DINPUT_SETUP_EPOUT_ADDR,
                           dinputSetupOutBuffer, DINPUT_EP_SIZE);
    return USBD_OK;
  }
  return USBD_FAIL;
}

static uint8_t *DINPUT_GetCfgDesc(uint16_t *length)
{
  *length = sizeof(dinputConfigDescriptor);
  return dinputConfigDescriptor;
}

static uint8_t *DINPUT_GetDeviceQualifierDesc(uint16_t *length)
{
  *length = sizeof(dinputDeviceQualifierDescriptor);
  return dinputDeviceQualifierDescriptor;
}

#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
static uint8_t *DINPUT_GetUserString(USBD_HandleTypeDef *pdev,
                                     uint8_t index,
                                     uint16_t *length)
{
  const char *text;
  UNUSED(pdev);
  if (index == 0x06U)
  {
    text = "Setup";
  }
  else if (index == 0x07U)
  {
    text = "Raw";
  }
  else
  {
    *length = 0U;
    return NULL;
  }
  USBD_GetString((uint8_t *)text, dinputStringBuffer, length);
  return dinputStringBuffer;
}
#endif