#include <rcc/rcc.h>
#include <gpio/pin.h>
#include <os/time.h>
#include <usb/usb.h>
#include <usb/descriptor.h>
#include <sdio/sdio.h>

static Pin sdio_pins[] = {
	// D0-3:
	PC8, PC9, PC10, PC11,
	// CK, CMD
	PC12, PD2
};

RBLog<64, 2> sd_rblog;

uint32_t sd_buf[512 / 4];

class SD_driver {
	private:
		SDIO_t& sdio;
		
		bool high_capacity;
		uint32_t rca;
	
	public:
		SD_driver(SDIO_t& s) : sdio(s) {}
		
		enum RespType {None, Short, Long = 3};
		
		void send_cmd(uint8_t cmd, uint32_t arg, RespType rt) {
			sd_rblog.log("Send CMD%d, arg=%#x", cmd, arg);
			
			sdio.reg.ARG = arg;
			sdio.reg.CMD = (1 << 10) | (rt << 6) | cmd;
			
			while(sdio.reg.STA & (1 << 11));
			
			sd_rblog.log("Command sent. STA=%#x, RESP1=%#x", sdio.reg.STA, sdio.reg.RESP1);
			
			sdio.reg.ICR = (1 << 7) | (1 << 6) | (1 << 2) | (1 << 0);
		}
		
		void read_data(uint32_t* buf, uint32_t len) {
			sd_rblog.log("Starting read. len=%d", len);
			
			sdio.reg.DLEN = len;
			
			sdio.reg.DCTRL = (9 << 4) | (1 << 1) | (1 << 0);
			
			while(!(sdio.reg.STA & (1 << 13)));
			sd_rblog.log("Receive started. STA=%#x", sdio.reg.STA);
			
			while(sdio.reg.STA & ((1 << 21) | (1 << 13))) {
				if(sdio.reg.STA & (1 << 21)) {
					*buf++ = sdio.reg.FIFO;
				}
			}
			
			sd_rblog.log("Receive finished. STA=%#x", sdio.reg.STA);
		}
		
		void write_data(uint32_t* buf, uint32_t len) {
			sd_rblog.log("Starting write. len=%d", len);
			
			sdio.reg.DLEN = len;
			
			sdio.reg.DCTRL = (9 << 4) | (0 << 1) | (1 << 0);
			
			while(len) {
				while(sdio.reg.STA & (1 << 16));
				
				sdio.reg.FIFO = *buf++;
				len -= 4;
			}
			
			while(sdio.reg.STA & (1 << 12));
			
			sd_rblog.log("Write finished. STA=%#x", sdio.reg.STA);
		}
		
		void initialize_card() {
			send_cmd(0, 0, None);
			send_cmd(8, 0x1aa, Short);
			
			send_cmd(55, 0, Short);
			send_cmd(41, 0, Short);
			
			do {
				send_cmd(55, 0, Short);
				send_cmd(41, (1 << 30) | (1 << 20), Short);
				
				Time::sleep(200);
			} while(!(sdio.reg.RESP1 & (1 << 31)));
			
			high_capacity = sdio.reg.RESP1 & (1 << 30);
			
			send_cmd(2, 0, Long);
			
			send_cmd(3, 0, Short);
			rca = sdio.reg.RESP1;
			
			send_cmd(7, rca, Short);
			
			send_cmd(55, rca, Short);
			send_cmd(6, 2, Short);
			
			sdio.reg.CLKCR = (1 << 11) | (1 << 9) | (1 << 8) | (48000000 / 12000000 - 1); // 4-bit, PWRSAV, CLKEN, 12 MHz
			
			read_block(0);
		}
		
		void enable() {
			sdio.reg.POWER = 3;
			sdio.reg.CLKCR = (1 << 11) | (1 << 9) | (1 << 8) | (48000000 / 400000 - 1); // 4-bit, PWRSAV, CLKEN, 400 kHz
			sdio.reg.DTIMER = 0xffffffff; // TODO: Find a sane value for this?
			
			initialize_card();
		}
		
		void read_block(uint32_t blocknum) {
			SDIO.reg.ICR = -1;
			
			if(!high_capacity) {
				blocknum *= 512;
			}
			
			send_cmd(17, blocknum, Short);
			
			send_cmd(13, rca, Short);
			
			read_data(sd_buf, 512);
			
			send_cmd(13, rca, Short);
		}
		
		void write_block(uint32_t blocknum) {
			SDIO.reg.ICR = -1;
			
			if(!high_capacity) {
				blocknum *= 512;
			}
			
			send_cmd(24, blocknum, Short);
			
			send_cmd(13, rca, Short);
			
			write_data(sd_buf, 512);
			
			send_cmd(13, rca, Short);
		}
};

SD_driver sd_driver(SDIO);

