#![no_std]
#![no_main]
extern crate alloc;

use crate::hal::peripherals::CAN1;
use canapi::id::{plant_id, IDENTIFY_CMD_OFFSET, MOISTURE_DATA_OFFSET};
use canapi::SensorSlot;
use ch32_hal::adc::{Adc, SampleTime, ADC_MAX};
use ch32_hal::{pac};
use ch32_hal::can::{Can, CanFifo, CanFilter, CanFrame, CanMode};
use ch32_hal::gpio::{Flex, Level, Output, Pull, Speed};
use ch32_hal::mode::NonBlocking;
use ch32_hal::peripherals::USBD;
use core::fmt::Write as _;
use core::sync::atomic::{AtomicBool, Ordering};
use embassy_executor::{task, Spawner};
use embassy_futures::yield_now;
use embassy_sync::blocking_mutex::raw::CriticalSectionRawMutex;
use embassy_sync::channel::Channel;
use embassy_time::{Duration, Instant, Timer};
use embassy_usb::class::cdc_acm::{CdcAcmClass, State};
use embassy_usb::{Builder, UsbDevice};
use embedded_alloc::LlffHeap as Heap;
use embedded_can::nb::Can as nb_can;
use embedded_can::{Id, StandardId};
use hal::bind_interrupts;
use hal::usbd::Driver;
use {ch32_hal as hal, panic_halt as _};

macro_rules! mk_static {
    ($t:ty,$val:expr) => {{
        static STATIC_CELL: static_cell::StaticCell<$t> = static_cell::StaticCell::new();
        #[deny(unused_attributes)]
        let x = STATIC_CELL.uninit().write(($val));
        x
    }};
}

bind_interrupts!(struct Irqs {
    USB_LP_CAN1_RX0 => hal::usbd::InterruptHandler<hal::peripherals::USBD>;
});

#[global_allocator]
static HEAP: Heap = Heap::empty();
static LOG_CH: Channel<CriticalSectionRawMutex, heapless::String<128>, 8> = Channel::new();
static CAN_TX_CH: Channel<CriticalSectionRawMutex, CanFrame, 4> = Channel::new();

static BEACON: AtomicBool = AtomicBool::new(false);

