#include "usb.h"
#include "usb_device.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) {
  size = ((size + 3) / 4) * 4;

  uint32_t subWordData;
  uint8_t subWordCount;

  usb_generic_fill_fifo_words(endpoint, data, size, fifo_tx_target, &subWordData, &subWordCount);
}

void 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 (endpoint->DIEPCTL & USB_OTG_DIEPCTL_EPENA) {
    // this is bad! Can't send the packet, this shouldn't get here, ever.
    while (endpoint->DIEPCTL & USB_OTG_DIEPCTL_EPENA);
    /* __nop(); */
  }

  if (size > 0) {
    endpoint->DIEPTSIZ = (1 << USB_OTG_DIEPTSIZ_MULCNT_Pos) | (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;
}

void usb_generic_send(USB_OTG_INEndpointTypeDef *endpoint,
                      uint8_t *data,
                      uint16_t size,
                      volatile uint32_t *fifo_tx_target) {
  if (endpoint->DIEPCTL & USB_OTG_DIEPCTL_EPENA) {
    // this is bad! Can't send the packet, this shouldn't get here, ever.
    while (endpoint->DIEPCTL & USB_OTG_DIEPCTL_EPENA);
    /* __nop(); */
  }

  uint32_t fifo_size;
  while ((fifo_size = endpoint->DTXFSTS & USB_OTG_DTXFSTS_INEPTFSAV) < (size + 3) / 4) {
    /* __nop(); */
  }

  // TODO: generic max packet size
  usb_generic_setup_in_endpoint(endpoint, size, 64);

  usb_data_t tx_data;
  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;
  }

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

  // After the fifo is filled...
  endpoint->DIEPCTL = USB_OTG_DIEPCTL_CNAK;
}

void 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 = 0;
    for (uint8_t i = 0; i < subWordBytes; i++) {
      *(data + i) = rx_data.bytes[i];
    }
  }
}

void 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;
  }

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

/* void usb_send_configuration_descriptor(USB_OTG_INEndpointTypeDef* endpoint, */
/*                                        usb_class_t *class, */
/*                                        volatile uint32_t *fifo_tx_target) { */
/*   class->configuration_descriptor.header.bDescriptorType = DESCRIPTOR_CONFIGURATION; */
/*   class->configuration_descriptor.header.bLength = sizeof(usb_configuration_descriptor_t); */
/*   class->interface_descriptor.header.bDescriptorType = DESCRIPTOR_INTERFACE; */
/*   class->interface_descriptor.header.bLength = sizeof(usb_interface_descriptor_t); */
/*   uint16_t total_length = */
/*     sizeof(usb_configuration_descriptor_t) + */
/*     sizeof(usb_interface_descriptor_t) + */
/*     class->interface_descriptor.bNumEndpoints * sizeof(usb_endpoint_descriptor_t); */

/*   for (int i = 0; i < class->interface_descriptor.bNumEndpoints; i++) { */
/*     class->endpoint_descriptors[i].header.bDescriptorType = DESCRIPTOR_ENDPOINT; */
/*     class->endpoint_descriptors[i].header.bLength = sizeof(usb_endpoint_descriptor_t); */
/*   } */

/*   // NOTE: since the memory layout is: configuration, interface, endpoints, */
/*   // we can send directly like this. */
/*   usb_generic_send(endpoint, */
/*                    (uint8_t*)&class->configuration_descriptor, */
/*                    total_length, fifo_tx_target); */
/* } */

void 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);
  usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}
void 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);
  usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}
void 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);
  usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}
void 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);
  usb_send_descriptor(endpoint, &descriptor->header, max_size, fifo_tx_target);
}

void 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;

  usb_generic_setup_in_endpoint(endpoint, string_descriptor->header.bLength, 64);

  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;
  }

  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;
}
void 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) {

  usb_generic_setup_in_endpoint(endpoint, string_descriptor->header.bLength, 64);

  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;
}