#include "stm32.h"
#include "rcc.h"
#include "interrupt.h"
#include "thread.h"
#include "time.h"

#include "pin.h"

#include "usb.h"
#include "i2c.h"

static Pin& led_status = PA4;
static Pin& led_error = PC4;

static Pin& usb_vbus = PB13;
static Pin& usb_dm = PB14;
static Pin& usb_dp = PB15;

static Pin& jtag_tdo = PC3;
static Pin& jtag_tms = PC13;
static Pin& jtag_tck = PC14;
static Pin& jtag_tdi = PC15;

static Pin& i2c_scl = PB10;
static Pin& i2c_sda = PB11;

bool jtag_tick(bool tdi, bool tms) {
	bool tdo = jtag_tdo.get();
	jtag_tdi.set(tdi);
	jtag_tms.set(tms);
	jtag_tck.on();
	for(uint32_t i = 0; i < 1000; i++) {
		asm volatile("nop");
	}
	jtag_tck.off();
	for(uint32_t i = 0; i < 1000; i++) {
		asm volatile("nop");
	}
	return tdo;
}

void usb_write(uint32_t ep, uint32_t* bufp, uint32_t len) {
	OTG_HS.dev_iep_reg[ep].DIEPTSIZ = (1 << 19) | len;
	//                               PKTCNT
	OTG_HS.dev_iep_reg[ep].DIEPCTL |= (1 << 31) | (1 << 26);
	//                               EPENA       CNAK
	
	len = (len + 3) >> 2;
	
	while(len--) {
		OTG_HS.fifo[ep].reg = *bufp++;
	}
}

bool usb_ep_ready(uint32_t ep) {
	return (OTG_HS.dev_iep_reg[1].DIEPCTL & 0x80008000) == 0x8000;
}

uint8_t dev_desc[] = {
	0x12, 0x01, 0x00, 0x02, 0xff, 0x00, 0x00, 0x40, 0x34, 0x12, 0x78, 0x56, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01,
};

uint8_t conf_desc[] = {
	0x09, 0x02, 0x19, 0x00, 0x01, 0x01, 0x00, 0xc0, 0x00,
	0x09, 0x04, 0x00, 0x00, 0x01, 0xff, 0x00, 0x00, 0x00,
	0x07, 0x05, 0x81, 0x02, 0x40, 0x00, 0x00,
};

bool get_descriptor(uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
	switch(wValue) {
		case 0x100:
			usb_write(0, (uint32_t*)dev_desc, wLength > sizeof(dev_desc) ? sizeof(dev_desc) : wLength);
			return true;
		case 0x200:
			usb_write(0, (uint32_t*)conf_desc, wLength > sizeof(conf_desc) ? sizeof(conf_desc) : wLength);
			return true;
		default:
			return false;
	}
}

bool set_address(uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
	OTG_HS.dev_reg.DCFG |= wValue << 4;
	usb_write(0, 0, 0);
	return true;
}

bool set_configuration(uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
	switch(wValue) {
		case 0:
			OTG_HS.dev_iep_reg[1].DIEPCTL = 0;
			break;
		
		case 1:
			OTG_HS.dev_iep_reg[1].DIEPCTL = (2 << 18) | (1 << 15) | 64;
			break;
		
		default:
			return false;
	}
	
	usb_write(0, 0, 0);
	return true;
}

bool jtag_shift(uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
	if(wLength > 16) {
		return false;
	}
	
	uint32_t tdo = 0;
	
	for(int16_t i = 0; i < wLength; i++) {
		tdo |= jtag_tick(wValue & 1, wIndex & 1) ? 1 << i : 0;
		wValue >>= 1;
		wIndex >>= 1;
	}
	
	usb_write(0, &tdo, (wLength + 7) >> 3);
	return true;
}

bool i2c_read(uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
	uint8_t buf[wLength];
	I2C2.read_reg(wValue, wIndex, wLength, buf);
	
	usb_write(0, (uint32_t*)buf, wLength);
	return true;
}

void handle_setup(const uint32_t* bufp) {
	led_error.toggle();
	
	uint8_t bmRequestType = bufp[0] & 0xff;
	uint8_t bRequest = (bufp[0] >> 8) & 0xff;
	uint16_t wValue = (bufp[0] >> 16) & 0xffff;
	uint16_t wIndex = bufp[1] & 0xffff;
	uint16_t wLength = (bufp[1] >> 16) & 0xffff;
	
	// GET_DESCRIPTOR
	if(bmRequestType == 0x80 && bRequest == 0x06) {
		if(get_descriptor(wValue, wIndex, wLength)) {
			return;
		}
	}
	
	// SET_ADDRESS
	if(bmRequestType == 0x00 && bRequest == 0x05) {
		if(set_address(wValue, wIndex, wLength)) {
			return;
		}
	}
	
	// SET_CONFIGURATION
	if(bmRequestType == 0x00 && bRequest == 0x09) {
		if(set_configuration(wValue, wIndex, wLength)) {
			return;
		}
	}
	
	// I2C_READ
	if(bmRequestType == 0xc0 && bRequest == 0xf0) {
		if(i2c_read(wValue, wIndex, wLength)) {
			return;
		}
	}
	
	// JTAG_SHIFT
	if(bmRequestType == 0xc0 && bRequest == 0xff) {
		if(jtag_shift(wValue, wIndex, wLength)) {
			return;
		}
	}
	
	OTG_HS.dev_iep_reg[0].DIEPCTL |= (1 << 21);
}

