#include <rcc/rcc.h>
#include <gpio/pin.h>
#include <os/time.h>
#include <usb/usb.h>
#include <usb/descriptor.h>
#include <spi/spi.h>
#include <dma/dma.h>
#include <i2c/i2c.h>

Pin& dac_nreset = PD4;

Pin& i2c_scl = PB6;
Pin& i2c_sda = PB9;

Pin& i2s_mck = PC7;
Pin& i2s_sck = PC10;
Pin& i2s_sd  = PC12;
Pin& i2s_ws  = PA4;

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

const uint32_t audio_buf_size = 1024;
uint16_t audio_buf[audio_buf_size];
uint32_t audio_buf_pos = 0;

struct AS_General_desc {
	uint8_t bLength;
	uint8_t bDescriptorType;
	uint8_t bDescriptorSubType;
	uint8_t bTerminalLink;
	uint8_t bDelay;
	uint16_t wFormatTag;
} __attribute__((packed));

constexpr AS_General_desc as_general_desc(
		uint8_t bTerminalLink,
		uint8_t bDelay,
		uint16_t wFormatTag
	) {
	
	return {
		sizeof(AS_General_desc),
		0x24,
		0x01,
		bTerminalLink,
		bDelay,
		wFormatTag
	};
}

template <uint32_t S>
struct foo_t {
	uint8_t a[S];
} __attribute__((packed));

const foo_t<43> acd = {{
	9, 36, 1, 0,1, 43,0, 1, 1,
	12, 36, 2, 1, 1,1, 0, 2, 3,0, 0, 0,
	13, 36, 6, 13, 1, 2, 1, 2, 2, 0, 2, 0, 0,
	9, 36, 3, 3, 1,3, 0, 13, 0
}};

const foo_t<11> asfd = {{11, 36, 2, 1, 2, 2, 16, 1, 48000 & 0xff, (48000 >> 8) & 0xff, 48000 >> 16}};

const foo_t<7> aced = {{7, 37, 1, 1, 0, 0, 0}};

auto dev_desc = device_desc(0x200, 0, 0, 0, 64, 0x1234, 0x5678, 0, 0, 0, 0, 1);
auto conf_desc = configuration_desc(2, 1, 0, 0xc0, 0,
	// AudioControl
	interface_desc(0, 0, 0, 0x01, 0x01, 0x00, 0,
		acd
	),
	// AudioStreaming
	interface_desc(1, 0, 0, 0x01, 0x02, 0x00, 0),
	interface_desc(1, 1, 1, 0x01, 0x02, 0x00, 0,
		as_general_desc(1, 1, 1),
		asfd,
		endpoint_desc(0x01, 0x09, 256, 1),
		aced
	)
);

class DAC {
	Pin& nreset;
	I2C& i2c;

	public:
		DAC(Pin& _nreset, I2C& _i2c) : nreset(_nreset), i2c(_i2c) {}
		
		void init() {
			// Configure CS43L22.
			nreset.set_mode(Pin::Output);
			nreset.on();
			
			i2c.write_reg(0x4a, 0x02, 0x9e);
			i2c.write_reg(0x4a, 0x06, (0 << 7) | (1 << 6) | (0 << 4) | (1 << 2) | (3 << 0));
		}
		uint8_t get_volume(uint8_t c) {
			uint8_t buf;
			i2c.read_reg(0x4a, 0x1a + c, 1, &buf);
			return ((buf & 0x7f) - 0x19) & 0x7f;
		}
		void set_volume(uint8_t c, uint8_t v) {
			i2c.write_reg(0x4a, 0x1a + c, (v + 0x19) & 0x7f);
		}
};

DAC dac(dac_nreset, I2C1);

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

USB_otg usb(OTG_FS, dev_desc_p, conf_desc_p);

class USB_Audio : public USB_class_driver {
	private:
		USB_generic& usb;
		
		enum ControlState {None, VolLeft, VolRight};
		ControlState control_state;
	
