#include "cusb.h"
#include "../addrmap.h"
#include "../bitmap.h"
#include "../interrupts.h"
#include "../gpio.h"
#include "proto.h"

u8 xferbuff[64]; // Free to use for transfer (IN)
u8 *xferptr;
u8 xferremain=0;
DeviceConfiguration *cusb_dc;

void cusb_init(DeviceConfiguration *dc){

  // Reset USB controller
  REG_WRITE_BITMAP_CLEAR(RESETS_RESET, 1<<24); // Reset USBCTRL
  while((REG_READ(RESETS_DONE) & 1<<24) == 0){} // Wait controller reset

  // Flush USB Controller DPSRAM
  memset((void*)USBCTRL_DPSRAM_BASE, 0, USBCTRL_DPSRAM_SIZE);

  // Enable usb controller interrupt
  interrupts_enable(5);

  // Setup controller
  REG_WRITE(USBCTRL_MUXING, BIT_USBCTRL_MUXING_TO_PHY|BIT_USBCTRL_MUXING_SOFTCON);
  REG_WRITE(USBCTRL_PWR, BIT_USBCTRL_PWR_VBUS_DETECT|BIT_USBCTRL_PWR_VBUS_DETECT_OVERRIDE_EN);
  REG_WRITE(USBCTRL_MAINCTRL, BIT_USBCTRL_MAIN_CTRL_CONTROLLER_EN); // Set device mode + enable controller
  REG_WRITE(USBCTRL_SIE_CTRL, BIT_USBCTRL_SIE_CTRL_EP0_INT_1BUF); // Enable usb interrupts on EP0
  REG_WRITE(USBCTRL_INTE, BIT_USBCTRL_INTE_SETUP_REQ|BIT_USBCTRL_INTE_BUS_RESET|BIT_USBCTRL_INTE_BUFF_STATUS); // Setup interrupt granularity

  // Initialize device configuration
  cusb_dc=dc; // Store controller device config
  cusb_init_device_configuration();

  // Initialize endpoints
  for(int i=0;i<USB_ENDPOINT_COUNT;i++){
	EndPointConfiguration *ep=&(cusb_dc->endpoints[i]);
	u32 dir=ep->endpoint_descriptor.bEndpointAddress & 0x80;
	u32 id=ep->endpoint_descriptor.bEndpointAddress & 0xF;
	// Endpoint and buffer control
	u32 endpoint_control=USBCTRL_EP1_ENDP_CTRL_IN+0x8*(id-1); // -1 since start at 1 (ep0 already covered in cusb_init_device_configuration())
	u32 buffer_control=USBCTRL_EP1_BUFFER_CTRL_IN+0x8*(id-1);
	if(dir == USB_DIRECTION_OUT){
	  endpoint_control+=4;
	  buffer_control+=4;
	}
	ep->endpoint_control=(void*)(endpoint_control);
	ep->buffer_control=(void*)(buffer_control);
	// Init buffers
	u32 buffer0=USBCTRL_DATA_BUFFER_START+(64*2)*(id-1); // 64 bits aligned warning see p388
	u32 buffer1=buffer0+64;
	ep->buffer0=(void*)buffer0;
	ep->buffer1=(void*)buffer1;
	// Init ep control
	u32 epctrl=BIT_USBCTRL_ENDP_CTRL_ENABLE |
	  BIT_USBCTRL_ENDP_CTRL_INT_PER_BUFF1 |
	  ep->endpoint_descriptor.bmAttributes << BIT_USBCTRL_ENDP_CTRL_TYPE_LSB|
	  USB_BUFFER_OFFSET(buffer0);
	*ep->endpoint_control=epctrl; // Apply buffer control setting
	ep->next_pid=0; // Maybe that is good?
  }

  // Pull-up usb line! (for the host :)
  REG_WRITE_BITMAP_SET(USBCTRL_SIE_CTRL, BIT_USBCTRL_SIE_CTRL_PULLUP_EN);
}

