#include "usb.h"
#include "usb_device.h"
#include "registers.h"
#include <stm32h747xx.h>

typedef union {
  uint32_t word;
  uint16_t halfwords[2];
  uint8_t bytes[4];
} usb_data_t;

uint32_t get_mask(uint8_t byte_count) {
  switch (byte_count) {
  case 0:
    return 0;
  case 1:
    return 0xFF;
  case 2:
    return 0xFFFF;
  case 3:
    return 0xFFFFFF;
  default:
    return ~0;
  }
}

void usb_generic_fill_fifo_words(USB_OTG_INEndpointTypeDef *endpoint,
                      uint8_t *data,
                      uint16_t size,
                      volatile uint32_t *fifo_tx_target,
                      uint32_t *sub_word_data,
                      uint8_t *sub_word_bytes) {
  usb_data_t tx_data;
  uint8_t total_bytes = size + *sub_word_bytes;
  if (total_bytes < 4) {
    uint8_t bytes = *sub_word_bytes;
    *sub_word_bytes = total_bytes;

    uint32_t mask = get_mask(bytes);
    tx_data.word = *sub_word_data & mask;

    for (int i = 0; i < size; i++) {
      tx_data.bytes[i + bytes] = *(data++);
    }

    *sub_word_data = tx_data.word;
    return;
  }

  if (*sub_word_bytes > 0) {
    // first send these bytes,
    uint8_t skip_bytes = *sub_word_bytes;
    tx_data.word = *sub_word_data;
    for (int i = 0; i < 4 - *sub_word_bytes; i++) {
      tx_data.bytes[skip_bytes + i] = *(data++);
      size--;
    }

    *fifo_tx_target = tx_data.word;
  }

  uint8_t subWordBytes = size % 4;
  uint8_t wordCount = size / 4;

  for (uint8_t i = 0; i < wordCount; i++) {
    tx_data.word = *((uint32_t*)data);
    *fifo_tx_target = tx_data.word;
    data += 4;
  }

  // NOTE: hm. This is not generally safe, but it's tempting...
  *sub_word_data = *((uint32_t*)data);
  *sub_word_bytes = subWordBytes;
}


void usb_generic_fill_fifo(USB_OTG_INEndpointTypeDef *endpoint,
                      uint8_t *data,
                      uint16_t size,
                      volatile uint32_t *fifo_tx_target) {
  usb_data_t txData;
  uint8_t subWordBytes = size % 4;
  uint8_t wordCount = size / 4;

  for (uint8_t i = 0; i < wordCount; i++) {
    txData.word = *((uint32_t*)data);
    *fifo_tx_target = txData.word;
    data += 4;
  }

  if (subWordBytes > 0) {
    txData.word = 0;
    for (int i = 0; i < subWordBytes; i++) {
      txData.bytes[i] = *(data++);
      size--;
    }

    *fifo_tx_target = txData.word;
  }
}

task_result_t usb_generic_setup_in_endpoint(USB_OTG_INEndpointTypeDef *endpoint, uint16_t size, uint16_t max_packet_size) {
  uint16_t packet_count = (size + max_packet_size - 1) / max_packet_size;

  if (usb_is_inendpoint_ready(endpoint)) {
    return RES_WOULD_BLOCK;
  }

  if (!usb_check_fifo_space(endpoint, size)) {
    return RES_WOULD_BLOCK;
  }

  if (size > 0) {
    endpoint->DIEPTSIZ = (packet_count << USB_OTG_DIEPTSIZ_PKTCNT_Pos) | (size << USB_OTG_DIEPTSIZ_XFRSIZ_Pos);
  } else { // just one packet, of zero length
    endpoint->DIEPTSIZ = (1 << USB_OTG_DIEPTSIZ_PKTCNT_Pos);
  }
  endpoint->DIEPCTL |= USB_OTG_DIEPCTL_CNAK | USB_OTG_DIEPCTL_EPENA;

  return RES_OK;
}

