#include "usb.h"
#include <stm32l0xx.h>
#include <stddef.h>

#define ENDPOINT_COUNT (8)

enum endpoint_state {EP_STATE_DISABLED, EP_STATE_NAK, EP_STATE_STALL, EP_STATE_VALID};

static void (*ep_rx_data_callback)(uint8_t endpoint, const uint8_t *buffer, uint32_t len) = NULL;
static enum control_state (*ep_rx_setup_received_callback)(uint8_t endpoint, const struct setup_packet *setup_pkg) = NULL;
static void (*ep_tx_complete_callback)(uint8_t endpoint);
static void (*usb_sof_callback)(void) = NULL;
static void (*usb_reset_callback)(void) = NULL;
static void (*usb_configured_callback)(uint16_t config_idx) = NULL;

struct usb_endpoint_info {
	uint16_t rx_pma_size;
	uint16_t tx_pma_size;
	enum usb_ep_type type;
	struct setup_packet last_setup;
	enum control_state ctrl_state;
	volatile uint16_t *ep_reg;

	uint32_t tx_count;
	const uint8_t *tx_ptr;

	uint32_t rx_size;
	uint8_t *rx_buffer;
};

static const struct usb_descriptor_entry *usb_descriptors = NULL;
static uint16_t * const pma_base_ptr = (uint16_t *)0x40006000UL;

struct usb_endpoint_info endpoints[ENDPOINT_COUNT];

static void setup_usb_clock(void)
{
	uint32_t reload_val;
	const uint32_t felim = 0x22;
	const uint32_t f_target = 48000000UL;
	const uint32_t f_usb_sof = 1000UL;

	RCC->APB2ENR |= RCC_APB2ENR_SYSCFGEN;
	/* Enable the internal voltage reference. This is needed for HSI48 */
	SYSCFG->CFGR3 |= SYSCFG_CFGR3_EN_VREFINT;
	while (!(SYSCFG->CFGR3 & SYSCFG_CFGR3_VREFINT_RDYF));
	SYSCFG->CFGR3 |= SYSCFG_CFGR3_ENREF_HSI48;

	RCC->CRRCR = RCC_CRRCR_HSI48ON;
	while(!(RCC->CRRCR & RCC_CRRCR_HSI48RDY));
	RCC->CCIPR |= RCC_CCIPR_HSI48SEL;

	/* Configure Clock recovery from USB SOF (1kHz) */
	RCC->APB1ENR |= RCC_APB1ENR_CRSEN;
	reload_val = (f_target / f_usb_sof) - 1;
	CRS->CFGR = (reload_val & 0xFFFF) | (felim << 16) | CRS_CFGR_SYNCSRC_1;
	CRS->CR = CRS_CR_CEN | CRS_CR_AUTOTRIMEN;
}

void usb_init(const struct usb_callbacks *callbacks)
{
	setup_usb_clock();

	/* Activate register clock for USB */
	RCC->APB1ENR |= RCC_APB1ENR_USBEN;

	/* Power up the USB interface but hold in reset */
	USB->CNTR = USB_CNTR_FRES;

	endpoints[0].ep_reg = &USB->EP0R;
	endpoints[1].ep_reg = &USB->EP1R;
	endpoints[2].ep_reg = &USB->EP2R;
	endpoints[3].ep_reg = &USB->EP3R;
	endpoints[4].ep_reg = &USB->EP4R;
	endpoints[5].ep_reg = &USB->EP5R;
	endpoints[6].ep_reg = &USB->EP6R;
	endpoints[7].ep_reg = &USB->EP7R;

	/* Setup the callbacks */
	ep_rx_data_callback = callbacks->ep_rx_data_callback;
	ep_tx_complete_callback = callbacks->ep_tx_complete_callback;
	ep_rx_setup_received_callback = callbacks->ep_rx_setup_received_callback;
	usb_configured_callback = callbacks->usb_configured_callback;
	usb_reset_callback = callbacks->usb_reset_callback;
	usb_sof_callback = callbacks->usb_sof_callback;
}

static void write_b_table(uint32_t addr, uint16_t value)
{
	volatile uint16_t *ptr = pma_base_ptr;

	ptr += (addr >> 1);
	*ptr = value;
}

static uint16_t read_b_table(uint32_t addr)
{
	volatile uint16_t *ptr = pma_base_ptr;

	ptr += (addr >> 1);
	return *ptr;
}

