#include "usbd_dfu_bootloader.h"

#include "stm32f1xx_hal.h"
#include "usbd_ctlreq.h"

#define DFU_DETACH          0x00U
#define DFU_DNLOAD          0x01U
#define DFU_UPLOAD          0x02U
#define DFU_GETSTATUS       0x03U
#define DFU_CLRSTATUS       0x04U
#define DFU_GETSTATE        0x05U
#define DFU_ABORT           0x06U

#define DFU_STATUS_OK       0x00U
#define DFU_STATUS_ERR_WRITE 0x03U
#define DFU_STATUS_ERR_ERASE 0x04U
#define DFU_STATUS_ERR_ADDRESS 0x08U
#define DFU_STATUS_ERR_STALLEDPKT 0x0FU

#define DFU_STATE_DFU_IDLE            0x02U
#define DFU_STATE_DFU_DNLOAD_SYNC     0x03U
#define DFU_STATE_DFU_DNLOAD_IDLE     0x05U
#define DFU_STATE_DFU_MANIFEST_SYNC   0x06U
#define DFU_STATE_DFU_MANIFEST_WAIT_RESET 0x08U
#define DFU_STATE_DFU_ERROR           0x0AU

#define DFU_DESCRIPTOR_TYPE 0x21U
#define STM32F103_FLASH_PAGE_SIZE 1024U

typedef struct
{
  uint8_t state;
  uint8_t status;
  uint8_t altSetting;
  uint8_t resetPending;
  uint16_t blockNumber;
  uint16_t downloadLength;
  uint32_t resetAt;
} USBD_DFU_BootloaderHandleTypeDef;

static uint8_t DFU_Init(USBD_HandleTypeDef *pdev, uint8_t cfgidx);
static uint8_t DFU_DeInit(USBD_HandleTypeDef *pdev, uint8_t cfgidx);
static uint8_t DFU_Setup(USBD_HandleTypeDef *pdev, USBD_SetupReqTypedef *req);
static uint8_t DFU_EP0RxReady(USBD_HandleTypeDef *pdev);
static uint8_t DFU_SOF(USBD_HandleTypeDef *pdev);
static uint8_t *DFU_GetCfgDesc(uint16_t *length);
static uint8_t *DFU_GetDeviceQualifierDesc(uint16_t *length);
static uint8_t DFU_WriteBlock(USBD_DFU_BootloaderHandleTypeDef *handle);
static uint32_t DFU_GetSlotBase(uint8_t altSetting);

__ALIGN_BEGIN static uint32_t dfuDownloadWords[DFU_TRANSFER_SIZE / 4U] __ALIGN_END;

USBD_ClassTypeDef USBD_DFU_BOOTLOADER =
{
  DFU_Init,
  DFU_DeInit,
  DFU_Setup,
  NULL,
  DFU_EP0RxReady,
  NULL,
  NULL,
  DFU_SOF,
  NULL,
  NULL,
  DFU_GetCfgDesc,
  DFU_GetCfgDesc,
  DFU_GetCfgDesc,
  DFU_GetDeviceQualifierDesc,
};

/* One DFU interface with two alternate settings: App1 and App2. */
__ALIGN_BEGIN static uint8_t dfuConfigDescriptor[USB_DFU_CONFIG_DESC_SIZ] __ALIGN_END =
{
  0x09, USB_DESC_TYPE_CONFIGURATION,
  LOBYTE(USB_DFU_CONFIG_DESC_SIZ), HIBYTE(USB_DFU_CONFIG_DESC_SIZ),
  0x01, 0x01, 0x00, 0x80, 0x32,

  0x09, USB_DESC_TYPE_INTERFACE,
  0x00, 0x00, 0x00,
  0xFE, 0x01, 0x02, 0x00,

  /* DFU 1.1 functional descriptor: download only, 1024-byte transfer. */
  0x09, DFU_DESCRIPTOR_TYPE,
  0x01,
  0xE8, 0x03,
  LOBYTE(DFU_TRANSFER_SIZE), HIBYTE(DFU_TRANSFER_SIZE),
  0x10, 0x01,

  0x09, USB_DESC_TYPE_INTERFACE,
  0x00, 0x01, 0x00,
  0xFE, 0x01, 0x02, 0x00,
};