task_result_t usb_generic_send(USB_OTG_INEndpointTypeDef *endpoint,
                      uint8_t *data,
                      uint16_t size,
                      volatile uint32_t *fifo_tx_target) {
  task_result_t result = usb_generic_setup_in_endpoint(endpoint, size, 64);
  if (result != RES_OK) {
    return result;
  }

  if (size > 0) {
    usb_generic_fill_fifo(endpoint, data, size, fifo_tx_target);
    // After the fifo is filled...
    /* endpoint->DIEPCTL |= USB_OTG_DIEPCTL_CNAK; */
  }

  return RES_OK;
}

task_result_t usb_generic_read(uint8_t *data, uint16_t size, volatile uint32_t *fifo_rx_source) {
  usb_data_t rx_data;
  uint8_t subWordBytes = size % 4;
  uint8_t wordCount = size / 4;

  for (uint8_t i = 0; i < wordCount; i++) {
    rx_data.word = *fifo_rx_source;
    *((uint32_t*)data) = rx_data.word;
    data += 4;
  }

  // TODO: check where to put it... beginning, end...
  if (subWordBytes > 0) {
    rx_data.word = *fifo_rx_source;
    for (uint8_t i = 0; i < subWordBytes; i++) {
      *(data + i) = rx_data.bytes[i];
    }
  }

  return RES_OK;
}

task_result_t usb_send_descriptor(USB_OTG_INEndpointTypeDef *endpoint,
                         usb_descriptor_t *descriptor,
                         uint16_t max_size,
                         volatile uint32_t *fifo_tx_target) {
  uint16_t size = descriptor->bLength;
  if (max_size != 0 && size > max_size) {
    size = max_size;
  }

  return usb_generic_send(endpoint, (uint8_t*)descriptor, size, fifo_tx_target);
}

task_result_t usb_send_device_descriptor(USB_OTG_INEndpointTypeDef *endpoint,
                                usb_device_descriptor_t *descriptor,
                                uint16_t max_size,
                                volatile uint32_t *fifo_tx_target) {
  descriptor->header.bDescriptorType = DESCRIPTOR_DEVICE;
  descriptor->header.bLength = sizeof(usb_device_descriptor_t);
  return usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}
task_result_t usb_send_device_qualifier_descriptor(USB_OTG_INEndpointTypeDef *endpoint,
                                          usb_device_qualifier_t *descriptor,
                                          uint16_t max_size,
                                   volatile uint32_t *fifo_tx_target) {
  descriptor->header.bDescriptorType = DESCRIPTOR_DEVICE_QUALIFIER;
  descriptor->header.bLength = sizeof(usb_device_qualifier_t);
  return usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}
task_result_t usb_send_interface_descriptor(USB_OTG_INEndpointTypeDef *endpoint,
                                   usb_interface_descriptor_t *descriptor,
                                   uint16_t max_size,
                                   volatile uint32_t *fifo_tx_target) {
  descriptor->header.bDescriptorType = DESCRIPTOR_INTERFACE;
  descriptor->header.bLength = sizeof(usb_interface_descriptor_t);
  return usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}
task_result_t usb_send_endpoint_descriptor(USB_OTG_INEndpointTypeDef *endpoint,
                                  usb_endpoint_descriptor_t *descriptor,
                                   uint16_t max_size,
                                  volatile uint32_t *fifo_tx_target) {
  descriptor->header.bDescriptorType = DESCRIPTOR_ENDPOINT;
  descriptor->header.bLength = sizeof(usb_endpoint_descriptor_t);
  return usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}

