#include "usbd_xinput.h"

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

static uint8_t USBD_XINPUT_Init(USBD_HandleTypeDef *pdev, uint8_t cfgidx);
static uint8_t USBD_XINPUT_DeInit(USBD_HandleTypeDef *pdev, uint8_t cfgidx);
static uint8_t USBD_XINPUT_Setup(USBD_HandleTypeDef *pdev,
                                 USBD_SetupReqTypedef *req);
static uint8_t USBD_XINPUT_DataIn(USBD_HandleTypeDef *pdev, uint8_t epnum);
static uint8_t USBD_XINPUT_DataOut(USBD_HandleTypeDef *pdev, uint8_t epnum);
static uint8_t *USBD_XINPUT_GetCfgDesc(uint16_t *length);
static uint8_t *USBD_XINPUT_GetDeviceQualifierDesc(uint16_t *length);
#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
static uint8_t *USBD_XINPUT_GetUserString(USBD_HandleTypeDef *pdev,
                                          uint8_t index,
                                          uint16_t *length);
#endif

static uint8_t xinputOutBuffer[XINPUT_EP_SIZE];
static uint8_t xinputSetupOutBuffer[XINPUT_EP_SIZE];
static uint8_t xinputInBuffer[XINPUT_REPORT_SIZE];
static volatile uint8_t rumbleLeft;
static volatile uint8_t rumbleRight;
#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
static uint8_t xinputStringBuffer[32];
#endif

USBD_ClassTypeDef USBD_XINPUT =
{
  USBD_XINPUT_Init,
  USBD_XINPUT_DeInit,
  USBD_XINPUT_Setup,
  NULL,
  NULL,
  USBD_XINPUT_DataIn,
  USBD_XINPUT_DataOut,
  NULL,
  NULL,
  NULL,
  USBD_XINPUT_GetCfgDesc,
  USBD_XINPUT_GetCfgDesc,
  USBD_XINPUT_GetCfgDesc,
  USBD_XINPUT_GetDeviceQualifierDesc,
#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
  USBD_XINPUT_GetUserString,
#endif
};

/* Original three-interface layout. Only IF0 implements device behavior. */
__ALIGN_BEGIN static uint8_t USBD_XINPUT_CfgDesc[USB_XINPUT_CONFIG_DESC_SIZ] __ALIGN_END =
{
  0x09, USB_DESC_TYPE_CONFIGURATION,
  LOBYTE(USB_XINPUT_CONFIG_DESC_SIZ), HIBYTE(USB_XINPUT_CONFIG_DESC_SIZ),
  0x03, /* bNumInterfaces */
  0x01, /* bConfigurationValue */
  0x00, /* iConfiguration */
  0x80, /* Bus powered */
  0x32, /* 100 mA */

  0x09, USB_DESC_TYPE_INTERFACE,
  0x00, /* bInterfaceNumber */
  0x00, /* bAlternateSetting */
  0x02, /* bNumEndpoints */
  0xFF, /* Vendor specific */
  0x5D, /* XUSB */
  0x01, /* Xbox 360 wired controller */
  0x00, /* iInterface */

  /* XUSB class-specific descriptor from the working controller firmware. */
  0x11, 0x21, 0x10, 0x01, 0x01, 0x25, 0x81, 0x14,
  0x03, 0x03, 0x03, 0x04, 0x13, 0x02, 0x08, 0x03, 0x03,

  0x07, USB_DESC_TYPE_ENDPOINT,
  XINPUT_EPIN_ADDR,
  0x03, /* Interrupt */
  LOBYTE(XINPUT_EP_SIZE), HIBYTE(XINPUT_EP_SIZE),
  0x01,

  0x07, USB_DESC_TYPE_ENDPOINT,
  XINPUT_EPOUT_ADDR,
  0x03, /* Interrupt */
  LOBYTE(XINPUT_EP_SIZE), HIBYTE(XINPUT_EP_SIZE),
  0x01,

  /* IF1: descriptor-compatible Setup HID; data path intentionally stalled. */
  0x09, USB_DESC_TYPE_INTERFACE,
  0x01, 0x00, 0x02,
  0x03, 0x00, 0x00, 0x06,
  0x09, 0x21,
  0x11, 0x01, 0x00, 0x01, 0x22,
  0x22, 0x00,
  0x07, USB_DESC_TYPE_ENDPOINT,
  XINPUT_SETUP_EPIN_ADDR, 0x03,
  LOBYTE(XINPUT_EP_SIZE), HIBYTE(XINPUT_EP_SIZE), 0x01,
  0x07, USB_DESC_TYPE_ENDPOINT,
  XINPUT_SETUP_EPOUT_ADDR, 0x03,
  LOBYTE(XINPUT_EP_SIZE), HIBYTE(XINPUT_EP_SIZE), 0x01,

  /* IF2: descriptor-compatible Raw HID; data path intentionally stalled. */
  0x09, USB_DESC_TYPE_INTERFACE,
  0x02, 0x00, 0x01,
  0x03, 0x00, 0x00, 0x07,
  0x09, 0x21,
  0x11, 0x01, 0x00, 0x01, 0x22,
  0x15, 0x00,
  0x07, USB_DESC_TYPE_ENDPOINT,
  XINPUT_RAW_EPIN_ADDR, 0x03,
  LOBYTE(XINPUT_EP_SIZE), HIBYTE(XINPUT_EP_SIZE), 0x01,
};