static uint16_t read_rx_pma(uint8_t ep, uint8_t *buff, uint32_t size)
{
	uint16_t rx_btable_entry;
	uint16_t pma_addr;
	uint8_t *app_addr;
	uint16_t count;
	uint16_t i;

	rx_btable_entry = read_b_table(ep * 8 + 6);
	count = rx_btable_entry & 0x3FF;

	pma_addr = read_b_table(ep * 8 + 4);
	app_addr = (uint8_t *)pma_base_ptr;
	app_addr += (pma_addr >> 0);

	/* Byte addresed writes seem to work just fine */
	for (i = 0; i < count && i < size; i++) {
		buff[i] = app_addr[i];
	}

	return count;
}

static void write_tx_pma(uint8_t ep, const uint8_t *src, uint32_t len)
{
	uint16_t tx_addr;
	uint16_t *pma = pma_base_ptr;
	uint32_t i;
	uint32_t halfword_cnt;
	bool padding;

	write_b_table(ep * 8 + 2, len);

	if (!len)
		return;

	tx_addr = read_b_table(ep * 8);
	padding = len % 2 ? true : false;
	halfword_cnt = len >> 1;

	pma += tx_addr >> 1;

	for (i = 0; i < halfword_cnt; i++) {
		pma[i] = (uint16_t)src[2 * i] | (((uint16_t)src[2 * i + 1]) << 8);
	}

	if (padding) {
		pma[i] = (uint16_t)src[2 * i];
	}

}

static const struct usb_descriptor_entry *find_descriptor_from_setup_request(uint16_t w_index, uint16_t w_value)
{
	int i;
	const struct usb_descriptor_entry *entry = NULL;

	for (i = 0; usb_descriptors[i].descriptor != NULL; i++) {
		if (usb_descriptors[i].w_index == w_index && usb_descriptors[i].w_value == w_value) {
			entry = &usb_descriptors[i];
			break;
		}
	}

	return entry;
}

void usb_endpoint_config(enum usb_ep_type type, uint8_t epnum, bool rx, bool tx)
{
	uint16_t epreg = 0;
	struct usb_endpoint_info *ep_info;

	if (epnum >= ENDPOINT_COUNT)
		return;

	ep_info = &endpoints[epnum];

	write_b_table(epnum * 8 + 0, 64 + epnum * 136);
	write_b_table(epnum * 8 + 2, 0);
	write_b_table(epnum * 8 + 4, 64 + epnum * 136 + 72);
	write_b_table(epnum * 8 + 6, (1 << 10) | (1<<15));

	epreg = *ep_info->ep_reg;
	epreg &= ~(1<<15);

	switch (type) {
	case EP_CONTROL:
		epreg |= (1<<9);
		ep_info->type = EP_CONTROL;
		ep_info->ctrl_state = CONTROL_SETUP;
		ep_info->tx_pma_size = 64;
		ep_info->rx_pma_size = 64;
	case EP_BULK:
		ep_info->type = EP_BULK;
		ep_info->ctrl_state = CONTROL_SETUP;
		ep_info->tx_pma_size = 64;
		ep_info->rx_pma_size = 64;
		break;
	case EP_INTERRUPT:
		epreg |= (1<<10) | (1<<9);
		ep_info->type = EP_INTERRUPT;
		ep_info->ctrl_state = CONTROL_SETUP;
		ep_info->tx_pma_size = 64;
		ep_info->rx_pma_size = 64;
		break;
	case EP_ISOCHRON:
		epreg |= (1<<10);
		ep_info->type = EP_ISOCHRON;
		ep_info->ctrl_state = CONTROL_SETUP;
		ep_info->tx_pma_size = 64;
		ep_info->rx_pma_size = 64;
		break;
	default:
		return;
		break;
	}

	if (tx) {
		epreg |= (1<<5);
	}

	if (rx) {
		epreg |= (1<<13);
	}

	epreg |= epnum;

	*ep_info->ep_reg = epreg;
}

static void usb_endpoint_reset_int_flags(uint8_t endpoint, bool rx, bool tx)
{
	uint16_t ep_reg;

	if (endpoint >= ENDPOINT_COUNT)
		return;

	ep_reg = *endpoints[endpoint].ep_reg;
	ep_reg &= ~((1<<14) | (1<<6) | (1<<5) | (1<<13) | (1<<12) | (1<<4));

	/* Dont clear flag */
	if (rx) {
		ep_reg &= ~(1<<15);
	}  else {
		ep_reg |= (1<<15);
	}
	if (tx) {
		ep_reg &= ~(1<<7);
	} else {
		ep_reg |= (1<<7);
	}


	*endpoints[endpoint].ep_reg = ep_reg;
}

