#![no_main]
#![no_std]
#![feature(type_alias_impl_trait)]
use defmt_brtt as _; // global logger

use panic_probe as _;
use stm32f4xx_hal as _;

mod bootloader;

// same panicking *behavior* as `panic-probe` but doesn't print a panic message
// this prevents the panic message being printed *twice* when `defmt::panic` is invoked
#[defmt::panic_handler]
fn panic() -> ! {
    cortex_m::asm::udf()
}

#[rtic::app(
    device = stm32f4xx_hal::pac,
    dispatchers = [SPI3]
)]
mod app {

    use as5048a::AS5048A;

    use core::fmt::Write;
    use heapless::{String, Vec};
    use num_traits::{Float, FloatConst};
    use postcard::{from_bytes_cobs, to_vec_cobs};
    use stm32f4xx_hal::{
        gpio::{gpioa, gpiob, gpioc, Output, PushPull},
        i2c,
        otg_fs::{UsbBus, UsbBusType, USB},
        pac::{I2C1, SPI1},
        prelude::*,
        signature, spi,
    };
    use usb_device::prelude::*;
    use usb_device::{class_prelude::UsbBusAllocator, device};
    use usbd_serial::SerialPort;

    use xca9548a::{SlaveAddr, Xca9548a};

    use qmc5883l::{self, QMC5883L};

    use radomctl_protocol::{HostMessage, *};

    use crate::bootloader;
    use rtic_monotonics::systick::prelude::*;

    systick_monotonic!(Mono, 4000);

    const USB_BUFFER_SIZE: usize = 64;

    // Shared resources go here
    #[shared]
    struct Shared {
        az_angle: i32,
        az_compass: i32,
        az_target: i32,

        el_angle: i32,
        el_target: i32,
    }

    // Local resources go here
    #[local]
    struct Local {
        i2cmux: Xca9548a<i2c::I2c<I2C1>>,
        board_led: gpioc::PC13<Output<PushPull>>,

        encoder_az: AS5048A<spi::Spi<SPI1>, gpiob::PB12<Output<PushPull>>>,
        encoder_el: AS5048A<spi::Spi<SPI1>, gpiob::PB13<Output<PushPull>>>,
        spi_cs2: gpiob::PB14<Output<PushPull>>,
        spi_cs3: gpiob::PB15<Output<PushPull>>,
        spi1: spi::Spi<SPI1>,

        az_enable: gpiob::PB8<Output<PushPull>>,
        az_dir: gpioa::PA15<Output<PushPull>>,
        az_step: gpiob::PB3<Output<PushPull>>,

        el_enable: gpiob::PB4<Output<PushPull>>,
        el_dir: gpioa::PA8<Output<PushPull>>,
        el_step: gpioa::PA9<Output<PushPull>>,

        usb_dev: UsbDevice<'static, UsbBusType>,
        usb_serial: SerialPort<'static, UsbBusType>,
        usb_buffer: Vec<u8, USB_BUFFER_SIZE>,
    }