#[embassy_executor::main(entry = "qingke_rt::entry")]
async fn main(spawner: Spawner) {
    ch32_hal::pac::AFIO.pcfr1().write(|w| w.set_can1_rm(2));

    unsafe {
        #[allow(static_mut_refs)]
        static mut HEAP_SPACE: [u8; 4096] = [0; 4096]; // 4 KiB heap
        #[allow(static_mut_refs)]
        HEAP.init(HEAP_SPACE.as_ptr() as usize, HEAP_SPACE.len());
    }

    let p = hal::init(hal::Config {
        rcc: hal::rcc::Config::SYSCLK_FREQ_144MHZ_HSI,
        ..Default::default()
    });

    // Build driver and USB stack using 'static buffers
    let driver = Driver::new(p.USBD, Irqs, p.PA12, p.PA11);

    let mut probe_gnd = Flex::new(p.PA2);
    probe_gnd.set_as_input(Pull::None);

    // Create GPIO for 555 Q output (PB0)
    let q_out = Output::new(p.PA0, Level::Low, Speed::Low);
    let info = mk_static!(Output, Output::new(p.PA10, Level::Low, Speed::Low));
    let warn = mk_static!(Output, Output::new(p.PA9, Level::Low, Speed::Low));

    // Read configuration switches on PB3..PB7 at startup with floating detection
    // PB3: Sensor A/B selector (Low=A, High=B)
    // PB4..PB7: address bits (1,2,4,8)
    let mut sensor_ab_pin = Flex::new(p.PA3);
    let mut sensor_address_bit_1_pin = Flex::new(p.PA4);
    let mut sensor_address_bit_2_pin = Flex::new(p.PA5);
    let mut sensor_address_bit_3_pin = Flex::new(p.PA6);
    let mut sensor_address_bit_4_pin = Flex::new(p.PA7);

    // Validate all config pins; if any is floating, stay in an error loop until fixed
    // Try read PB3..PB7
    let sensor_ab_config = detect_stable_pin(&mut sensor_ab_pin).await;
    let sensor_address_bit_1_config = detect_stable_pin(&mut sensor_address_bit_1_pin).await;
    let sensor_address_bit_2_config = detect_stable_pin(&mut sensor_address_bit_2_pin).await;
    let sensor_address_bit_3_config = detect_stable_pin(&mut sensor_address_bit_3_pin).await;
    let sensor_address_bit_4_config = detect_stable_pin(&mut sensor_address_bit_4_pin).await;

    let slot = if sensor_ab_config.unwrap_or(false) {
        SensorSlot::B
    } else {
        SensorSlot::A
    };
    let mut addr: u8 = 0;
    if sensor_address_bit_1_config.unwrap_or(false) {
        addr |= 1;
    }
    if sensor_address_bit_2_config.unwrap_or(false) {
        addr |= 2;
    }
    if sensor_address_bit_3_config.unwrap_or(false) {
        addr |= 4;
    }
    if sensor_address_bit_4_config.unwrap_or(false) {
        addr |= 8;
    }
    let moisture_id = plant_id(MOISTURE_DATA_OFFSET, slot, addr as u16);
    let identify_id = plant_id(IDENTIFY_CMD_OFFSET, slot, addr as u16);
    let standard_identify_id = StandardId::new(identify_id).unwrap();

    //is any floating, or invalid addr (only 1-8 are valid)
    let invalid_config = sensor_ab_config.is_none()
        || sensor_address_bit_1_config.is_none()
        || sensor_address_bit_2_config.is_none()
        || sensor_address_bit_3_config.is_none()
        || sensor_address_bit_4_config.is_none() 
        || addr == 0 
        || addr > 8;

    let mut config = embassy_usb::Config::new(0xC0DE, 0xCAFE);
    config.manufacturer = Some("Can Sensor v0.2");
    let msg = mk_static!(heapless::String<128>, heapless::String::new());
    if invalid_config {
        let _ = core::fmt::Write::write_fmt(
            msg,
            format_args!(
                "CFG err: {:?} {:?} {:?} {:?} {:?}",
                to_info(sensor_ab_config), to_info(sensor_address_bit_1_config), to_info(sensor_address_bit_2_config), to_info(sensor_address_bit_3_config), to_info(sensor_address_bit_4_config)
            ),
        );

    } else {
        let _ = core::fmt::Write::write_fmt(msg, format_args!("Sensor {:?} plant {}", slot, addr));
    }

    config.product = Some(msg.as_str());
    config.serial_number = Some("12345678");
    config.max_power = 100;
    config.max_packet_size_0 = 64;

    // Windows compatibility requires these; CDC-ACM
    config.device_class = 0x02;
    config.device_sub_class = 0x02;
    config.device_protocol = 0x00;
    config.composite_with_iads = false;

    let mut builder = Builder::new(
        driver,
        config,
        mk_static!([u8; 256], [0; 256]),
        mk_static!([u8; 256], [0; 256]),
        &mut [], // no msos descriptors
        mk_static!([u8; 64], [0; 64]),
    );
    // Initialize CDC state and create CDC-ACM class
    let class = mk_static!(
        CdcAcmClass<'static, Driver<'static, hal::peripherals::USBD>>,
        CdcAcmClass::new(&mut builder, mk_static!(State, State::new()), 64)
    );

    // Build USB device
    let usb = mk_static!(UsbDevice<Driver<USBD>>, builder.build());

    spawner.spawn(usb_task(usb)).unwrap();
    spawner.spawn(usb_writer(class)).unwrap();

    if invalid_config {
        // At least one floating: report and blink code for the first one found.
        let mut msg: heapless::String<128> = heapless::String::new();
        let code = if sensor_address_bit_1_config.is_none() {
            1
        } else if sensor_address_bit_2_config.is_none() {
            2
        } else if sensor_address_bit_3_config.is_none() {
            3
        } else if sensor_address_bit_4_config.is_none() {
            4
        } else if sensor_ab_config.is_none() {
            5
        } else {
            6 // Invalid address (0 or > 8)
        };
        let which = match code {
            1 => "PB4 (bit 1)",
            2 => "PB5 (bit 2)",
            3 => "PB6 (bit 3)",
            4 => "PB7 (bit 4)",
            5 => "PB3 (A/B)",
            _ => "Address (0 or > 8)",
        };
        if code == 6 {
            let _ = core::fmt::Write::write_fmt(
                &mut msg,
                format_args!(
                    "Invalid address {} (only 1-8 allowed) -> blinking code {}. Fix jumpers.\r\n",
                    addr, code
                ),
            );
        } else {
            let _ = core::fmt::Write::write_fmt(
                &mut msg,
                format_args!(
                    "Config pin floating detected on {} -> blinking code {}. Fix jumpers.\r\n",
                    which, code
                ),
            );
        }
        log(msg);
        spawner.spawn(blink_error_task(warn, info, 2, code)).unwrap();
    } else {
        // Log startup configuration and derived CAN IDs

        {
            let mut msg: heapless::String<128> = heapless::String::new();
            let slot_chr = match slot {
                SensorSlot::A => 'a',
                SensorSlot::B => 'b',
            };
            let _ = core::fmt::Write::write_fmt(
                &mut msg,
                format_args!(
                    "Startup: slot={} addr={} moisture_id=0x{:03X} identity_id=0x{:03X}\r\n",
                    slot_chr, addr, moisture_id, identify_id
                ),
            );
            log(msg);
        }

        // Create ADC on ADC1 and use PA1 as analog input (Threshold/Trigger)
        let adc = Adc::new(p.ADC1, Default::default());
        let ain = p.PA1;
        let config = Default::default();
        let can: Can<CAN1, NonBlocking> = Can::new_nb(
            p.CAN1,
            p.PB8,
            p.PB9,
            CanFifo::Fifo0,
            CanMode::Normal,
            20_000,
            config,
        )
        .expect("Valid");
        ch32_hal::pac::AFIO.pcfr1().write(|w| w.set_can1_rm(2));

        can.add_filter(CanFilter::accept_all());
        
        // Improve CAN robustness for longer cables:
        // 1. Enable Automatic Bus-Off Management (ABOM)
        // 2. Disable Automatic Retransmission (NART) as we send regular measurements anyway
        // 3. Enable Receive FIFO Overwrite Mode (RFLM = 0, default)
        // 4. Increase Resync Jump Width (SJW) if possible by patching BTIMR
        hal::pac::CAN1.ctlr().modify(|w| {
            w.set_abom(false);
            w.set_nart(true); 
        });
        
        // SJW is bits 24-25 of BTIMR. HAL sets it to 0 (SJW=1). 
        // Let's try to set it to 3 (SJW=4) for better jitter tolerance.
        hal::pac::CAN1.btimr().modify(|w| {
            w.set_sjw(3); // 3 means 4TQ
        });

       // let mut filter = CanFilter::new_id_list();
       // filter.get(0).unwrap().set(Id::Standard(standard_identify_id), Default::default());
       // can.add_filter(filter);
        let standard_moisture_id = StandardId::new(moisture_id).unwrap();
        spawner
            .spawn(can_task(can,info, warn, standard_identify_id, standard_moisture_id))
            .unwrap();

        // move Q output, LED, ADC and analog input into worker task
        spawner
            .spawn(worker(
                probe_gnd,
                q_out,
                adc,
                ain,
                standard_moisture_id,
                standard_identify_id,
            ))
            .unwrap();
    }
    // Prevent main from exiting
    core::future::pending::<()>().await;
}