static uint32_t buf[16];

void handle_rxfifo() {
	led_status.toggle();
	
	uint32_t status = OTG_HS.reg.GRXSTSP;
	
	uint32_t len = (status & 0x7ff0) >> 6;
	
	for(uint32_t i = 0; i < len; i++) {
		buf[i] = OTG_HS.fifo[0].reg;
	}
	
	//if(status == 0x000c0080) {
	if((status & (0xf << 17)) == (0x4 << 17)) {
		handle_setup(buf);
		OTG_HS.dev_oep_reg[0].DOEPCTL |= (1 << 26);
	}
}

void usb_process() {
	// USB reset.
	if(OTG_HS.reg.GINTSTS & (1 << 12)) {
		//led_yellow.toggle();
		OTG_HS.dev_oep_reg[0].DOEPCTL = (1 << 27);
		OTG_HS.dev_reg.DAINTMSK = (1 << 16) | 1;
		OTG_HS.dev_reg.DOEPMSK = (1 << 3) | 1;
		OTG_HS.dev_reg.DIEPEMPMSK = (1 << 3) | 1;
		OTG_HS.reg.GRXFSIZ = 256;
		OTG_HS.reg.DIEPTXF0 = (64 << 16) | 256;
		OTG_HS.reg.DIEPTXF1 = (64 << 16) | 320;
		OTG_HS.dev_oep_reg[0].DOEPTSIZ = (3 << 29);
	}
	
	// OTG interrupt.
	if(OTG_HS.reg.GINTSTS & (1 << 2)) {
		led_status.toggle();
		
		OTG_HS.reg.GOTGINT = (1 << 2); // SEDET
	}
	
	// RxFIFO non-empty.
	if(OTG_HS.reg.GINTSTS & (1 << 4)) {
		handle_rxfifo();
	}
	
	OTG_HS.reg.GINTSTS = 0xffffffff;
}

template<>
void interrupt<(Interrupt::IRQ)77>() {
	usb_process();
}

void usb_main() {
	usb_vbus.set_mode(Pin::Input);
	usb_dm.set_mode(Pin::AF);
	usb_dm.set_pull(Pin::PullNone);
	usb_dm.set_af(12);
	usb_dp.set_mode(Pin::AF);
	usb_dp.set_pull(Pin::PullNone);
	usb_dp.set_af(12);
	
	RCC.enable(RCC.OTGHS);
	//Interrupt::enable((Interrupt::IRQ)77);
	
	// Set PHYSEL.
	OTG_HS.reg.GUSBCFG |= (1 << 6);
	
	Time::sleep(10);
	
	while(!(OTG_HS.reg.GRSTCTL & (1 << 31)));
	OTG_HS.reg.GRSTCTL |= 1;
	while(OTG_HS.reg.GRSTCTL & 1);
	
	OTG_HS.reg.GAHBCFG = 0;
	
	// USB configuration
	OTG_HS.reg.GUSBCFG = (1 << 30) | (0xf << 10) | (0 << 9) | (0 << 8) | (1 << 6);
	//                   FDMOD       TRDT          HNPCAP     SRPCAP     PHYSEL
	
	// interrupt mask
	OTG_HS.reg.GINTMSK = (1 << 13) | (1 << 12) | (1 << 11) | (1 << 10) | (1 << 3) | (1 << 2) | (1 << 1) | (1 << 4);
	//                   ENUMDNEM    USBRST      USBSUSPM    ESUSPM      SOFM       OTGINT     MMISM
	
	// device configuration
	OTG_HS.dev_reg.DCFG = (1 << 2) | 3;
	//                    NZLSOHSK   DSPD
	
	// core configuration
	OTG_HS.reg.GCCFG = (1 << 19) | (1 << 16);
	//                 VBUSBSEN    PWRDWN
	
	
	while(1) {
		usb_process();
		Thread::yield();
	}
}

uint32_t usb_stack[1024];

Thread usb_thread(usb_stack, sizeof(usb_stack), usb_main);

int main() {
	// Initialize system timer.
	STK.LOAD = 168000000 / 8 / 1000; // 1000 Hz.
	STK.CTRL = 0x03;
	
	RCC.enable(RCC.GPIOA);
	RCC.enable(RCC.GPIOB);
	RCC.enable(RCC.GPIOC);
	RCC.enable(RCC.GPIOD);
	
	led_status.set_mode(Pin::Output);
	led_status.off();
	led_error.set_mode(Pin::Output);
	led_error.off();
	
	jtag_tdi.set_mode(Pin::Output);
	jtag_tms.set_mode(Pin::Output);
	jtag_tck.set_mode(Pin::Output);
	jtag_tdo.set_mode(Pin::Input);
	
	RCC.enable(RCC.I2C2);
	I2C2.enable(i2c_scl, i2c_sda);
	
	usb_thread.start();
	
	while(1) {
		led_error.toggle();
		Time::sleep(100);
		
		int16_t buf[] = {
		};
		
		if(usb_ep_ready(1)) {
			usb_write(1, (uint32_t*)buf, sizeof(buf));
		}
	}
}