void cusb_init_device_configuration(){

  // Init device descriptor
  cusb_dc->device_descriptor.bLength=sizeof(USB_DEVICE_DESCRIPTOR);
  cusb_dc->device_descriptor.bDescriptorType=USB_DESCRIPTOR_TYPE_DEVICE;
  cusb_dc->device_descriptor.bNumConfigurations = 1; // We support only 1
  cusb_dc->device_descriptor.bcdUSB=BCD_USB;
  cusb_dc->device_descriptor.bMaxPacketSize0=64;

  // Init configuration descriptor
  cusb_dc->configuration_descriptor.bLength=sizeof(USB_CONFIGURATION_DESCRIPTOR);
  cusb_dc->configuration_descriptor.bDescriptorType=USB_DESCRIPTOR_TYPE_CONFIGURATION;
  cusb_dc->configuration_descriptor.bConfigurationValue = 1; // Configuration id
  cusb_dc->configuration_descriptor.iConfiguration = 0; // No string
  cusb_dc->configuration_descriptor.bmAttributes = 0xC0; // attributes: self powered, no remote wakeup
  cusb_dc->configuration_descriptor.bMaxPower = 0x32; // 100ma

  // Init device state
  cusb_dc->devaddr=0; // Just in case
  cusb_dc->setdevaddr=0; // Just in case
  cusb_dc->setdevaddr=0; // Just in case

  // Init string zero descriptor
  cusb_dc->supported_languages.bLength = sizeof(USB_STRING_DESCRIPTOR_ZERO);
  cusb_dc->supported_languages.bDescriptorType = USB_DESCRIPTOR_TYPE_STRING;

  // Init ep0 in
  EndPointConfiguration *ep0_in= &(cusb_dc->endpoints[USB_ENDPOINT_COUNT]);
  ep0_in->buffer0=(void*)USBCTRL_EP0_BUFFER0;
  ep0_in->buffer1=(void*)USBCTRL_EP0_BUFFER0+0x40;
  ep0_in->buffer_control=(void*)USBCTRL_EP0_BUFFER_CTRL_IN;
  ep0_in->handler=cusb_ep0_in_handler;
  ep0_in->endpoint_descriptor.bEndpointAddress = B_ENDPOINT_ADDRESS(0, USB_DIRECTION_IN);
  ep0_in->endpoint_descriptor.bmAttributes = USB_TRANSFERT_TYPE_CONTROL;
  ep0_in->endpoint_descriptor.wMaxPacketSize=64;

  // Init ep0 out
  EndPointConfiguration *ep0_out= &(cusb_dc->endpoints[USB_ENDPOINT_COUNT+1]);
  ep0_out->buffer0=(void*)USBCTRL_EP0_BUFFER0;
  ep0_out->buffer1=(void*)USBCTRL_EP0_BUFFER0+0x40;
  ep0_out->buffer_control=(void*)USBCTRL_EP0_BUFFER_CTRL_OUT;
  ep0_out->handler=cusb_ep0_out_handler;
  ep0_out->endpoint_descriptor.bEndpointAddress = B_ENDPOINT_ADDRESS(0, USB_DIRECTION_OUT);
  ep0_out->endpoint_descriptor.bmAttributes = USB_TRANSFERT_TYPE_CONTROL;
  ep0_out->endpoint_descriptor.wMaxPacketSize=64;

  // Init bLength
  for(char i=0;i<USB_ENDPOINT_COUNT+2;i++){
	cusb_dc->endpoints[i].endpoint_descriptor.bLength=sizeof(USB_ENDPOINT_DESCRIPTOR);
	cusb_dc->endpoints[i].endpoint_descriptor.bDescriptorType=USB_DESCRIPTOR_TYPE_ENDPOINT;
  }

  // Init interfaces
  for(char i=0;i<USB_INTERFACE_COUNT;i++){
	cusb_dc->interface_descriptors[i].bLength=sizeof(USB_INTERFACE_DESCRIPTOR);
	cusb_dc->interface_descriptors[i].bDescriptorType=USB_DESCRIPTOR_TYPE_INTERFACE;
  }
}

int cusb_check_interrupt(int int_bit){
  if(REG_READ(USBCTRL_INTS) & int_bit)
	return 1;
  return 0;
}