fn to_info(res: Option<bool>) -> i8 {
    match res {
        Some(true) => 1,
        Some(false) => -1,
        None => 0,
    }
}

// Helper closure: detect stable pin by comparing readings under Pull::Down and Pull::Up
async fn detect_stable_pin(pin: &mut Flex<'static>) -> Option<bool> {
    pin.set_as_input(Pull::Down);
    Timer::after_millis(2).await;
    let low_read = pin.is_high();
    pin.set_as_input(Pull::Up);
    Timer::after_millis(2).await;
    let high_read = pin.is_high();
    if low_read == high_read {
        Some(high_read)
    } else {
        None
    }
}
async fn blink_error_loop(info_led: &mut Output<'static>, warn_led: &mut Output<'static>, c_i: u8, c_w: u8) -> ! {
    for _loop_count in 0..5 {
        // code: 1-4 for PB4..PB7, 5 for PB3 (A/B), 7 for CAN address collision
        for _ in 0..c_i {
            info_led.set_high();
            Timer::after_millis(200).await;
            info_led.set_low();
            Timer::after_millis(200).await;
        }
        for _ in 0..c_w {
            warn_led.set_high();
            Timer::after_millis(200).await;
            warn_led.set_low();
            Timer::after_millis(200).await;
        }
        // Pause between sequences
        Timer::after_millis(400).await;
    }

    for _ in 0..5 {
        info_led.set_high();
        Timer::after_millis(50).await;
        info_led.set_low();
        Timer::after_millis(50).await;
        warn_led.set_high();
        Timer::after_millis(50).await;
        warn_led.set_low();
        Timer::after_millis(50).await;
    }

    pac::PFIC.cfgr().modify(|w| {
        w.set_resetsys(true);
        w.set_keycode(pac::pfic::vals::Keycode::KEY3); // KEY3 is 0xBEEF, the System Reset key
    });
    loop{

    }
}