    #[init]
    fn init(cx: init::Context) -> (Shared, Local) {
        bootloader::init();

        defmt::info!("init");

        let rcc = cx.device.RCC.constrain();

        // Freeze the configuration of all the clocks in the system and store the frozen frequencies in
        // `clocks`
        let clocks = rcc
            .cfgr
            .use_hse(25.MHz())
            .sysclk(84.MHz())
            .require_pll48clk()
            .freeze();

        Mono::start(cx.core.SYST, clocks.sysclk().to_Hz());

        defmt::info!("Clock Setup done");

        // Acquire the GPIO peripherials
        let gpioa = cx.device.GPIOA.split();
        let gpiob = cx.device.GPIOB.split();
        let gpioc = cx.device.GPIOC.split();

        let board_led = gpioc.pc13.into_push_pull_output();

        defmt::info!("Basic gpio setup done");

        static mut EP_MEMORY: [u32; 1024] = [0; 1024];
        static mut USB_BUS: Option<usb_device::bus::UsbBusAllocator<UsbBusType>> = None;

        let usb = USB::new(
            (
                cx.device.OTG_FS_GLOBAL,
                cx.device.OTG_FS_DEVICE,
                cx.device.OTG_FS_PWRCLK,
            ),
            (gpioa.pa11, gpioa.pa12),
            &clocks,
        );

        unsafe {
            USB_BUS.replace(UsbBus::new(usb, &mut EP_MEMORY));
        }

        let usb_serial = usbd_serial::SerialPort::new(unsafe { USB_BUS.as_ref().unwrap() });

        let serial = unsafe {
            let u_id0 = 0x1FFF_7A10 as *const u32;
            let u_id2 = 0x1FFF_7A18 as *const u32;

            defmt::debug!("UID0: {:x}", u_id0.read());
            defmt::debug!("UID2: {:x}", u_id2.read());

            // See https://community.st.com/t5/stm32-mcus-products/usb-bootloader-serial-number/td-p/432148
            (u_id0.read() as u64 + u_id2.read() as u64) << 16
                | (u_id2.read() as u64 & 0xFF00) >> 8
                | (u_id2.read() as u64 & 0x00FF) << 8
        };

        static mut SERIAL: String<16> = String::new();
        unsafe {
            write!(SERIAL, "{:X}", serial).unwrap();
        }

        let usb_dev = unsafe {
            UsbDeviceBuilder::new(USB_BUS.as_ref().unwrap(), UsbVidPid(0x16c0, 0x27dd))
                .device_class(usbd_serial::USB_CLASS_CDC)
                .strings(&[StringDescriptors::default()
                    .manufacturer("Amteurfunk Forschungs Gruppe")
                    .product("Radom Controler")
                    .serial_number(SERIAL.as_ref())])
                .unwrap()
                .build()
        };

        defmt::info!("USB Setup done");

        // Todo: Check if internal pullups work here
        let scl = gpiob.pb6.into_alternate_open_drain();
        let sda = gpiob.pb7.into_alternate_open_drain();
        let i2c = i2c::I2c::new(
            cx.device.I2C1,
            (scl, sda),
            i2c::Mode::Standard {
                frequency: 400.kHz(),
            },
            &clocks,
        );

        defmt::info!("I2C Setup done");

        let mut i2cmux = Xca9548a::new(i2c, SlaveAddr::default());
        i2cmux.select_channels(0b0000_0001).unwrap();

        defmt::info!("I2C MUX Setup done");

        let spi_cs0 = gpiob.pb12.into_push_pull_output();
        let encoder_az = AS5048A::new(spi_cs0);

        let spi_cs1 = gpiob.pb13.into_push_pull_output();
        let encoder_el = AS5048A::new(spi_cs1);

        let spi_cs2 = gpiob.pb14.into_push_pull_output();
        let spi_cs3 = gpiob.pb15.into_push_pull_output();

        let sck = gpioa.pa5.into_push_pull_output();
        let poci = gpioa.pa6;
        let pico = gpioa.pa7.into_push_pull_output();
        let spi1 = spi::Spi::new(
            cx.device.SPI1,
            (sck, poci, pico),
            spi::Mode {
                polarity: spi::Polarity::IdleLow,
                phase: spi::Phase::CaptureOnSecondTransition,
            },
            8.MHz(),
            &clocks,
        );

        defmt::info!("SPI Setup done");

        let mut az_enable = gpiob.pb8.into_push_pull_output();
        az_enable.set_high();
        let az_dir = gpioa.pa15.into_push_pull_output();
        let az_step = gpiob.pb3.into_push_pull_output();

        let mut el_enable = gpiob.pb4.into_push_pull_output();
        el_enable.set_high();
        let el_dir = gpioa.pa8.into_push_pull_output();
        let el_step = gpioa.pa9.into_push_pull_output();

        defmt::info!("Motor Setup done");

        poll_i2c::spawn().ok();
        poll_spi::spawn().ok();
        move_az::spawn().ok();
        move_el::spawn().ok();

        (
            Shared {
                az_angle: 0,
                az_target: 0,
                el_angle: 0,
                el_target: 0,
                az_compass: 0,
            },
            Local {
                i2cmux,
                board_led,
                encoder_az,
                encoder_el,
                spi_cs2,
                spi_cs3,
                spi1,

                az_enable,
                az_dir,
                az_step,

                el_enable,
                el_dir,
                el_step,

                usb_dev,
                usb_serial,
                usb_buffer: Vec::new(),
            },
        )
    }

