#include <rcc/rcc.h>
#include <interrupt/interrupt.h>
#include <os/thread.h>
#include <os/time.h>

#include <gpio/pin.h>

#include <usb/usb.h>
#include <usb/descriptor.h>

#include <i2c/i2c.h>

#include <timer/timer.h>

#include "lsm303dlm.h"
#include "l3gd20.h"

#include "ppmsum.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;

static Pin& cs_pressure = PB4;
static Pin& cs_gyro = PB5;

static Pin& esc_power = PA1;
static Pin pwm_outputs[] = {PC6, PC7, PC8, PC9};

volatile float g_throttle = 0;

class PID {
		float i_acc;
		float last_error;
		
		const float dt = (1.0 / 400);
	
	public:
		float Kp;
		float Ki;
		float Kd;
		
		PID() : i_acc(0), last_error(0), Kp(0), Ki(0), Kd(0) {}
		
		float update(float error) {
			float corr = Kp * error + (i_acc += Ki * error * dt) + Kd * (error - last_error) / dt;
			
			last_error = error;
			
			return corr;
		}
		
		void reset() {
			i_acc = 0;
			last_error = 0;
		}
};

PID pid_x;
PID pid_y;
PID pid_z;

volatile uint32_t gyro_calib = 0;
float gyro_calib_x = 0;
float gyro_calib_y = 0;
float gyro_calib_z = 0;

auto dev_desc = device_desc(0x200, 0, 0, 0, 64, 0x1234, 0x5678, 0, 0, 0, 0, 1);
auto conf_desc = configuration_desc(3, 1, 0, 0xc0, 0,
	interface_desc(0, 0, 3, 0xff, 0x00, 0x00, 0,
		endpoint_desc(0x01, 0x02, 64, 0), // OUT
		endpoint_desc(0x81, 0x02, 64, 0), // IN
		endpoint_desc(0x82, 0x01, 64, 1)  // IN ISO
	)
);

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_HS, dev_desc_p, conf_desc_p);

class USB_I2C : public USB_class_driver {
	private:
		USB_generic& usb;
		I2C_t& i2c;
		
		bool i2c_read(uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
			uint8_t buf[wLength];
			i2c.read_reg(wValue, wIndex, wLength, buf);
			
			usb.write(0, (uint32_t*)buf, wLength);
			return true;
		}
	
	public:
		USB_I2C(USB_generic& usbd, I2C_t& i2c_bus) : usb(usbd), i2c(i2c_bus) {
			usb.register_driver(this);
		}
	
	protected:
		virtual SetupStatus handle_setup(uint8_t bmRequestType, uint8_t bRequest, uint16_t wValue, uint16_t wIndex, uint16_t wLength) {
			if(bmRequestType == 0xc0 && bRequest == 0xf0) {
				return i2c_read(wValue, wIndex, wLength) ? SetupStatus::Ok : SetupStatus::Stall;
			}
			
			return SetupStatus::Unhandled;
		}
};

USB_I2C usb_i2c(usb, I2C2);

class USB_JTAG : public USB_class_driver {
	private:
		USB_generic& usb;
		
		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;
		}
		
		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;
		}
	
	public:
		USB_JTAG(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) {
			if(bmRequestType == 0xc0 && bRequest == 0xff) {
				return jtag_shift(wValue, wIndex, wLength) ? SetupStatus::Ok : SetupStatus::Stall;
			}
			
			return SetupStatus::Unhandled;
		}
};

USB_JTAG usb_jtag(usb);

class USB_TM : public USB_class_driver {
	private:
		USB_generic& usb;
		
		uint32_t buf[16];
		
	public:
		USB_TM(USB_generic& usbd) : usb(usbd) {
			usb.register_driver(this);
		}
	
	protected:
		virtual void handle_set_configuration(uint8_t configuration) {
			if(configuration) {
				usb.hw_conf_ep(0x01, EPType::Bulk, 64);
				usb.hw_conf_ep(0x81, EPType::Bulk, 64);
				usb.hw_conf_ep(0x82, EPType::Isochronous, 64);
				usb.register_out_handler(this, 1);
			}
		}
		