static void usb_endpoint_set_rx_state(uint8_t endpoint, enum endpoint_state state)
{
	uint16_t ep_reg;
	uint16_t ep_rx_target = 0;

	if (endpoint >= ENDPOINT_COUNT)
		return;

	ep_reg = *endpoints[endpoint].ep_reg;

	/* Mask out all toggle flags that we don't want to change */
	ep_reg &= ~((1<<14) | (1<<6) | (1<<5) | (1<<4));

	/* Dont clear flag */
	ep_reg |= (1<<15) | (1<<7);

	/* Prepare the target value */
	switch (state) {
	case EP_STATE_DISABLED:
		ep_rx_target = 0;
		break;
	case EP_STATE_NAK:
		ep_rx_target = (1<<13);
		break;
	case EP_STATE_VALID:
		ep_rx_target = (1<<13) | (1<<12);
		break;
	case EP_STATE_STALL:
		ep_rx_target = (1<<12);
	default:
		return;
	}

	/* Generate toggle mask for the TX bits */
	ep_reg ^= ep_rx_target;

	*endpoints[endpoint].ep_reg = ep_reg;
}

static void usb_endpoint_set_tx_state(uint8_t endpoint, enum endpoint_state state)
{
	uint16_t ep_reg;
	uint16_t ep_tx_target = 0;

	if (endpoint >= ENDPOINT_COUNT)
		return;

	ep_reg = *endpoints[endpoint].ep_reg;

	/* Mask out all toggle flags that we don't want to change */
	ep_reg &= ~((1<<14) | (1<<13) | (1<<12) | (1<<6));

	/* Dont clear flag */
	ep_reg |= (1<<15) | (1<<7);

	/* Prepare the target value */
	switch (state) {
	case EP_STATE_DISABLED:
		ep_tx_target = 0;
		break;
	case EP_STATE_NAK:
		ep_tx_target = (1<<5);
		break;
	case EP_STATE_VALID:
		ep_tx_target = (1<<5) | (1<<4);
		break;
	case EP_STATE_STALL:
		ep_tx_target = (1<<4);
	default:
		return;
	}

	/* Generate toggle mask for the TX bits */
	ep_reg ^= ep_tx_target;

	*endpoints[endpoint].ep_reg = ep_reg;
}

void usb_endpoint_stall(uint8_t endpoint, bool tx_dir, bool rx_dir)
{
	if (tx_dir)
		usb_endpoint_set_tx_state(endpoint, EP_STATE_STALL);
}


int usb_endpoint_send(uint8_t endpoint, const uint8_t *data, uint32_t len)
{
	struct usb_endpoint_info *epinfo;
	uint32_t tx_cnt;

	if (endpoint >= ENDPOINT_COUNT)
		return -1001;

	epinfo = &endpoints[endpoint];
	if (epinfo->tx_pma_size == 0)
		return -1;

	tx_cnt = len < epinfo->tx_pma_size ? len : epinfo->tx_pma_size;

	/* Copy data to pma and save the data pointer to the next word to be transmitted */
	if (tx_cnt < len) {
		epinfo->tx_ptr = &data[tx_cnt];
		epinfo->tx_count = len - tx_cnt;
	} else {
		epinfo->tx_count = 0;
		epinfo->tx_ptr = NULL;
	}

	write_tx_pma(endpoint, data, tx_cnt);

	/* Prepare endpoint for TX */
	usb_endpoint_set_tx_state(endpoint, EP_STATE_VALID);

	return 0;
}

int usb_endpoint_prepare_receive(uint8_t endpoint, uint8_t *buffer, uint32_t bufflen)
{
	struct usb_endpoint_info *info;

	if (endpoint >= ENDPOINT_COUNT)
		return -1001;

	info = &endpoints[endpoint];

	if (info->rx_buffer || info->rx_size) {
		return -1;
	}

	__disable_irq();
	info->rx_buffer = buffer;
	info->rx_size = bufflen;
	usb_endpoint_set_rx_state(endpoint, EP_STATE_VALID);
	__enable_irq();

	return 0;
}