__ALIGN_BEGIN static uint8_t dfuDeviceQualifierDescriptor[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 DFU_Init(USBD_HandleTypeDef *pdev, uint8_t cfgidx)
{
  USBD_DFU_BootloaderHandleTypeDef *handle;

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

  handle = (USBD_DFU_BootloaderHandleTypeDef *)pdev->pClassData;
  handle->state = DFU_STATE_DFU_IDLE;
  handle->status = DFU_STATUS_OK;
  handle->altSetting = 0U;
  handle->resetPending = 0U;
  handle->blockNumber = 0U;
  handle->downloadLength = 0U;
  handle->resetAt = 0U;
  return USBD_OK;
}

static uint8_t DFU_DeInit(USBD_HandleTypeDef *pdev, uint8_t cfgidx)
{
  UNUSED(cfgidx);
  if (pdev->pClassData != NULL)
  {
    USBD_free(pdev->pClassData);
    pdev->pClassData = NULL;
  }
  return USBD_OK;
}

static uint8_t DFU_Setup(USBD_HandleTypeDef *pdev, USBD_SetupReqTypedef *req)
{
  USBD_DFU_BootloaderHandleTypeDef *handle =
      (USBD_DFU_BootloaderHandleTypeDef *)pdev->pClassData;
  static uint8_t statusResponse[6];
  uint16_t statusInfo = 0U;

  if (handle == NULL)
  {
    USBD_CtlError(pdev, req);
    return USBD_FAIL;
  }

  if ((req->bmRequest & USB_REQ_TYPE_MASK) == USB_REQ_TYPE_CLASS)
  {
    switch (req->bRequest)
    {
      case DFU_DNLOAD:
        if ((handle->state != DFU_STATE_DFU_IDLE) &&
            (handle->state != DFU_STATE_DFU_DNLOAD_IDLE))
        {
          break;
        }
        if (req->wLength == 0U)
        {
          handle->state = DFU_STATE_DFU_MANIFEST_SYNC;
          return USBD_OK;
        }
        if (req->wLength > DFU_TRANSFER_SIZE)
        {
          handle->status = DFU_STATUS_ERR_STALLEDPKT;
          handle->state = DFU_STATE_DFU_ERROR;
          break;
        }
        handle->blockNumber = req->wValue;
        handle->downloadLength = req->wLength;
        USBD_CtlPrepareRx(pdev, (uint8_t *)dfuDownloadWords, req->wLength);
        return USBD_OK;

      case DFU_UPLOAD:
        break;

      case DFU_GETSTATUS:
        if (handle->state == DFU_STATE_DFU_DNLOAD_SYNC)
        {
          handle->state = DFU_STATE_DFU_DNLOAD_IDLE;
        }
        else if (handle->state == DFU_STATE_DFU_MANIFEST_SYNC)
        {
          handle->state = DFU_STATE_DFU_MANIFEST_WAIT_RESET;
          handle->resetPending = 1U;
          handle->resetAt = HAL_GetTick() + 100U;
        }
        statusResponse[0] = handle->status;
        statusResponse[1] = 0U;
        statusResponse[2] = 0U;
        statusResponse[3] = 0U;
        statusResponse[4] = handle->state;
        statusResponse[5] = 0U;
        USBD_CtlSendData(pdev, statusResponse, sizeof(statusResponse));
        return USBD_OK;

      case DFU_CLRSTATUS:
        handle->status = DFU_STATUS_OK;
        handle->state = DFU_STATE_DFU_IDLE;
        return USBD_OK;

      case DFU_GETSTATE:
        USBD_CtlSendData(pdev, &handle->state, 1U);
        return USBD_OK;

      case DFU_ABORT:
        handle->status = DFU_STATUS_OK;
        handle->state = DFU_STATE_DFU_IDLE;
        handle->downloadLength = 0U;
        return USBD_OK;

      case DFU_DETACH:
        handle->resetPending = 1U;
        handle->resetAt = HAL_GetTick() + 100U;
        return USBD_OK;

      default:
        break;
    }
  }
  else if ((req->bmRequest & USB_REQ_TYPE_MASK) == USB_REQ_TYPE_STANDARD)
  {
    switch (req->bRequest)
    {
      case USB_REQ_GET_STATUS:
        USBD_CtlSendData(pdev, (uint8_t *)&statusInfo, sizeof(statusInfo));
        return USBD_OK;

      case USB_REQ_GET_INTERFACE:
        USBD_CtlSendData(pdev, &handle->altSetting, 1U);
        return USBD_OK;

      case USB_REQ_SET_INTERFACE:
        if (req->wValue <= 1U)
        {
          handle->altSetting = (uint8_t)req->wValue;
          handle->state = DFU_STATE_DFU_IDLE;
          handle->status = DFU_STATUS_OK;
          return USBD_OK;
        }
        break;

      case USB_REQ_GET_DESCRIPTOR:
        if ((req->wValue >> 8) == DFU_DESCRIPTOR_TYPE)
        {
          USBD_CtlSendData(pdev, &dfuConfigDescriptor[18],
                           MIN(9U, req->wLength));
          return USBD_OK;
        }
        break;

      default:
        break;
    }
  }

  USBD_CtlError(pdev, req);
  return USBD_FAIL;
}

static uint8_t DFU_EP0RxReady(USBD_HandleTypeDef *pdev)
{
  USBD_DFU_BootloaderHandleTypeDef *handle =
      (USBD_DFU_BootloaderHandleTypeDef *)pdev->pClassData;

  if ((handle == NULL) || (handle->downloadLength == 0U))
  {
    return USBD_OK;
  }

  if (DFU_WriteBlock(handle) == USBD_OK)
  {
    handle->status = DFU_STATUS_OK;
    handle->state = DFU_STATE_DFU_DNLOAD_SYNC;
  }
  else
  {
    handle->state = DFU_STATE_DFU_ERROR;
  }
  handle->downloadLength = 0U;
  return USBD_OK;
}

static uint8_t DFU_SOF(USBD_HandleTypeDef *pdev)
{
  USBD_DFU_BootloaderHandleTypeDef *handle =
      (USBD_DFU_BootloaderHandleTypeDef *)pdev->pClassData;

  if ((handle != NULL) && (handle->resetPending != 0U) &&
      ((int32_t)(HAL_GetTick() - handle->resetAt) >= 0))
  {
    NVIC_SystemReset();
  }
  return USBD_OK;
}

static uint8_t DFU_WriteBlock(USBD_DFU_BootloaderHandleTypeDef *handle)
{
  FLASH_EraseInitTypeDef erase = {0};
  uint32_t eraseError = 0U;
  uint32_t slotBase = DFU_GetSlotBase(handle->altSetting);
  uint32_t address = slotBase + ((uint32_t)handle->blockNumber * DFU_TRANSFER_SIZE);
  uint32_t slotEnd = slotBase + DFU_APP_SLOT_SIZE;
  uint32_t index;

  if ((address < slotBase) || (address >= slotEnd) ||
      (handle->downloadLength > (slotEnd - address)))
  {
    handle->status = DFU_STATUS_ERR_ADDRESS;
    return USBD_FAIL;
  }

  HAL_FLASH_Unlock();
  __HAL_FLASH_CLEAR_FLAG(FLASH_FLAG_EOP | FLASH_FLAG_PGERR | FLASH_FLAG_WRPERR);

  erase.TypeErase = FLASH_TYPEERASE_PAGES;
  erase.PageAddress = address;
  erase.NbPages = 1U;
  if (HAL_FLASHEx_Erase(&erase, &eraseError) != HAL_OK)
  {
    handle->status = DFU_STATUS_ERR_ERASE;
    HAL_FLASH_Lock();
    return USBD_FAIL;
  }

  for (index = 0U; index < handle->downloadLength; index += 2U)
  {
    uint8_t *bytes = (uint8_t *)dfuDownloadWords;
    uint16_t halfword = bytes[index];
    if ((index + 1U) < handle->downloadLength)
    {
      halfword |= (uint16_t)((uint16_t)bytes[index + 1U] << 8U);
    }
    else
    {
      halfword |= 0xFF00U;
    }

    if (HAL_FLASH_Program(FLASH_TYPEPROGRAM_HALFWORD,
                          address + index, halfword) != HAL_OK)
    {
      handle->status = DFU_STATUS_ERR_WRITE;
      HAL_FLASH_Lock();
      return USBD_FAIL;
    }
  }

  HAL_FLASH_Lock();
  return USBD_OK;
}

static uint32_t DFU_GetSlotBase(uint8_t altSetting)
{
  return (altSetting == 0U) ? DFU_APP1_ADDRESS : DFU_APP2_ADDRESS;
}

static uint8_t *DFU_GetCfgDesc(uint16_t *length)
{
  *length = sizeof(dfuConfigDescriptor);
  return dfuConfigDescriptor;
}

static uint8_t *DFU_GetDeviceQualifierDesc(uint16_t *length)
{
  *length = sizeof(dfuDeviceQualifierDescriptor);
  return dfuDeviceQualifierDescriptor;
}