auto dev_desc = device_desc(0x200, 0, 0, 0, 64, 0x1234, 0x5678, 0, 0, 0, 0, 1);
auto conf_desc = configuration_desc(1, 1, 0, 0xc0, 0,
	// MSC BBB
	interface_desc(0, 0, 2, 0x08, 0x06, 0x50, 0,
		endpoint_desc(0x81, 0x02, 64, 0), // IN
		endpoint_desc(0x01, 0x02, 64, 0)  // OUT
	)
);

desc_t dev_desc_p = {sizeof(dev_desc), (void*)&dev_desc};
desc_t conf_desc_p = {sizeof(conf_desc), (void*)&conf_desc};

Pin& usb_vbus = PA9;
Pin& usb_dm   = PA11;
Pin& usb_dp   = PA12;

USB_otg usb(OTG_FS, dev_desc_p, conf_desc_p);

struct CBW {
	uint32_t dCBWSignature;
	uint32_t dCBWTag;
	uint32_t dCBWDataTransferLength;
	uint8_t bmCBWFlags;
	uint8_t bCBWLUN;
	uint8_t bCBWCBLength;
	uint8_t CBWCB[16];
} __attribute__((packed));

struct CSW {
	uint32_t dCSWSignature;
	uint32_t dCSWTag;
	uint32_t dCSWDataResidue;
	uint8_t bCSWStatus;
} __attribute__((packed));

uint32_t nullbuf[16];
uint32_t capbuf[] = {
	0x00002000,
	0x00020000
};
uint32_t inquirybuf[] = {
	0x02068020,
	32,
	
	0x20202020, 0x20202020,
	
	0x20202020, 0x20202020, 0x20202020, 0x20202020,
	
	0x34333231,
};

uint32_t sensebuf[] = {
	0x00050070,
	0x0a000000,
	
	0,
	0x00000020,
	0
};

class USB_MSC_BBB : public USB_class_driver {
	private:
		USB_generic& usb;
		
		uint32_t buf[16];
		
		uint32_t pending_data_in;
		
		bool pending_write;
		uint32_t pending_write_addr;
		uint32_t pending_write_num;
		uint32_t pending_write_recv;
		
		CBW cbw;
		CSW csw;
	
	public:
		USB_MSC_BBB(USB_generic& usbd) : usb(usbd) {
			usb.register_driver(this);
		}
	
	protected:
		virtual SetupStatus handle_setup(uint8_t bmRequestType, uint8_t bRequest, uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
			// Get max LUN
			if(bmRequestType == 0xa1 && bRequest == 0xfe) {
				uint32_t max_lun = 0;
				usb.write(0, &max_lun, 1);
				return SetupStatus::Ok;
			}
			
			// Mass Storage Reset
			if(bmRequestType == 0x21 && bRequest == 0xff) {
				usb.write(0, nullptr, 0);
				return SetupStatus::Ok;
			}
			
			return SetupStatus::Unhandled;
		}
		
		virtual void handle_set_configuration(uint8_t configuration) {
			if(configuration) {
				usb.register_out_handler(this, 1);
				usb.hw_conf_ep(0x01, EPType::Bulk, 64);
				usb.hw_conf_ep(0x81, EPType::Bulk, 64);
				
				pending_data_in = 0;
				pending_write = false;
			}
		}
		
		virtual void handle_out(uint8_t ep, uint32_t len) {
			if(ep == 0) {
				if(len) {
					usb.write(0, nullptr, 0);
				}
				
			} else if(ep == 1 && pending_data_in != 0) {
				if(pending_write) {
					usb_rblog.log("Handling write packet. (%d bytes)", len);
					
					handle_write_block_packet();
					
				} else {
					usb_rblog.log("Ignoring %d bytes of received data.", len);
				}
				
				if(len > pending_data_in) {
					pending_data_in = 0;
				} else {
					pending_data_in -= len;
				}
				
				if(pending_data_in == 0) {
					usb.write(1, (uint32_t*)&csw, sizeof(CSW));
				}
				
			} else if(ep == 1) {
				
				// Read CBW.
				if(len != sizeof(CBW) || usb.read(ep, (uint32_t*)&cbw, len) != sizeof(CBW)) {
					usb_rblog.log("Invalid CBW size.");
					return; // FIXME: Indicate error.
				}
				
				if(cbw.dCBWSignature != 0x43425355) {
					usb_rblog.log("Invalid CBW signature.");
					return; // FIXME: Indicate error.
				}
				
				usb_rblog.log("Received CBW, tl=%d, flags=%#02x", cbw.dCBWDataTransferLength, cbw.bmCBWFlags);
				
				bool cmd_ok = handle_scsi_cmd(cbw.CBWCB, cbw.bCBWCBLength);
				
				if(!cmd_ok && cbw.dCBWDataTransferLength != 0 && cbw.bmCBWFlags & 0x80) {
					write_zero();
				}
				
				if(cmd_ok) {
					csw = {0x53425355, cbw.dCBWTag, 0, 0x00};
				} else {
					csw = {0x53425355, cbw.dCBWTag, cbw.dCBWDataTransferLength, 0x01};
				}
				
				if(cbw.dCBWDataTransferLength != 0 && !(cbw.bmCBWFlags & 0x80)) {
					pending_data_in = cbw.dCBWDataTransferLength;
					return;
				}
				
				usb.write(1, (uint32_t*)&csw, sizeof(CSW));
			}
		}
		
