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

use panic_probe as _;
use stm32f1xx_hal as _;

mod nvstate;
mod si5153;

// 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()
}

use rtic::app;

#[app(device = stm32f1xx_hal::pac, peripherals = true, dispatchers = [SPI3])]
mod app {

    use cortex_m::asm::delay;
    use rtic_monotonics::systick::*;

    use stm32f1xx_hal::{
        flash::{self, FlashWriter},
        gpio::{self, gpioa, gpioc, Alternate, Output, PushPull},
        i2c, pac,
        pac::{RCC, TIM2, TIM3, TIM4},
        prelude::*,
        rcc::Enable,
        rcc::Reset,
        timer::{self, Channel, PwmHz, Tim4NoRemap},
    };

    use stm32f1xx_hal::usb::{Peripheral, UsbBus, UsbBusType};
    use usb_device::prelude::*;

    use heapless::Vec;
    use postcard::{from_bytes_cobs, to_vec_cobs};

    use cheapsdo_protocol::{DeviceMessage, HostMessage, PLLSettings, StatusMessage};

    use crate::nvstate::{self, NVState};
    use crate::si5153;

    const USB_BUFFER_SIZE: usize = 64;

    #[local]
    struct Local {
        board_led: gpioc::PC13<Output<PushPull>>,
        tim2: TIM2,
        tim3: TIM3,
        pwm: PwmHz<TIM4, Tim4NoRemap, timer::Ch<0>, gpio::Pin<'B', 6, Alternate>>,
    }

    #[shared]
    struct Shared {
        usb_dev: UsbDevice<'static, UsbBusType>,
        serial: usbd_serial::SerialPort<'static, UsbBusType>,
        device_status: StatusMessage,
        buffer: Vec<u8, USB_BUFFER_SIZE>,
        nvstate: NVState,
        flash: flash::Parts,
    }

    const TARGET_FREQ: u64 = 10_000_000_000; // in millihertz