void cusb_handle_bus_reset(void){
  // https://github.com/raspberrypi/pico-examples/blob/master/usb/device/dev_lowlevel/dev_lowlevel.c#L469
  REG_WRITE(USBCTRL_ADDR_ENDP, 0); // Remove device address from controller
  cusb_dc->devaddr=0; // No more device address
  cusb_dc->setdevaddr=0; // Ensure no device address will be set
  cusb_dc->configured=0;
}

EndPointConfiguration* cusb_get_endpoint(char num, u32 direction){
  EndPointConfiguration *ep;
  for(char i=0;i<(USB_ENDPOINT_COUNT+2);i++){
	ep=&(cusb_dc->endpoints[i]);
	u32 bEndpointAddress=ep->endpoint_descriptor.bEndpointAddress;
	// Bit 7 (mask 0x80) is IN or OUT and first 4 bits is addr (see p269)
	if((bEndpointAddress & 0x80) == direction && (bEndpointAddress & 0xF) == num)
	  return ep;
  }

  return 0;
}

void cusb_handle_setup(void){
  UD_SETUP *pkt=(UD_SETUP*)USBCTRL_DPSRAM_SETUP_PACKET;

  // Always responds with DATA1 PID
  EndPointConfiguration *ep=cusb_get_endpoint(0, USB_DIRECTION_IN);
  ep->next_pid=1;

  u8 direction=pkt->bmRequestType & BM_REQUEST_TYPE_DIR_BIT;
  if(direction == USB_DIRECTION_OUT){
	if (pkt->bRequest == USB_REQUEST_CODE_SET_ADDRESS){
	  cusb_dc->devaddr= pkt->wValue & 0xff;
	  cusb_dc->setdevaddr=1;
	  // Since we should acknowledge (status phase) before setting the address
	  // we use setdevaddr boolean. When status done, buffer_status interrupt will be triggered
	  // and ep0_in_handler will set the address
	  cusb_status_xfer(USB_DIRECTION_IN);
	} else if (pkt->bRequest == USB_REQUEST_CODE_SET_CONFIGURATION) {
	  cusb_status_xfer(USB_DIRECTION_IN);
	  cusb_dc->configured=1;
	} else {
	  if(cusb_dc->setup_handler){
		cusb_dc->setup_handler(pkt);
	  }
	  // Acknowledge whatever other requests
	  cusb_status_xfer(USB_DIRECTION_IN);
	}
  } else if(direction == USB_DIRECTION_IN){
	if(pkt->bRequest == USB_REQUEST_CODE_GET_DESCRIPTOR){
	  int desc_type=pkt->wValue>>8; // See USB SPecification (wValue contains descriptor type+index)
	  if(desc_type == USB_DESCRIPTOR_TYPE_DEVICE ){
		// Send descriptor
		cusb_start_xfer(&cusb_dc->device_descriptor, sizeof(USB_DEVICE_DESCRIPTOR), ep);
	  } else if(desc_type == USB_DESCRIPTOR_TYPE_CONFIGURATION ){

		USB_CONFIGURATION_DESCRIPTOR *conf=cusb_dc->full_configuration_descriptor;
		int size=conf->bLength;

		if(pkt->wLength > size)
		  size=conf->wTotalLength;

		// Send descriptors!!
		cusb_start_xfer(conf, size, ep);
	  } else if(desc_type == USB_DESCRIPTOR_TYPE_STRING ){
		u8 id = pkt->wValue & 0xff; // Get string id
		if(id==0){ // This is the special string descriptor for supported language
		  cusb_start_xfer(&(cusb_dc->supported_languages), sizeof(USB_STRING_DESCRIPTOR_ZERO), ep);
		}
		else {
		  char *str=cusb_dc->descriptor_strings[id-1]; // Remember id 0 taken by ZERO DESCriptor
		  u8 *ptr=xferbuff;
		  USB_UNICODE_STRING_DESCRIPTOR *u=(USB_UNICODE_STRING_DESCRIPTOR*)ptr;
		  u->bLength = sizeof(USB_UNICODE_STRING_DESCRIPTOR);
		  u->bDescriptorType=USB_DESCRIPTOR_TYPE_STRING;
		  char c;
		  ptr+=sizeof(USB_UNICODE_STRING_DESCRIPTOR); // String first 2 descriptor entries
		  do {
			c = *str++;
			*ptr++ = c;
			*ptr++ = 0; // Unicode
			u->bLength+=2; // Unicode
		  } while(c != '\0');
		  cusb_start_xfer(xferbuff, u->bLength, ep);
		}
	  }
	  else {
		cusb_dc->setup_handler(pkt);
	  }
	}
  }
}