#[task]
async fn blink_error_task(info_led: &'static mut Output<'static>, warn_led: &'static mut Output<'static>, c_i: u8, c_w: u8) -> ! {
    blink_error_loop(info_led, warn_led, c_i, c_w).await
}

#[task]
async fn can_task(
    mut can: Can<'static, CAN1, NonBlocking>,
    info: &'static mut Output<'static>,
    warn: &'static mut Output<'static>,
    identify_id: StandardId,
    moisture_id: StandardId,
) {
    // Non-blocking beacon blink timing.
    // We keep this inside the CAN task so it can't stall other tasks (like `worker`) with `await`s.
    let mut next_beacon_toggle = Instant::now();
    let beacon_period = Duration::from_millis(50);

    loop {
        match can.receive() {
            Ok(frame) => {
                match frame.id() {
                    Id::Standard(s_frame) => {
                        let mut msg: heapless::String<128> = heapless::String::new();
                        let _ = write!(
                            &mut msg,
                            "Received from canbus: {:?} ident is {:?} \r\n",
                            s_frame.as_raw(),
                            identify_id.as_raw()
                        );
                        log(msg);
                        if s_frame.as_raw() == moisture_id.as_raw() {
                            //trigger collision detection on other side as well
                            let _ = can.transmit(&frame);
                            // We should never receive moisture packets addressed to ourselves.
                            // If we do, another node likely uses the same jumper configuration.
                            blink_error_loop(info, warn, 1,2).await;
                        }
                        if s_frame.as_raw() == identify_id.as_raw() {
                            BEACON.store(true, Ordering::Relaxed);
                        }
                    }
                    Id::Extended(_) => {}
                }
            }
            Err(nb::Error::WouldBlock) => {
                // No frame available
            }
            Err(nb::Error::Other(err)) => {
                for _ in 0..3 {
                    warn.set_high();
                    Timer::after_millis(100).await;
                    warn.set_low();
                    Timer::after_millis(100).await;
                }
                let mut msg: heapless::String<128> = heapless::String::new();
                let _ = write!(&mut msg, "rx err {:?}", err);
                log(msg);
            }
        }

        if BEACON.load(Ordering::Relaxed) {
            let now = Instant::now();
            if now >= next_beacon_toggle {
                info.toggle();
                // Move the schedule forward; if we fell behind, resync to "now".
                next_beacon_toggle = now + beacon_period;
            }
        }

        while let Ok(mut frame) = CAN_TX_CH.try_receive() {
            match can.transmit(&mut frame) {
                Ok(..) => {
                }
                Err(nb::Error::WouldBlock) => {
                    for _ in 0..2 {
                        warn.set_high();
                        Timer::after_millis(100).await;
                        warn.set_low();
                        Timer::after_millis(100).await;
                    }
                    let mut msg: heapless::String<128> = heapless::String::new();
                    let _ = write!(&mut msg, "canbus out buffer full");
                    log(msg);
                }
                Err(nb::Error::Other(err)) => {
                    for _ in 0..3 {
                        warn.set_high();
                        Timer::after_millis(100).await;
                        warn.set_low();
                        Timer::after_millis(100).await;
                    }
                    let mut msg: heapless::String<128> = heapless::String::new();
                    let _ = write!(&mut msg, "tx err {:?}", err);
                    log(msg);
                }
            }
        }
        yield_now().await;
    }
}