void usb_endpoint_send_status_stage(uint8_t endpoint)
{
	struct usb_endpoint_info *info;

	if (endpoint >= ENDPOINT_COUNT)
		return;

	info = &endpoints[endpoint];
	info->ctrl_state = CONTROL_STATUS_TX;
	usb_endpoint_send(endpoint, NULL, 0);
}

void usb_enable(const struct usb_descriptor_entry *descriptors)
{
	if (!descriptors)
		return;

	usb_descriptors = descriptors;

	/* Enable device with address 0 */
	USB->DADDR = USB_DADDR_EF;

	/* Clear all interrupts */
	USB->ISTR = 0;

	/* Set buffer description table to beginning of PMA */
	USB->BTABLE = 0x0U;

	/* Actiovate USB module (clear reset bit) and setup the interrupt masks */
	USB->CNTR = USB_CNTR_CTRM | USB_CNTR_ERRM | USB_CNTR_RESETM | USB_CNTR_SOFM;

	/* Enable internal D+ pullup to tell the host we're here */
	USB->BCDR = USB_BCDR_DPPU;

	NVIC_EnableIRQ(USB_IRQn);
}


static void usb_handle_reset(void)
{
	int i;

	/* Reset USB address */
	USB->DADDR = USB_DADDR_EF;

	for (i = 0; i < ENDPOINT_COUNT; i++) {
		endpoints[i].rx_size = 0;
		endpoints[i].rx_buffer = NULL;
		endpoints[i].tx_count = 0;
		endpoints[i].tx_ptr = NULL;
	}

	usb_endpoint_config(EP_CONTROL, 0, true, true);
	usb_endpoint_set_rx_state(0, EP_STATE_VALID);
}

void usb_ep0_send_status(void)
{
	struct usb_endpoint_info *ep_info;

	ep_info = &endpoints[0];
	ep_info->ctrl_state = CONTROL_STATUS_TX;

	usb_endpoint_send(0, NULL, 0);
}

void usb_ep0_handle_setup(void)
{
	uint16_t cnt;
	enum control_state next_state;
	const struct usb_descriptor_entry *descriptor;
	struct usb_endpoint_info *ep_info;
	struct setup_packet *setup;

	ep_info = &endpoints[0];
	ep_info->ctrl_state = CONTROL_SETUP;

	setup = &ep_info->last_setup;

	cnt = read_rx_pma(0, (uint8_t *)setup, 8);

	if (cnt != 8) {
		usb_endpoint_stall(0, true, true);
		return;
	}

	if ((setup->bm_req_type == 0x80 || setup->bm_req_type == 0x81 || setup->bm_req_type == 0x82) && setup->b_request == 6) {
		/* Get descriptor request */
		descriptor = find_descriptor_from_setup_request(setup->w_index, setup->w_value);
		if (!descriptor) {
			usb_endpoint_stall(0, true, true);
		}

		cnt = setup->w_length < descriptor->size ? setup->w_length : descriptor->size;

		usb_endpoint_send(0, descriptor->descriptor, cnt);
		ep_info->ctrl_state = CONTROL_DATA_TX;
	} else if (setup->bm_req_type == 0x00 && setup->b_request == 5) {
		/* Set address command
		 * Send out status stage and set address after status has finished
		 */
		usb_ep0_send_status();
	} else if (setup->bm_req_type == 0x00 && setup->b_request == 9) {
		/* Set configuration */
		if (usb_configured_callback)
			usb_configured_callback(setup->w_value);
		usb_ep0_send_status();
	} else {
		next_state = CONTROL_NOT_HANDLED;

		/* Check if the callback for setup requests is set and try to handle it this way */
		if (ep_rx_setup_received_callback)
			next_state = ep_rx_setup_received_callback(0, setup);

		/* We could not handle the setup request if handle state != 0 */
		if (next_state == CONTROL_NOT_HANDLED) {
			usb_endpoint_stall(0, true, true);
		} else {
			/* Control handled */
			switch (next_state) {
			case CONTROL_STATUS_TX:
				usb_ep0_send_status();
				break;
			case CONTROL_DATA_TX:
				break;
			case CONTROL_DATA_RX:
				break;
			default:
				usb_endpoint_stall(0, true, true);
				break;
			}
		}

		ep_info->ctrl_state = (next_state != CONTROL_NOT_HANDLED ? next_state : CONTROL_SETUP);
	}

	return;
}