		void write_zero() {
			uint32_t len = cbw.dCBWDataTransferLength;
			while(len > 64) {
				usb.write(1, nullbuf, 64);
				while(!usb.ep_ready(1));
				len -= 64;
			}
			
			usb.write(1, nullbuf, len);
			while(!usb.ep_ready(1));
		}
		
		void handle_read_block() {
			uint32_t block = (cbw.CBWCB[2] << 24) | (cbw.CBWCB[3] << 16) | (cbw.CBWCB[4] << 8) | cbw.CBWCB[5];
			uint16_t num_blocks = (cbw.CBWCB[7] << 8) | cbw.CBWCB[8];
			
			while(num_blocks--) {
				sd_driver.read_block(block++);
				
				for(uint32_t i = 0; i < 512; i += 64) {
					usb.write(1, sd_buf + i / 4, 64);
					while(!usb.ep_ready(1));
				}
			}
		}
		
		void handle_write_block() {
			pending_write_addr = (cbw.CBWCB[2] << 24) | (cbw.CBWCB[3] << 16) | (cbw.CBWCB[4] << 8) | cbw.CBWCB[5];
			pending_write_num = (cbw.CBWCB[7] << 8) | cbw.CBWCB[8];
			pending_write_recv = 0;
			
			pending_write = true;
		}
		
		void handle_write_block_packet() {
			usb_rblog.log("Received write packet. offset=%d, block=%d", pending_write_recv, pending_write_addr);
			
			usb.read(1, sd_buf + pending_write_recv / 4, 64);
			pending_write_recv += 64;
			
			if(pending_write_recv < 512) {
				return;
			}
			
			usb_rblog.log("Got full block, starting write.");
			
			sd_driver.write_block(pending_write_addr++);
			pending_write_num--;
			pending_write_recv = 0;
			
			pending_write = pending_write_num > 0;
		}
		
		bool handle_scsi_cmd(uint8_t* cmd, uint32_t len) {
			pending_write = false;
			
			if(!len) {
				return false;
			}
			
			switch(cmd[0]) {
				// TEST UNIT READY
				case 0x00:
					return true;
				
				// REQUEST SENSE
				case 0x03:
					usb.write(1, sensebuf, cbw.dCBWDataTransferLength);
					while(!usb.ep_ready(1));
					return true;
				
				// FORMAT UNIT
				case 0x04:
					return false; // TODO
				
				// INQUIRY
				case 0x12:
					usb.write(1, inquirybuf, cbw.dCBWDataTransferLength);
					while(!usb.ep_ready(1));
					return true;
				
				// MODE SENSE
				case 0x1a:
					usb.write(1, nullbuf, cbw.dCBWDataTransferLength);
					while(!usb.ep_ready(1));
					return true;
				
				// SEND DIAGNOSTIC
				case 0x1d:
					return false; // TODO
				
				// READ CAPACITY
				case 0x25:
					usb.write(1, capbuf, cbw.dCBWDataTransferLength);
					while(!usb.ep_ready(1));
					return true;
				
				// READ (10)
				case 0x28:
					handle_read_block();
					return true;
				
				// WRITE (10)
				case 0x2a:
					handle_write_block();
					return true; // TODO
				
				// READ (16)
				case 0x88:
					return false; // TODO
				
				// REPORT LUNS
				case 0xa0:
					return false; // TODO
				
				default:
					return false;
			}
			
			return true;
		}
};

USB_MSC_BBB usb_msc_bbb(usb);

int main() {
	// Initialize system timer.
	STK.LOAD = 168000000 / 8 / 1000; // 1000 Hz.
	STK.CTRL = 0x03;
	
	RCC.enable(RCC.GPIOA);
	RCC.enable(RCC.GPIOC);
	RCC.enable(RCC.GPIOD);
	
	for(Pin& p : sdio_pins) {
		p.set_mode(Pin::AF);
		p.set_af(12);
		p.set_speed(Pin::High);
	}
	
	RCC.enable(RCC.SDIO);
	
	sd_driver.enable();
	
	usb_vbus.set_mode(Pin::Input);
	usb_dm.set_mode(Pin::AF);
	usb_dm.set_af(10);
	usb_dp.set_mode(Pin::AF);
	usb_dp.set_af(10);
	
	RCC.enable(RCC.OTGFS);
	
	usb.init();
	
	while(1) {
		usb.process();
	}
}