#[task]
async fn worker(
    mut probe_gnd: Flex<'static>,
    mut q: Output<'static>,
    mut adc: Adc<'static, hal::peripherals::ADC1>,
    mut ain: hal::peripherals::PA1,
    moisture_id: StandardId,
    identify_id: StandardId,
) {
    // 555 emulation state: Q initially Low
    let mut q_high = false;
    let low_th: u16 = (ADC_MAX as u16) / 3; // ~1/3 Vref
    let high_th: u16 = ((ADC_MAX as u32 * 2) / 3) as u16; // ~2/3 Vref

    loop {
        // Count rising edges of Q in a 100 ms window
        let start = Instant::now();
        let mut pulses: u32 = 0;
        let mut last_q = q_high;

        probe_gnd.set_as_output(Speed::Low);
        probe_gnd.set_low();
        let probe_duration = Duration::from_millis(100);
        while Instant::now()
            .checked_duration_since(start)
            .unwrap_or(Duration::from_millis(0))
            < probe_duration
        {
            // Sample the analog input (Threshold/Trigger on A1)
            let val: u16 = adc.convert(&mut ain, SampleTime::CYCLES28_5);

            // 555 core behavior:
            // - If input <= 1/3 Vref => set Q high (trigger)
            // - If input >= 2/3 Vref => set Q low (threshold)
            // - Otherwise keep previous Q state (hysteresis)
            if val <= low_th {
                q_high = true;
            } else if val >= high_th {
                q_high = false;
            }

            // Drive output pin accordingly
            if q_high {
                q.set_high();
            } else {
                q.set_low();
            }

            // Count rising edges
            if !last_q && q_high {
                pulses = pulses.saturating_add(1);
            }
            last_q = q_high;

            // Yield to allow USB and other tasks to run
            yield_now().await;
        }
        probe_gnd.set_as_input(Pull::None);

        let freq_hz: u32 = pulses * (1000 / probe_duration.as_millis()) as u32; // pulses per 0.1s => Hz

        let mut msg: heapless::String<128> = heapless::String::new();
        let _ = write!(
            &mut msg,
            "555 window={}ms pulses={} freq={} Hz (A1->Q on PB0) id={:?}\r\n",
            probe_duration.as_millis(),
            pulses,
            freq_hz,
            identify_id.as_raw()
        );
        log(msg);

        let moisture = CanFrame::new(moisture_id, &(freq_hz as u32).to_be_bytes()).unwrap();
        CAN_TX_CH.send(moisture).await;
    }
}

fn log(message: heapless::String<128>) {
    match LOG_CH.try_send(message) {
        Ok(_) => {}
        Err(_) => {}
    }
}

#[task]
async fn usb_task(usb: &'static mut UsbDevice<'static, Driver<'static, hal::peripherals::USBD>>) {
    usb.run().await;
}

#[task]
async fn usb_writer(
    class: &'static mut CdcAcmClass<'static, Driver<'static, hal::peripherals::USBD>>,
) {
    loop {
        class.wait_connection().await;
        printer(class).await;
    }
}

async fn printer(class: &mut CdcAcmClass<'static, Driver<'static, USBD>>) {
    loop {
        let msg = LOG_CH.receive().await;
        match class.write_packet(msg.as_bytes()).await {
            Ok(_) => {}
            Err(_) => {
                // Disconnected or endpoint disabled
                return;
            }
        }
    }
}