void usb_ep0_handle_tx(void)
{
	struct usb_endpoint_info *info = &endpoints[0];

	if (info->tx_count) {
		/* Do the rest */
		usb_endpoint_send(0, info->tx_ptr, info->tx_count);
	} else {
		if (info->ctrl_state == CONTROL_DATA_TX) {
			info->ctrl_state = CONTROL_STATUS_RX;
			/* Listen for new setup packets */
			usb_endpoint_set_rx_state(0, EP_STATE_VALID);
		} else if (info->ctrl_state == CONTROL_STATUS_TX) {
			/* we've sent a status back. Check if we should set our address now */
			if (info->last_setup.bm_req_type == 0x00 && info->last_setup.b_request == 0x5) {
				USB->DADDR = USB_DADDR_EF | info->last_setup.w_value;
			}

			info->ctrl_state = CONTROL_SETUP;
			/* Listen for new setup packets */
			usb_endpoint_set_rx_state(0, EP_STATE_VALID);
		} else {
			if (ep_tx_complete_callback)
				ep_tx_complete_callback(0);
		}
	}
}

void usb_handle_rx_packet(uint8_t ep)
{
	uint16_t epreg;
	struct usb_endpoint_info *info;
	uint16_t pkg_len;
	uint8_t *buffer;

	info = &endpoints[ep];
	epreg = *info->ep_reg;

	usb_endpoint_reset_int_flags(ep, true, false);


	if (epreg & (1<<11) && ep == 0) {
		info->ctrl_state = CONTROL_SETUP;
		usb_ep0_handle_setup();
	} else if (ep == 0 && info->ctrl_state == CONTROL_STATUS_RX) {
		info->ctrl_state = CONTROL_SETUP;

		/* Check the received status frame */
		pkg_len = read_rx_pma(0, NULL, 0);
		if (pkg_len != 0) {
			usb_endpoint_stall(0, true, true);
		} else {
			/* Ready to receive next datum */
			usb_endpoint_set_rx_state(0, EP_STATE_VALID);
		}		
	} else {
		if (epreg & (1<<11)) {
			(void)read_rx_pma(ep, (uint8_t *)&info->last_setup, 8);
			if (ep_rx_setup_received_callback)
				ep_rx_setup_received_callback(ep, &info->last_setup);
		}

		if (info->rx_buffer && info->rx_size) {
			pkg_len = read_rx_pma(ep, info->rx_buffer, info->rx_size);
			buffer = info->rx_buffer;
			info->rx_size = 0;
			info->rx_buffer = 0;
			ep_rx_data_callback(ep, buffer, pkg_len);
		}

	}
}

void usb_handle_tx_packet_sent(uint8_t ep)
{
	struct usb_endpoint_info *info = &endpoints[ep];

	usb_endpoint_reset_int_flags(ep, false, true);

	if (ep == 0) {
		usb_ep0_handle_tx();
	} else {
		if (info->tx_count && info->tx_ptr) {
			usb_endpoint_send(ep, info->tx_ptr, info->tx_count);
		} else {
			if (ep_tx_complete_callback)
				ep_tx_complete_callback(ep);
		}
	}
}

void USB_IRQHandler(void)
{
	uint16_t istr = USB->ISTR;
	uint8_t endpoint;
	bool dir_tx;

	if (istr & USB_ISTR_RESET) {
		USB->ISTR = (uint16_t)(~USB_ISTR_RESET);
		usb_handle_reset();
		if (usb_reset_callback)
			usb_reset_callback();
	}

	if (istr & USB_ISTR_CTR) {
		/* correct transfer interrupt */
		endpoint = istr & USB_ISTR_EP_ID;
		dir_tx = (istr & USB_ISTR_DIR) ? false : true;

		if (dir_tx) {
			usb_handle_tx_packet_sent(endpoint);
		} else {
			usb_handle_rx_packet(endpoint);
		}

		/* Clear interrupt */
		USB->ISTR = (uint16_t)(~USB_ISTR_CTR);

	}

	if (istr & USB_ISTR_SOF) {
		USB->ISTR = (uint16_t)(~USB_ISTR_SOF);
		if (usb_sof_callback)
			usb_sof_callback();
	}

	if (istr & USB_ISTR_ERR) {
		USB->ISTR = (uint16_t)(~USB_ISTR_ERR);
	}

	__DSB();
}