		virtual void handle_out(uint8_t ep, uint32_t len) {
			if(ep == 1 && len) {
				usb.read(ep, buf, len);
				
				if(len == 2) {
					uint8_t* bufp = (uint8_t*)buf;
					
					switch(bufp[0]) {
						case 2:
							g_throttle = float(bufp[1]) / 128.0f * 0.7;
							break;
						
						case 3:
							pid_x.Kp = float(bufp[1]) / 128.0f * 0.4;
							pid_y.Kp = float(bufp[1]) / 128.0f * 0.4;
							break;
						
						case 4:
							pid_x.Ki = float(bufp[1]) / 128.0f * 0.4;
							pid_y.Ki = float(bufp[1]) / 128.0f * 0.4;
							break;
						
						case 5:
							pid_x.Kd = float(bufp[1]) / 128.0f * 0.4;
							pid_y.Kd = float(bufp[1]) / 128.0f * 0.4;
							break;
						
						case 6:
							pid_z.Kp = float(bufp[1]) / 128.0f * 0.4;
							break;
						
						case 7:
							pid_z.Kp = float(bufp[1]) / 128.0f * 0.4;
							break;
						
						case 8:
							pid_z.Kp = float(bufp[1]) / 128.0f * 0.4;
							break;
						
						case 23:
							if(bufp[1]) {
								gyro_calib = 256;
								gyro_calib_x = 0;
								gyro_calib_y = 0;
								gyro_calib_z = 0;
							}
							break;
					}
				}
			}
		}
};

USB_TM usb_tm(usb);

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);
	
	usb.init();
	
	while(1) {
		usb.process();
		Thread::yield();
	}
}

uint32_t usb_stack[1024];

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

L3GD20 gyro(cs_gyro, SPI1);
LSM303DLM_A accel(I2C2);
LSM303DLM_M magn(I2C2);

PPMSum ppmsum;

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();
	
	RCC.enable(RCC.SPI1);
	
	PA5.set_mode(Pin::AF);
	PA5.set_af(5);
	PA6.set_mode(Pin::AF);
	PA6.set_af(5);
	PA7.set_mode(Pin::AF);
	PA7.set_af(5);
	
	cs_gyro.on();
	cs_gyro.set_mode(Pin::Output);
	cs_pressure.on();
	cs_pressure.set_mode(Pin::Output);
	
	SPI1.reg.CR1 = (1 << 9) | (1 << 8) | (1 << 6) | (3 << 3) | (1 << 2) | (1 << 1) | (1 << 0); // SSM, SSI, SPE, 84/16MHz, MSTR, CPOL, CPHA
	
	Time::sleep(1000);
	
	gyro.init();
	accel.init();
	magn.init();
	
	PA0.set_mode(Pin::AF);
	PA0.set_af(1);
	ppmsum.enable();
	
	for(Pin& p : pwm_outputs) {
		p.set_af(3);
		p.set_mode(Pin::AF);
	}
	
	RCC.enable(RCC.TIM8);
	TIM8.PSC = 168 - 1;
	TIM8.ARR = 2500;
	TIM8.CCER = 0x1111;
	TIM8.CCMR1 = 0x6868;
	TIM8.CCMR2 = 0x6868;
	TIM8.BDTR = (1 << 15);
	TIM8.CCR1 = 1000;
	TIM8.CCR2 = 1000;
	TIM8.CCR3 = 1000;
	TIM8.CCR4 = 1000;
	TIM8.CR1 = 0x05;
	
	float abst = 0;
	
	while(1) {
		led_error.toggle();
		
		// Wait for new period.
		while(!(TIM8.SR & 1)) {
			Thread::yield();
		}
		
		TIM8.SR = 0;
		
		float throttle = g_throttle;
		
		gyro.update();
		//accel.update();
		//magn.update();
		
		if(gyro_calib) {
			gyro_calib--;
			gyro_calib_x += gyro.x / 256;
			gyro_calib_y += gyro.y / 256;
			gyro_calib_z += gyro.z / 256;
		}
		
		gyro.x -= gyro_calib_x;
		gyro.y -= gyro_calib_y;
		gyro.z -= gyro_calib_z;
		
		float stab_x = pid_x.update(-gyro.x);
		float stab_y = pid_y.update(-gyro.y);
		float stab_z = pid_z.update(-gyro.z);
		
		if(throttle < 0.05) {
			stab_x = 0;
			stab_y = 0;
			stab_z = 0;
			pid_x.reset();
			pid_y.reset();
			pid_z.reset();
		}
		
		float motors[] = {
			throttle - stab_x - stab_y + stab_z,
			throttle - stab_x + stab_y - stab_z,
			throttle + stab_x + stab_y + stab_z,
			throttle + stab_x - stab_y - stab_z
		};
		
		TIM8.CCR1 = 1000 + motors[0] * 1000;
		TIM8.CCR2 = 1000 + motors[1] * 1000;
		TIM8.CCR3 = 1000 + motors[2] * 1000;
		TIM8.CCR4 = 1000 + motors[3] * 1000;
		
		float buf[] = {
			abst,
			gyro.x,
			gyro.y,
			gyro.z,
			stab_x,
			stab_y,
			stab_z,
			throttle,
			0,
		};
		
		if(usb.ep_ready(2)) {
			usb.write(2, (uint32_t*)buf, sizeof(buf));
		} else {
			usb_rblog.log("Busy.");
		}
		
		abst += 1.0 / 400.0;
	}
}