void cusb_handle_buffer_status(void){
  u32 status=REG_READ(USBCTRL_BUFF_STATUS);
  for(u8 num=0;num<16;num++){
	for(u8 i=0;i<2;i++){
	  u32 bit=1u << (num*2+i); // Shift register for IN and OUT for endpoint num
	  if(status & bit){
		u32 dir=i ? USB_DIRECTION_OUT : USB_DIRECTION_IN;
		EndPointConfiguration *ep=cusb_get_endpoint(num,dir);
		REG_WRITE_BITMAP_CLEAR(USBCTRL_BUFF_STATUS, bit); // Clear buffer status
		ep->handler((u8*)ep->buffer0,*(ep->buffer_control)&MASK_USBCTRL_BUFFER0_LENGTH);
	  }
	}
  }
}

void cusb_status_xfer(u32 dir){
  // If dir == USB_DIRECTION_OUT
  // controller will receive the out token from host then send acknowledgement
  // otherwise if we do not perform xfer when receiving out token, then controller do not send acknowledgement which do not complete the control transaction
  // If dir == USB_DIRECTION_IN
  // controller will receive the in token from host then send zlp and receive acknowledgement from host
  EndPointConfiguration *ep=cusb_get_endpoint(0, dir);
  cusb_start_xfer(0,0,ep);
}

void cusb_ep0_in_handler(u8 *buffer, int len) {
  // This function is trigger when buff_status interrupt is handled
  // it is called inside  cusb_handle_buffer_status
  if(cusb_dc->setdevaddr){
	REG_WRITE(USBCTRL_ADDR_ENDP, cusb_dc->devaddr);
	cusb_dc->setdevaddr=0;
  }
  else if(xferremain>0){
	EndPointConfiguration *ep=cusb_get_endpoint(0, USB_DIRECTION_IN);
	cusb_start_xfer(xferptr,xferremain,ep);
  }
  else {
	cusb_status_xfer(USB_DIRECTION_OUT); // Acknowledge with zlp when setup transaction ends
  }
}

void cusb_ep0_out_handler(u8 *buffer, int len) {
  // This function is trigger when buff_status interrupt is handled
  // it is called inside  cusb_handle_buffer_status
}



void cusb_start_xfer(void* data, u32 size, EndPointConfiguration *ep){
  u32 buffer_ctrl = size; // Set data size
  buffer_ctrl|=ep->next_pid ? BIT_USBCTRL_BUFFER0_PID : 0; // Set DATA0 or DATA1
  ep->next_pid ^= 1u; // For next transfert

  // Move user data to usb buffer if needed
  u32 direction=ep->endpoint_descriptor.bEndpointAddress & 0x80;
  if(direction == USB_DIRECTION_IN){
	memcpy((void *)ep->buffer0, data, size);
	buffer_ctrl |= BIT_USBCTRL_BUFFER0_FULL; // Set buffer full for controller
	// Support for data > 64 bytes (see also cusb_ep0_in_handler())
	if(size>64){
	  xferremain=size-64;
	  xferptr=((u8*) data)+64;
	}
	else {
	  xferremain=0;
	}
  }

  // Setup buffer (not available bit yet)
  *(ep->buffer_control)=buffer_ctrl;

  // Now do available bit (because USB and cores have different clock speed (race condition see p389)
  asm volatile("nop;nop;nop;"); // ceil(125MHz/48MHz) Should wait 3 cycles see warning p392
  *(ep->buffer_control)=buffer_ctrl| BIT_USBCTRL_BUFFER0_AVAILABLE; // Apply available bit

}

void cusb_eoi(){
  if(cusb_dc->eoi)
	cusb_dc->eoi();
}

void cusb_isr_trigger(){
  REG_WRITE(PPB_NVIC_ISPR, 1u<<5); // Trigger interrupt 5 (USB Controller)
}