	public:
		USB_Audio(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) {
			control_state = ControlState::None;
			
			// GET_CUR
			if(bmRequestType == 0xa1 && bRequest == 0x81) {
				uint32_t res = 0;
				usb.write(0, &res, wLength);
				return SetupStatus::Ok;
			}
			
			// GET_MIN
			if(bmRequestType == 0xa1 && bRequest == 0x82) {
				uint32_t res = 0;
				usb.write(0, &res, wLength);
				return SetupStatus::Ok;
			}
			
			// GET_MAX
			if(bmRequestType == 0xa1 && bRequest == 0x83) {
				uint32_t res = 127;
				usb.write(0, &res, wLength);
				return SetupStatus::Ok;
			}
			
			// GET_RES
			if(bmRequestType == 0xa1 && bRequest == 0x84) {
				uint32_t res = 1;
				usb.write(0, &res, wLength);
				return SetupStatus::Ok;
			}
			
			// SET_CUR
			if(bmRequestType == 0x21 && bRequest == 0x01) {
				switch(wValue) {
					case 0x0201:
						control_state = ControlState::VolLeft;
						break;
					
					case 0x0202:
						control_state = ControlState::VolRight;
						break;
				}
				
				return SetupStatus::Ok;
			}
			
			// SET_CUR
			if(bmRequestType == 0x22 && bRequest == 0x01) {
				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::Isochronous, 256);
			}
		}
		
		virtual void handle_out(uint8_t ep, uint32_t len) {
			if(ep == 0 && len > 0) {
				uint32_t buf;
				usb.read(0, &buf, len > 4 ? 4 : len);
				
				switch(control_state) {
					case ControlState::VolLeft:
						dac.set_volume(0, buf & 0xff);
						break;
					
					case ControlState::VolRight:
						dac.set_volume(1, buf & 0xff);
						break;
					
					default:
						break;
				}
				
				usb.write(0, nullptr, 0);
				
			} else if(ep == 1) {
				usb_rblog.log("Received audio frame with %d bytes.", len);
				
				if(audio_buf_pos + (len >> 1) >= audio_buf_size) {
					usb.read(1, (uint32_t*)&audio_buf[audio_buf_pos], (audio_buf_size - audio_buf_pos) << 1);
					len -= (audio_buf_size - audio_buf_pos) << 1;
					audio_buf_pos = 0;
				}
				
				if(len) {
					usb.read(1, (uint32_t*)&audio_buf[audio_buf_pos], len);
					audio_buf_pos += len >> 1;
				}
			}
		}
};

USB_Audio usb_audio(usb);

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);
	
	// Initialize I2S.
	i2s_mck.set_mode(Pin::AF);
	i2s_mck.set_af(6);
	i2s_sck.set_mode(Pin::AF);
	i2s_sck.set_af(6);
	i2s_sd.set_mode(Pin::AF);
	i2s_sd.set_af(6);
	i2s_ws.set_mode(Pin::AF);
	i2s_ws.set_af(6);
	
	RCC.PLLI2SCFGR = (3 << 28) | (258 << 6);
	RCC.CR |= 1 << 26;
	while(!(RCC.CR & (1 << 27)));
	
	RCC.enable(RCC.SPI3);
	
	SPI3.reg.I2SPR = 0x303;
	SPI3.reg.I2SCFGR = 0xe00;
	SPI3.reg.CR2 = 2;
	
	// Initialize DMA.
	RCC.enable(RCC.DMA1);
	
	DMA1.reg.S[7].CR = (0 << 25) | (1 << 13) | (1 << 11) | (1 << 10) | (1 << 8) | (1 << 6);
	DMA1.reg.S[7].PAR = uint32_t(&SPI3.reg.DR);
	DMA1.reg.S[7].M0AR = uint32_t(&audio_buf);
	DMA1.reg.S[7].NDTR = audio_buf_size;
	DMA1.reg.S[7].CR |= 1;
	
	// Initialize I2C.
	I2C1.enable(i2c_scl, i2c_sda);
	
	dac.init();
	
	// Initialize USB.
	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.set_rxfifo_size(512);
	usb.init();
	
	while(1) {
		usb.process();
	}
}