__ALIGN_BEGIN static uint8_t xinputSetupReportDescriptor[] __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 xinputRawReportDescriptor[] __ALIGN_END =
{
  0x06, 0x00, 0xFF, 0x09, 0x02, 0xA1, 0x01,
  0x09, 0x20, 0x15, 0x00, 0x26, 0xFF, 0x00,
  0x75, 0x08, 0x95, 0x18, 0x81, 0x02,
  0xC0
};

_Static_assert(sizeof(xinputSetupReportDescriptor) == 34U,
               "Setup report descriptor length changed");
_Static_assert(sizeof(xinputRawReportDescriptor) == 21U,
               "Raw report descriptor length changed");

__ALIGN_BEGIN static uint8_t USBD_XINPUT_DeviceQualifierDesc[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 uint8_t USBD_XINPUT_Init(USBD_HandleTypeDef *pdev, uint8_t cfgidx)
{
  USBD_XINPUT_HandleTypeDef *handle;

  UNUSED(cfgidx);

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

  USBD_LL_OpenEP(pdev, XINPUT_EPOUT_ADDR, USBD_EP_TYPE_INTR, XINPUT_EP_SIZE);
  pdev->ep_out[XINPUT_EPOUT_ADDR & 0x0FU].is_used = 1U;

  USBD_LL_OpenEP(pdev, XINPUT_SETUP_EPIN_ADDR,
                 USBD_EP_TYPE_INTR, XINPUT_EP_SIZE);
  pdev->ep_in[XINPUT_SETUP_EPIN_ADDR & 0x0FU].is_used = 1U;
  USBD_LL_OpenEP(pdev, XINPUT_SETUP_EPOUT_ADDR,
                 USBD_EP_TYPE_INTR, XINPUT_EP_SIZE);
  pdev->ep_out[XINPUT_SETUP_EPOUT_ADDR & 0x0FU].is_used = 1U;
  USBD_LL_OpenEP(pdev, XINPUT_RAW_EPIN_ADDR,
                 USBD_EP_TYPE_INTR, XINPUT_EP_SIZE);
  pdev->ep_in[XINPUT_RAW_EPIN_ADDR & 0x0FU].is_used = 1U;

  pdev->pClassData = USBD_malloc(sizeof(USBD_XINPUT_HandleTypeDef));
  if (pdev->pClassData == NULL)
  {
    return USBD_FAIL;
  }

  handle = (USBD_XINPUT_HandleTypeDef *)pdev->pClassData;
  memset(handle, 0, sizeof(*handle));
  handle->state = XINPUT_IDLE;
  rumbleLeft = 0U;
  rumbleRight = 0U;

  USBD_LL_PrepareReceive(pdev, XINPUT_EPOUT_ADDR,
                         xinputOutBuffer, XINPUT_EP_SIZE);

  /* Stub HID data paths: unused IN endpoints naturally answer NAK. Setup OUT
   * packets are accepted and discarded so Windows can start both interfaces. */
  USBD_LL_PrepareReceive(pdev, XINPUT_SETUP_EPOUT_ADDR,
                         xinputSetupOutBuffer, XINPUT_EP_SIZE);
  return USBD_OK;
}

static uint8_t USBD_XINPUT_DeInit(USBD_HandleTypeDef *pdev, uint8_t cfgidx)
{
  UNUSED(cfgidx);

  USBD_LL_CloseEP(pdev, XINPUT_EPIN_ADDR);
  pdev->ep_in[XINPUT_EPIN_ADDR & 0x0FU].is_used = 0U;
  USBD_LL_CloseEP(pdev, XINPUT_EPOUT_ADDR);
  pdev->ep_out[XINPUT_EPOUT_ADDR & 0x0FU].is_used = 0U;
  USBD_LL_CloseEP(pdev, XINPUT_SETUP_EPIN_ADDR);
  pdev->ep_in[XINPUT_SETUP_EPIN_ADDR & 0x0FU].is_used = 0U;
  USBD_LL_CloseEP(pdev, XINPUT_SETUP_EPOUT_ADDR);
  pdev->ep_out[XINPUT_SETUP_EPOUT_ADDR & 0x0FU].is_used = 0U;
  USBD_LL_CloseEP(pdev, XINPUT_RAW_EPIN_ADDR);
  pdev->ep_in[XINPUT_RAW_EPIN_ADDR & 0x0FU].is_used = 0U;

  if (pdev->pClassData != NULL)
  {
    USBD_free(pdev->pClassData);
    pdev->pClassData = NULL;
  }

  return USBD_OK;
}

static uint8_t USBD_XINPUT_Setup(USBD_HandleTypeDef *pdev,
                                 USBD_SetupReqTypedef *req)
{
  static uint8_t capability20[20] =
  {
    0x00, 0x14, 0xF7, 0xFF, 0xFF, 0xFF, 0xFF, 0xC0, 0xFF, 0xC0,
    0xFF, 0xC0, 0xFF, 0xC0, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
  };
  static uint8_t capability8[8] =
  {
    0x00, 0x08, 0x00, 0xFF, 0xFF, 0x00, 0x00, 0x00
  };
  static uint8_t status4[4] =
  {
    0x28, 0x00, 0x00, 0x00
  };
  USBD_XINPUT_HandleTypeDef *handle =
      (USBD_XINPUT_HandleTypeDef *)pdev->pClassData;
  uint16_t statusInfo = 0U;
  uint8_t interface = (uint8_t)req->wIndex;
  uint8_t *descriptor = NULL;
  uint16_t descriptorLength = 0U;

  if ((req->bmRequest == 0xC1U) && (req->bRequest == 0x01U) &&
      (req->wValue == 0x0100U) && (req->wIndex == 0x0000U) &&
      (req->wLength == sizeof(capability20)))
  {
    USBD_CtlSendData(pdev, capability20, sizeof(capability20));
    return USBD_OK;
  }

  if ((req->bmRequest == 0xC1U) && (req->bRequest == 0x01U) &&
      (req->wValue == 0x0000U) && (req->wIndex == 0x0000U) &&
      (req->wLength == sizeof(capability8)))
  {
    USBD_CtlSendData(pdev, capability8, sizeof(capability8));
    return USBD_OK;
  }

  if ((req->bmRequest == 0xC0U) && (req->bRequest == 0x01U) &&
      (req->wValue == 0x0000U) && (req->wIndex == 0x0000U) &&
      (req->wLength == sizeof(status4)))
  {
    USBD_CtlSendData(pdev, status4, sizeof(status4));
    return USBD_OK;
  }

  /* No Microsoft OS descriptors or compatible-ID vendor request exists. */
  if ((req->bmRequest & USB_REQ_TYPE_MASK) == USB_REQ_TYPE_CLASS)
  {
    USBD_CtlError(pdev, req);
    return USBD_FAIL;
  }

  if (((req->bmRequest & USB_REQ_TYPE_MASK) == USB_REQ_TYPE_STANDARD) &&
      (interface <= 2U))
  {
    switch (req->bRequest)
    {
      case USB_REQ_GET_STATUS:
        USBD_CtlSendData(pdev, (uint8_t *)&statusInfo, sizeof(statusInfo));
        return USBD_OK;

      case USB_REQ_GET_INTERFACE:
        if (handle == NULL)
        {
          break;
        }
        USBD_CtlSendData(pdev, &handle->AltSetting[interface], 1U);
        return USBD_OK;

      case USB_REQ_SET_INTERFACE:
        if (handle == NULL)
        {
          break;
        }
        if (req->wValue == 0U)
        {
          handle->AltSetting[interface] = 0U;
          return USBD_OK;
        }
        break;

      case USB_REQ_GET_DESCRIPTOR:
        if ((interface == 1U) && ((req->wValue >> 8) == 0x22U))
        {
          descriptor = xinputSetupReportDescriptor;
          descriptorLength = sizeof(xinputSetupReportDescriptor);
        }
        else if ((interface == 2U) && ((req->wValue >> 8) == 0x22U))
        {
          descriptor = xinputRawReportDescriptor;
          descriptorLength = sizeof(xinputRawReportDescriptor);
        }
        else if (((interface == 1U) || (interface == 2U)) &&
                 ((req->wValue >> 8) == 0x21U))
        {
          descriptor = &USBD_XINPUT_CfgDesc[(interface == 1U) ? 58U : 90U];
          descriptorLength = 9U;
        }
        if (descriptor != NULL)
        {
          USBD_CtlSendData(pdev, descriptor,
                           MIN(descriptorLength, req->wLength));
          return USBD_OK;
        }
        break;

      default:
        break;
    }
  }

  USBD_CtlError(pdev, req);
  return USBD_FAIL;
}

uint8_t USBD_XINPUT_SendReport(USBD_HandleTypeDef *pdev,
                               uint8_t *report,
                               uint16_t len)
{
  USBD_XINPUT_HandleTypeDef *handle =
      (USBD_XINPUT_HandleTypeDef *)pdev->pClassData;

  if ((pdev->dev_state == USBD_STATE_CONFIGURED) &&
      (handle != NULL) && (handle->state == XINPUT_IDLE) &&
      (len == XINPUT_REPORT_SIZE))
  {
    handle->state = XINPUT_BUSY;
    memcpy(xinputInBuffer, report, XINPUT_REPORT_SIZE);
    USBD_LL_Transmit(pdev, XINPUT_EPIN_ADDR, xinputInBuffer, len);
  }

  return USBD_OK;
}

static uint8_t USBD_XINPUT_DataIn(USBD_HandleTypeDef *pdev, uint8_t epnum)
{
  if ((epnum == (XINPUT_EPIN_ADDR & 0x0FU)) &&
      (pdev->pClassData != NULL))
  {
    ((USBD_XINPUT_HandleTypeDef *)pdev->pClassData)->state = XINPUT_IDLE;
    return USBD_OK;
  }
  return USBD_FAIL;
}

static uint8_t USBD_XINPUT_DataOut(USBD_HandleTypeDef *pdev, uint8_t epnum)
{
  uint32_t length;

  if (epnum == (XINPUT_SETUP_EPOUT_ADDR & 0x0FU))
  {
    /* No setup protocol is implemented. Discard data and keep the endpoint
     * alive rather than halting the whole HID interface. */
    USBD_LL_PrepareReceive(pdev, XINPUT_SETUP_EPOUT_ADDR,
                           xinputSetupOutBuffer, XINPUT_EP_SIZE);
    return USBD_OK;
  }

  if (epnum != (XINPUT_EPOUT_ADDR & 0x0FU))
  {
    return USBD_FAIL;
  }

  length = USBD_LL_GetRxDataSize(pdev, XINPUT_EPOUT_ADDR);
  if ((length == 8U) && (xinputOutBuffer[0] == 0x00U) &&
      (xinputOutBuffer[1] == 0x08U))
  {
    rumbleLeft = xinputOutBuffer[3];
    rumbleRight = xinputOutBuffer[4];
  }

  USBD_LL_PrepareReceive(pdev, XINPUT_EPOUT_ADDR,
                         xinputOutBuffer, XINPUT_EP_SIZE);
  return USBD_OK;
}

uint8_t USBD_XINPUT_GetRumbleLeft(void)
{
  return rumbleLeft;
}

uint8_t USBD_XINPUT_GetRumbleRight(void)
{
  return rumbleRight;
}

static uint8_t *USBD_XINPUT_GetCfgDesc(uint16_t *length)
{
  *length = sizeof(USBD_XINPUT_CfgDesc);
  return USBD_XINPUT_CfgDesc;
}

static uint8_t *USBD_XINPUT_GetDeviceQualifierDesc(uint16_t *length)
{
  *length = sizeof(USBD_XINPUT_DeviceQualifierDesc);
  return USBD_XINPUT_DeviceQualifierDesc;
}

#if (USBD_SUPPORT_USER_STRING_DESC == 1U)
static uint8_t *USBD_XINPUT_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, xinputStringBuffer, length);
  return xinputStringBuffer;
}
#endif