    #[task(local = [i2cmux, board_led], shared = [az_compass])]
    async fn poll_i2c(mut cx: poll_i2c::Context) {
        let i2cmux = cx.local.i2cmux;
        let board_led = cx.local.board_led;

        let parts = i2cmux.split();

        let mut compass1 = QMC5883L::new(parts.i2c0).unwrap();
        compass1.reset().unwrap();
        compass1.continuous().unwrap();

        let mut compass2 = QMC5883L::new(parts.i2c1).unwrap();
        compass2.reset().unwrap();
        compass2.continuous().unwrap();

        let declination_rads: f32 = 65.0 / 180.0 * f32::PI();

        loop {
            board_led.toggle();

            loop {
                defmt::info!("Compass 1");
                match compass1.mag() {
                    Ok((x, y, z)) => {
                        defmt::info!("x1: {} y1: {} z1: {}", x, y, z);

                        let mut heading = (y as f32).atan2(x as f32); //+ declination_rads;
                        if heading < 0.0 {
                            heading += 2.0 * f32::PI();
                        } else if heading > 2.0 * f32::PI() {
                            heading -= 2.0 * f32::PI();
                        }
                        let heading_degrees = heading * 180.0 / f32::PI();

                        cx.shared.az_compass.lock(|az_compass| {
                            *az_compass = heading_degrees as i32 * 10;
                        });

                        defmt::info!("Heading1 {}", heading_degrees);
                        break;
                    }
                    Err(qmc5883l::Error::NotReady) => {
                        Mono::delay(1000.millis()).await;
                    }
                    e => {
                        let _ = e.unwrap();
                    }
                }
            }

            loop {
                defmt::info!("Compass 2");
                match compass2.mag() {
                    Ok((x, y, z)) => {
                        defmt::info!("x2: {} y2: {} z2: {}", x, y, z);

                        let mut heading = (y as f32).atan2(x as f32); //+ declination_rads;
                        if heading < 0.0 {
                            heading += 2.0 * f32::PI();
                        } else if heading > 2.0 * f32::PI() {
                            heading -= 2.0 * f32::PI();
                        }
                        let heading_degrees = heading * 180.0 / f32::PI();
                        defmt::info!("Heading2 {}", heading_degrees);
                        break;
                    }
                    Err(qmc5883l::Error::NotReady) => {
                        Mono::delay(1000.millis()).await;
                    }
                    e => {
                        let _ = e.unwrap();
                    }
                }
            }

            Mono::delay(100.millis()).await;
        }
    }

    #[task(local = [spi1, encoder_az, encoder_el, spi_cs2, spi_cs3], shared = [az_angle, el_angle])]
    async fn poll_spi(mut cx: poll_spi::Context) {
        let spi1 = cx.local.spi1;
        let encoder_az = cx.local.encoder_az;
        let encoder_el = cx.local.encoder_el;

        loop {
            /*
            let (diag, gain) = encoder_az.diag_gain(spi1).unwrap();
            defmt::info!("diag: {:08b} gain: {}", diag, gain);
            defmt::info!("magnitude: {:?}", encoder_az.magnitude(spi1).unwrap());
            */
            let raw_angle = encoder_az.angle(spi1).unwrap();
            let angle_deg = raw_angle as i32 * 3600 / 16384;

            cx.shared.az_angle.lock(|az_angle| {
                *az_angle = angle_deg;
            });

            defmt::info!("az angle: {:?}", angle_deg);

            let raw_angle = encoder_el.angle(spi1).unwrap();
            let angle_deg = raw_angle as i32 * 3600 / 16384;

            cx.shared.el_angle.lock(|el_angle| {
                *el_angle = angle_deg;
            });

            defmt::info!("el angle: {:?}", angle_deg);

            Mono::delay(1.millis()).await;
        }
    }

