#include "project.h"

/* Commands sent with wBlockNum == 0 as per ST implementation. */
#define CMD_SETADDR 0x21
#define CMD_ERASE 0x41


/* We need a special large control buffer for this device: */

static enum dfu_state usbdfu_state = STATE_DFU_IDLE;

static uint32_t sector_bases[] = {
  0x08000000,
  0x08004000,
  0x08008000,
  0x0800C000,
  0x08010000,
  0x08020000,
  0x08040000,
  0x08060000,
  0x08080000,
  0x080A0000,
  0x080C0000,
  0x080E0000,
};

#define N_SECTORS (sizeof(sector_bases)/sizeof(sector_bases[0]))

static struct {
  uint8_t buf[sizeof (usbd_control_buffer)];
  uint16_t len;
  uint32_t addr;
  uint16_t blocknum;
} prog;

const struct usb_dfu_descriptor dfu_function = {
  .bLength = sizeof (struct usb_dfu_descriptor),
  .bDescriptorType = DFU_FUNCTIONAL,
  .bmAttributes = USB_DFU_CAN_DOWNLOAD | USB_DFU_WILL_DETACH,
  .wDetachTimeout = 255,
  .wTransferSize = 1024,
  .bcdDFUVersion = 0x011A,
};

const struct usb_interface_descriptor dfu_iface = {
  .bLength = USB_DT_INTERFACE_SIZE,
  .bDescriptorType = USB_DT_INTERFACE,
  .bInterfaceNumber = 0,
  .bAlternateSetting = 0,
  .bNumEndpoints = 0,
  .bInterfaceClass = 0xFE, /* Device Firmware Upgrade */
  .bInterfaceSubClass = 1,
  .bInterfaceProtocol = 2,

  /* The ST Microelectronics DfuSe application needs this string.
   * The format isn't documented... */
  .iInterface = 4,

  .extra = &dfu_function,
  .extralen = sizeof (dfu_function),
};

static uint8_t usbdfu_getstatus (usbd_device *usbd_dev, uint32_t *bwPollTimeout)
{
  (void)usbd_dev;

  switch (usbdfu_state) {
  case STATE_DFU_DNLOAD_SYNC:
    usbdfu_state = STATE_DFU_DNBUSY;
    *bwPollTimeout = 100;
    return DFU_STATUS_OK;

  case STATE_DFU_MANIFEST_SYNC:
    /* Device will reset when read is complete. */
    usbdfu_state = STATE_DFU_MANIFEST;
    return DFU_STATUS_OK;

  default:
    return DFU_STATUS_OK;
  }
}

static void usbdfu_getstatus_complete (usbd_device *usbd_dev, struct usb_setup_data *req)
{
  unsigned i;
  (void)req;
  (void)usbd_dev;

  switch (usbdfu_state) {
  case STATE_DFU_DNBUSY:
    flash_unlock();

    if (prog.blocknum == 0) {
      switch (prog.buf[0]) {
      case CMD_ERASE: {
        uint32_t *dat = (uint32_t *) (prog.buf + 1);

        for (i = 0; i < N_SECTORS; ++i) {
          if (*dat == sector_bases[i])
            flash_erase_sector (i, FLASH_CR_PROGRAM_X32);
        }
      }

      /*Fall through*/
      case CMD_SETADDR: {
        uint32_t *dat = (uint32_t *) (prog.buf + 1);
        prog.addr = *dat;
      }
      }
    } else {
      uint32_t baseaddr = prog.addr + ((prog.blocknum - 2) *
                                       dfu_function.wTransferSize);

      for (i = 0; i < prog.len; i += 4) {
        uint32_t *dat = (uint32_t *) (prog.buf + i);
        flash_program_word (baseaddr + i, *dat);
      }

      max7219_report_addr (baseaddr);
    }

    flash_lock();

    /* Jump straight to dfuDNLOAD-IDLE, skipping dfuDNLOAD-SYNC. */
    usbdfu_state = STATE_DFU_DNLOAD_IDLE;
    return ;

  case STATE_DFU_MANIFEST:
    /* USB device must detach, we just reset... */
    scb_reset_system();
    return ; /* Will never return. */

  default:
    ;
  }
}

enum usbd_request_return_codes usbdfu_control_request (usbd_device *usbd_dev, struct usb_setup_data *req, uint8_t **buf,
    uint16_t *len, usbd_control_complete_callback *complete)
{

  if ((req->bmRequestType & 0x7F) != 0x21)
    return USBD_REQ_NEXT_CALLBACK; /* Only accept class request. */

  switch (req->bRequest) {
  case DFU_DNLOAD:
    if ((len == NULL) || (*len == 0)) {
      usbdfu_state = STATE_DFU_MANIFEST_SYNC;
      return USBD_REQ_HANDLED;
    } else {
      /* Copy download data for use on GET_STATUS. */
      prog.blocknum = req->wValue;
      prog.len = *len;
      memcpy (prog.buf, *buf, *len);
      usbdfu_state = STATE_DFU_DNLOAD_SYNC;
      return USBD_REQ_HANDLED;
    }

  case DFU_CLRSTATUS:

    /* Clear error and return to dfuIDLE. */
    if (usbdfu_state == STATE_DFU_ERROR)
      usbdfu_state = STATE_DFU_IDLE;

    return USBD_REQ_HANDLED;

  case DFU_ABORT:
    /* Abort returns to dfuIDLE state. */
    usbdfu_state = STATE_DFU_IDLE;
    return USBD_REQ_HANDLED;

  case DFU_UPLOAD:
    /* Upload not supported for now. */
    return USBD_REQ_NOTSUPP;

  case DFU_GETSTATUS: {
    uint32_t bwPollTimeout = 0; /* 24-bit integer in DFU class spec */
    (*buf)[0] = usbdfu_getstatus (usbd_dev, &bwPollTimeout);
    (*buf)[1] = bwPollTimeout & 0xFF;
    (*buf)[2] = (bwPollTimeout >> 8) & 0xFF;
    (*buf)[3] = (bwPollTimeout >> 16) & 0xFF;
    (*buf)[4] = usbdfu_state;
    (*buf)[5] = 0; /* iString not used here */
    *len = 6;
    *complete = usbdfu_getstatus_complete;
    return USBD_REQ_HANDLED;
  }

  case DFU_GETSTATE:
    /* Return state with no state transision. */
    *buf[0] = usbdfu_state;
    *len = 1;
    return USBD_REQ_HANDLED;
  }

  return USBD_REQ_NEXT_CALLBACK;
}