    #[init]
    fn init(cx: init::Context) -> (Shared, Local) {
        let rcc = cx.device.RCC.constrain();
        let mut flash = cx.device.FLASH.constrain();

        let clocks = rcc
            .cfgr
            .use_hse(12.MHz())
            .sysclk(48.MHz())
            .pclk1(24.MHz())
            .freeze(&mut flash.acr);

        assert!(clocks.usbclk_valid());
        defmt::info!("Clock Setup done");

        // Initialize the systick interrupt & obtain the token to prove that we did
        let systick_mono_token = rtic_monotonics::create_systick_token!();
        Systick::start(cx.core.SYST, clocks.sysclk().to_Hz(), systick_mono_token);

        let mut gpioc = cx.device.GPIOC.split();

        // Configure gpio C pin 13 as a push-pull output. The `crh` register is passed to the function
        // in order to configure the port. For pins 0-7, crl should be passed instead.
        let board_led = gpioc.pc13.into_push_pull_output(&mut gpioc.crh);

        let mut afio = cx.device.AFIO.constrain();
        let mut gpiob = cx.device.GPIOB.split();
        let pwm_pin = gpiob.pb6.into_alternate_push_pull(&mut gpiob.crl);
        let mut pwm =
            cx.device
                .TIM4
                .pwm_hz::<Tim4NoRemap, _, _>(pwm_pin, &mut afio.mapr, 16.kHz(), &clocks);
        pwm.enable(Channel::C1);

        defmt::info!("Max PWM is {}", pwm.get_max_duty());
        defmt::info!("PWM Setup done");

        let tim2 = cx.device.TIM2;
        unsafe {
            let rcc = &*RCC::ptr();
            TIM2::enable(rcc);
            TIM2::reset(rcc);
        }

        // Enable external clocking
        tim2.smcr.write(|w| {
            w.etf().no_filter(); // No filter for to 10Mhz clock
            w.etps().div1(); // No divider
            w.etp().not_inverted(); // on rising edege at ETR pin
            w.ece().enabled() // mode 2 (use ETR pin)
        });

        tim2.ccmr1_input().write(|w| {
            w.cc1s().ti2(); // Input capture using TI2 input
            w.ic1f().no_filter() // No filter on input capture input
                                 //w.ic1psc().bits(0) // Disable prescaler, not safely implement by HAL yet
        });

        tim2.ccer.write(|w| {
            w.cc1p().set_bit(); // Use rising edge on TI
            w.cc1e().set_bit() // Enable input capture
        });

        tim2.cr2.write(|w| {
            w.mms().update() // Trigger output on update/overflow
        });

        tim2.ccer.write(|w| {
            w.cc1p().set_bit(); // Use rising edge on TI
            w.cc1e().set_bit() // Enable input capture
        });

        tim2.cr2.write(|w| {
            w.mms().update() // Trigger output on update/overflow
        });

        // Counting up to 10^7 should need 24 bits
        // Clock tim2 by tim1s overflow to make a 32bit timer
        let tim3 = cx.device.TIM3;
        unsafe {
            let rcc = &*RCC::ptr();
            TIM3::enable(rcc);
            TIM3::reset(rcc);
        }

        tim3.smcr.write(|w| {
            w.ts().itr1(); // Trigger from internal trigger 1
            w.sms().ext_clock_mode() // Use trigger as clock
        });

        tim3.ccmr1_input().write(|w| {
            w.cc1s().ti1(); // Input capture using TI1 input
            w.ic1f().no_filter() // No filter on input capture input
                                 //w.ic1psc().bits(0) // Disable prescaler, not safely implement by HAL yet
        });

        tim3.ccer.write(|w| {
            w.cc1p().set_bit(); // Use rising edge on TI
            w.cc1e().set_bit() // Enable input capture
        });

        tim2.cr1.write(|w| w.cen().enabled());
        tim3.cr1.write(|w| w.cen().enabled());

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

        static mut USB_BUS: Option<usb_device::bus::UsbBusAllocator<UsbBusType>> = None;

        let mut gpioa = cx.device.GPIOA.split();

        let mut usb_dp = gpioa.pa12.into_push_pull_output(&mut gpioa.crh);
        usb_dp.set_low();
        delay(clocks.sysclk().raw() / 100);

        let usb_dm = gpioa.pa11;
        let usb_dp = usb_dp.into_floating_input(&mut gpioa.crh);

        let usb = Peripheral {
            usb: cx.device.USB,
            pin_dm: usb_dm,
            pin_dp: usb_dp,
        };

        unsafe {
            USB_BUS.replace(UsbBus::new(usb));
        }

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

        let usb_dev = UsbDeviceBuilder::new(
            unsafe { USB_BUS.as_ref().unwrap() },
            UsbVidPid(0x16c0, 0x27dd),
        )
        .manufacturer("Arbitrary Precision Instruments")
        .product("cheapsdo")
        .serial_number("1337")
        .device_class(usbd_serial::USB_CLASS_CDC)
        .build();

        let scl = gpiob.pb8.into_alternate_open_drain(&mut gpiob.crh);
        let sda = gpiob.pb9.into_alternate_open_drain(&mut gpiob.crh);
        let mut i2c1 = i2c::BlockingI2c::i2c1(
            cx.device.I2C1,
            (scl, sda),
            &mut afio.mapr,
            i2c::Mode::Fast {
                frequency: 400_000.Hz(),
                duty_cycle: i2c::DutyCycle::Ratio2to1,
            },
            clocks,
            1000,
            10,
            1000,
            1000,
        );

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

        let mut si_pll = si5153::Si5153::new(&i2c1);
        si_pll.init(&mut i2c1, 10_000_000, 800_000_000, 800_000_000);
        si_pll.set_ms_source(&mut i2c1, si5153::Multisynth::MS0, si5153::PLL::A);
        defmt::info!("si5153 Setup done");

        si_pll.set_ms_freq(&mut i2c1, si5153::Multisynth::MS0, 100_000_000);
        si_pll.enable_ms_output(&mut i2c1, si5153::Multisynth::MS0);

        let nvstate = nvstate::load(&mut flash);

        defmt::info!("read nvstate from flash");

        update_pwm::spawn().unwrap();

        (
            Shared {
                serial,
                usb_dev,
                device_status: StatusMessage::default(),
                buffer: Vec::new(),
                nvstate,
                flash,
            },
            Local {
                board_led,
                tim2,
                tim3,
                pwm,
            },
        )
    }

    const WINDOW_LEN: usize = 100;