    #[task(local = [az_enable, az_dir, az_step], shared = [az_angle, az_target])]
    async fn move_az(mut cx: move_az::Context) {
        let az_enable = cx.local.az_enable;
        let az_dir = cx.local.az_dir;
        let az_step = cx.local.az_step;

        loop {
            let az_target = cx.shared.az_target.lock(|az_target| *az_target);
            let az_angle = cx.shared.az_angle.lock(|az_angle| *az_angle);
            let diff = az_angle - az_target;

            defmt::info!(
                "angle diff/target/actual: {:?}/{:?}/{:?}",
                diff,
                az_target,
                az_angle
            );

            let delay = if diff.abs() < 10 {
                10.millis()
            } else if diff < 100 {
                5.millis()
            } else {
                1.millis()
            };

            if diff.abs() > 50 {
                az_enable.set_low();
                if diff < 0 {
                    az_dir.set_high();
                } else {
                    az_dir.set_low();
                }

                az_step.set_low();
                Mono::delay(delay / 2).await;
                az_step.set_high();
                Mono::delay(delay / 2).await;
            } else {
                az_enable.set_high();
                Mono::delay(delay).await;
            }
        }
    }

    #[task(local = [el_enable, el_dir, el_step], shared = [el_angle, el_target])]
    async fn move_el(mut cx: move_el::Context) {
        let el_enable = cx.local.el_enable;
        let el_dir = cx.local.el_dir;
        let el_step = cx.local.el_step;

        loop {
            let el_target = cx.shared.el_target.lock(|el_target| *el_target);
            let el_angle = cx.shared.el_angle.lock(|el_angle| *el_angle);
            let diff = el_angle - el_target;

            defmt::info!(
                "angle diff/target/actual: {:?}/{:?}/{:?}",
                diff,
                el_target,
                el_angle
            );

            let delay = if diff.abs() < 10 {
                10.millis()
            } else if diff < 100 {
                5.millis()
            } else {
                1.millis()
            };

            if diff.abs() > 50 {
                el_enable.set_low();
                if diff < 0 {
                    el_dir.set_high();
                } else {
                    el_dir.set_low();
                }

                el_step.set_low();
                Mono::delay(delay / 2).await;
                el_step.set_high();
                Mono::delay(delay / 2).await;
            } else {
                el_enable.set_high();
                Mono::delay(delay).await;
            }
        }
    }

    #[task(binds=OTG_FS, local=[usb_dev, usb_serial, usb_buffer], shared=[az_target, el_target, az_angle, el_angle])]
    fn usb_fs(mut cx: usb_fs::Context) {
        let usb_dev = cx.local.usb_dev;
        let serial = cx.local.usb_serial;
        let buffer = cx.local.usb_buffer;

        if !usb_dev.poll(&mut [serial]) {
            return;
        }

        let mut tmp = [0u8; 16];
        match serial.read(&mut tmp) {
            Ok(count) if count > 0 => {
                if buffer.extend_from_slice(&tmp[0..count]).is_err() {
                    buffer.clear();
                    defmt::error!("Buffer overflow while waiting for the end of the packet");
                }
            }
            _ => {}
        }

        loop {
            if let Some(idx) = buffer.iter().position(|&x| x == 0) {
                let (msg, rest) = buffer.split_at(idx + 1);

                let mut message = [0u8; 128];
                message[0..msg.len()].clone_from_slice(msg);
                let host_msg = from_bytes_cobs::<HostMessage>(&mut message);

                match host_msg {
                    Ok(host_msg) => match host_msg {
                        HostMessage::RequestStatus => {
                            let status = StatusMessage {
                                position: Position {
                                    az: cx.shared.az_angle.lock(|az| (*az / 10) as f32),
                                    el: cx.shared.el_angle.lock(|el| (*el / 10) as f32),
                                    az_endcoder: 0.0,
                                    el_encoder: 0.0,
                                    az_magnetic: 0.0,
                                    el_magnetic: 0.0,
                                },
                                alarms: Vec::new(),
                            };
                            let device_msg = RadomMessage::Status(status);
                            let bytes =
                                to_vec_cobs::<RadomMessage, USB_BUFFER_SIZE>(&device_msg).unwrap();
                            serial.write(bytes.as_slice()).unwrap();
                        }
                        HostMessage::SetTarget(pos) => {
                            cx.shared.az_target.lock(|az| *az = (pos.az * 10.0) as i32);
                            cx.shared.el_target.lock(|el| *el = (pos.el * 10.0) as i32);
                        }
                        HostMessage::TriggerDFUBootloader => {
                            bootloader::reboot_to_bootloader();
                        }
                    },
                    Err(err) => defmt::error!("Unable to parse host message"),
                };

                *buffer = Vec::<u8, USB_BUFFER_SIZE>::from_slice(rest).unwrap();
            } else {
                break;
            }
        }
    }
}