task_result_t usb_send_string_descriptor_zero(USB_OTG_INEndpointTypeDef *endpoint,
                                     usb_string_descriptor_zero_t *string_descriptor,
                                     uint16_t max_size,
                              volatile uint32_t *fifo_tx_target) {
  usb_data_t data;
  data.word = 0;

  task_result_t result = usb_generic_setup_in_endpoint(endpoint, string_descriptor->header.bLength, 64);
  if (result != RES_OK) {
    return result;
  }

  data.bytes[0] = string_descriptor->header.bLength;
  data.bytes[1] = string_descriptor->header.bDescriptorType;

  if (string_descriptor->header.bLength == 2) {
    usb_generic_fill_fifo(endpoint, &data.bytes[0], 2, fifo_tx_target);
    return RES_OK;
  }

  data.halfwords[1] = string_descriptor->wLANGID[0];

  usb_generic_fill_fifo(endpoint, &data.bytes[0], 4, fifo_tx_target);

  if (string_descriptor->header.bLength > 4) {
    usb_generic_fill_fifo(endpoint, (uint8_t*)&string_descriptor->wLANGID[1], string_descriptor->header.bLength - 4, fifo_tx_target);
  }

  /* endpoint->DIEPCTL |= USB_OTG_DIEPCTL_CNAK; */

  return RES_OK;
}

task_result_t usb_send_unicode_string_descriptor(
    USB_OTG_INEndpointTypeDef *endpoint,
    usb_unicode_string_descriptor_t *string_descriptor,
    uint16_t max_size,
    volatile uint32_t *fifo_tx_target) {

  task_result_t result = usb_generic_setup_in_endpoint(endpoint, string_descriptor->header.bLength, 64);
  if (result != RES_OK) {
    return result;
  }

  usb_data_t data = { .word = 0 };

  data.bytes[0] = string_descriptor->header.bLength;
  data.bytes[1] = string_descriptor->header.bDescriptorType;

  if (string_descriptor->bString != 0 && string_descriptor->header.bLength > 2) {
    // NOTE: if the string had length of just one, there would be two bytes
    // read here. That shouldn't usually matter much, so we don't care about
    // that here.
    data.halfwords[1] = *((uint16_t*)string_descriptor->bString);
  }
  usb_generic_fill_fifo(endpoint, &data.bytes[0], 4, fifo_tx_target);

  if (string_descriptor->header.bLength > 4) {
    usb_generic_fill_fifo(endpoint, &string_descriptor->bString[2], string_descriptor->header.bLength - 4, fifo_tx_target);
  }

  /* endpoint->DIEPCTL |= USB_OTG_DIEPCTL_CNAK; */
  return RES_OK;
}

bool usb_is_inendpoint_ready(USB_OTG_INEndpointTypeDef *endpoint) {
  return (endpoint->DIEPCTL & USB_OTG_DIEPCTL_EPENA) != 0;
}

bool usb_check_fifo_space(USB_OTG_INEndpointTypeDef *endpoint, uint16_t size) {
  return ((endpoint->DTXFSTS & (USB_OTG_DTXFSTS_INEPTFSAV_Msk)) >> USB_OTG_DTXFSTS_INEPTFSAV_Pos) >= (size + 3) / 4;
}

void usb_flush_tx_fifo(USB_OTG_GlobalTypeDef *core, uint16_t fifo_num) {
  reg_write_bits_pos(&core->GRSTCTL, fifo_num, USB_OTG_GRSTCTL_TXFNUM_Pos, 0xF);
  reg_write_bits_pos(&core->GRSTCTL, fifo_num, USB_OTG_GRSTCTL_TXFFLSH_Pos, 1);
  while (reg_read_bits_pos(&core->GRSTCTL, USB_OTG_GRSTCTL_TXFFLSH_Pos, 1));
}
void usb_flush_rx_fifo(USB_OTG_GlobalTypeDef *core) {
  reg_set_bits_pos(&core->GRSTCTL, USB_OTG_GRSTCTL_RXFFLSH_Pos, 1);
  while (reg_read_bits_pos(&core->GRSTCTL, USB_OTG_GRSTCTL_TXFFLSH_Pos, 1));
}