    #[task(local=[tim2, tim3, pwm, board_led], shared=[device_status, nvstate, flash])]
    async fn update_pwm(mut cx: update_pwm::Context) {
        defmt::info!("Update Task started");

        let tim2 = cx.local.tim2;
        let tim3 = cx.local.tim3;
        let pwm = cx.local.pwm;
        let board_led = cx.local.board_led;

        let max_pwm = pwm.get_max_duty() as i32;

        let mut cur_pwm = 0i32;
        cx.shared.nvstate.lock(|nvstate| {
            cur_pwm = nvstate.pwm as i32;
        });

        cur_pwm = if cur_pwm < 0 { 0 } else { cur_pwm };
        cur_pwm = if cur_pwm > max_pwm { max_pwm } else { cur_pwm };

        pwm.set_duty(Channel::C1, cur_pwm as u16);
        defmt::info!("pwm:\t{}", cur_pwm);

        // Inialize last_ic
        while !tim2.sr.read().cc1if().bit_is_set() || !tim3.sr.read().cc1if().bit_is_set() {
            Systick::delay(10.millis()).await;
        }
        let ic1 = tim2.ccr1().read().bits();
        let ic2 = tim3.ccr1().read().bits();

        let mut last_ic = ic2 << 16 | ic1;

        let mut samples: [u64; WINDOW_LEN] = [10_000_000_000; WINDOW_LEN];
        let mut short_avg: u64 = 10_000_000_000;

        loop {
            let mut last_freq: u64 = 10_000_000_000;

            let mut count = 0;
            while count < WINDOW_LEN {
                while !tim3.sr.read().cc1if().bit_is_set() || !tim3.sr.read().cc1if().bit_is_set() {
                    Systick::delay(10.millis()).await;
                }
                let ic1 = tim2.ccr1().read().bits();
                let ic2 = tim3.ccr1().read().bits();

                let sum_ic = ic2 << 16 | ic1;

                let diff_ic = if sum_ic > last_ic {
                    sum_ic - last_ic
                } else {
                    u32::MAX - last_ic + sum_ic
                };

                defmt::info!("ic1:\t{}", ic1);
                defmt::info!("ic2:\t{}", ic2);
                defmt::info!("sum_ic:\t{}", sum_ic);
                defmt::info!("diff_ic:\t{}", diff_ic);

                last_ic = sum_ic;

                let freq = (diff_ic as u64) * 1000;
                defmt::info!("freq:\t{} mHz", freq as i64);
                defmt::info!("last_freq:\t{} mHz", last_freq as i64);

                let diff = freq as i64 - last_freq as i64;
                last_freq = freq;

                if diff.abs() > 50_000 {
                    defmt::info!("Out of range, dropping sample.");
                    continue;
                }

                samples[count] = freq;
                short_avg = 0;
                for i in 0..WINDOW_LEN {
                    short_avg += samples[i];
                }
                short_avg = short_avg / WINDOW_LEN as u64;
                defmt::info!("short_avg:\t{} mHz", short_avg);

                cx.shared.device_status.lock(|device_status| {
                    device_status.measured_frequency = freq;
                    device_status.average_frequency = short_avg;
                    device_status.pwm = cur_pwm as u16;
                });

                count += 1;
                board_led.toggle();
            }

            let diff = TARGET_FREQ as i64 - short_avg as i64;
            if diff.abs() > 100 {
                cur_pwm += (diff * 30 / 1000) as i32;
            } else if diff < -10 {
                cur_pwm -= 1;
            } else if diff > 10 {
                cur_pwm += 1;
            }

            cur_pwm = if cur_pwm < 0 { 0 } else { cur_pwm };
            cur_pwm = if cur_pwm > max_pwm { max_pwm } else { cur_pwm };

            pwm.set_duty(Channel::C1, cur_pwm as u16);
            defmt::info!("pwm:\t{}", cur_pwm);

            {
                (&mut cx.shared.nvstate, &mut cx.shared.flash).lock(|nvstate, flash| {
                    nvstate.pwm = cur_pwm as u16;
                    nvstate.save(flash);
                });
            }

            Systick::delay(500.millis()).await;
        }
    }

    #[task(binds = USB_HP_CAN_TX, shared = [usb_dev, serial, buffer, device_status, nvstate, flash])]
    fn usb_tx(cx: usb_tx::Context) {
        let mut usb_dev = cx.shared.usb_dev;
        let mut serial = cx.shared.serial;
        let mut buffer = cx.shared.buffer;
        let mut device_status = cx.shared.device_status;
        let mut nvstate = cx.shared.nvstate;
        let mut flash = cx.shared.flash;

        (
            &mut usb_dev,
            &mut serial,
            &mut buffer,
            &mut device_status,
            &mut nvstate,
            &mut flash,
        )
            .lock(|usb_dev, serial, buffer, device_status, nvstate, flash| {
                usb_poll(usb_dev, serial, buffer, device_status, nvstate, flash);
            });
    }

    #[task(binds = USB_LP_CAN_RX0, shared = [usb_dev, serial, buffer, device_status, nvstate, flash])]
    fn usb_rx0(cx: usb_rx0::Context) {
        let mut usb_dev = cx.shared.usb_dev;
        let mut serial = cx.shared.serial;
        let mut buffer = cx.shared.buffer;
        let mut device_status = cx.shared.device_status;
        let mut nvstate = cx.shared.nvstate;
        let mut flash = cx.shared.flash;

        (
            &mut usb_dev,
            &mut serial,
            &mut buffer,
            &mut device_status,
            &mut nvstate,
            &mut flash,
        )
            .lock(|usb_dev, serial, buffer, device_status, nvstate, flash| {
                usb_poll(usb_dev, serial, buffer, device_status, nvstate, flash);
            });
    }

    fn usb_poll<B: usb_device::bus::UsbBus>(
        usb_dev: &mut usb_device::prelude::UsbDevice<'static, B>,
        serial: &mut usbd_serial::SerialPort<'static, B>,
        buffer: &mut Vec<u8, USB_BUFFER_SIZE>,
        device_status: &StatusMessage,
        nvstate: &mut NVState,
        flash: &mut flash::Parts,
    ) {
        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 device_msg = DeviceMessage::Status(device_status.clone());
                            let bytes =
                                to_vec_cobs::<DeviceMessage, USB_BUFFER_SIZE>(&device_msg).unwrap();
                            serial.write(bytes.as_slice()).unwrap();
                        }
                        HostMessage::SetPLLSettings(settings) => {
                            nvstate.pll_settings = settings;
                            nvstate.save(flash);
                        }
                        HostMessage::GetPllSettings => {
                            defmt::error!("PLL output is not implemented yet")
                        }
                    },
                    Err(err) => defmt::error!("Unable to parse host message: {}", err),
                };

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