#![no_std]
#![no_main]
#![feature(type_alias_impl_trait, future_join, abi_avr_interrupt, impl_trait_in_assoc_type)]
#![feature(int_roundings)]
#![feature(isqrt)]

extern crate alloc;

use alloc::rc::Rc;
use core::convert::Infallible;
use core::fmt::Write;
use core::mem::MaybeUninit;
use core::ops::Deref;
use core::panic::PanicInfo;
use core::sync::atomic::Ordering;
use embedded_io::Write as _;
use arduino_hal::{DefaultClock, Pins};
use atmega_hal::pac::USART0;
use atmega_hal::port::{Dynamic, PB5, PD4, Pin};
use avr_async_serial::UsartDriver;
use avr_hal_generic::avr_device;
use avr_hal_generic::port::mode::Output;
use avr_hal_generic::usart::Baudrate;
use avr_tc1_embassy_time::{define_interrupt, init_system_time};
use ehlcd2d::{DisplayControl, EntryMode, HalfWidthBus, LcdPinConfiguration, Lines};
use ehlcd2d::nonblocking::Lcd;
use embassy_executor::Spawner;
use embassy_sync::blocking_mutex::raw::NoopRawMutex;
use embassy_sync::mutex::Mutex;
use embassy_time::{Duration, Instant, Ticker, Timer, with_timeout};
use embedded_alloc::Heap;
use embedded_hal::blocking::delay::DelayUs;
use embedded_hal::digital::v2::OutputPin;
use embedded_hal_async::delay::DelayNs;
use ufmt::uWrite;
use embedded_io_async::{Read, Seek, SeekFrom, Write as AsyncWrite};
use itertools::Itertools;
use portable_atomic::AtomicBool;


type Se = UsartDriver<USART0>;
define_interrupt!(atmega328p);


const HEAP_SIZE: usize = 64;
static mut HEAP_MEM: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
#[global_allocator]
static HEAP: Heap = Heap::empty();
#[embassy_executor::task]
async fn serial(mut serial: Se, mut led: Pin<Output>, status: Rc<Mutex<NoopRawMutex, (u16, u16)>>) {
    led.set_high();
    let mut write = serial.write().unwrap();
    let mut read = serial.read().unwrap();

    let mut last_status_update = Instant::now();
    let mut last_receive: u8 = 0;
    let mut update_drops = 0;
    let mut updates = 0;
    loop {
        let idx;
        {
            let mut header = [0u8; 1];
            while header[0] & 0b10000000 == 0 {
                let result = with_timeout(Duration::from_secs(1), read.read_exact(&mut header)).await;
                if result.is_err() {
                    DISCONNECTED.store(true, Ordering::SeqCst);
                }
            }
            idx = header[0] & 0b01111111;
            if (last_receive + 1) & 0b01111111 != idx {
                update_drops += 1;
            }
        }

        let mut data = [0u8; 12];
        read.read_exact(&mut data).await.unwrap();
        let mut decoded = [0u32; 3];
        'a: {
            for i in 0..3 {
                for j in (i * 4)..(i * 4 + 4) {
                    if data[j] & 0b10000000 != 0 {
                        break 'a;
                    }
                    decoded[i] <<= 7;
                    decoded[i] |= data[j] as u32;
                }
            }
            last_receive = idx;
            updates += 1;
            *status.deref().lock().await = (decoded[0] as u16, decoded[1] as u16);
        }
        if Instant::now() - last_status_update > Duration::from_secs(1) {
            //*status.deref().lock().await = (updates, update_drops);
            updates = 0;
            update_drops = 0;
            last_status_update = Instant::now();
        }

    }
}

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

// #[embassy_executor::task]
// async fn blink(mut led: Pin<Output, PB5>) {
//     let mut ticker = Ticker::every(Duration::from_millis(50));
//     loop {
//         led.toggle();
//         ticker.next().await;
//     }
// }

#[embassy_executor::task]
async fn lcd_update(mut lcd: Lcd<
    Pin<Output, atmega_hal::port::PB5>,
    Pin<Output, atmega_hal::port::PB6>,
    HalfWidthBus<
        Pin<Output, atmega_hal::port::PB4>,
        Pin<Output, atmega_hal::port::PH6>,
        Pin<Output, atmega_hal::port::PH5>,
        Pin<Output, atmega_hal::port::PH4>
    >,
    EmbassyDelayNs,
    Infallible
>, status: Rc<Mutex<NoopRawMutex, (u16, u16)>>) {

    let mut ticker = Ticker::every(Duration::from_millis(100));
    let mut b0 = [0u8; 32];
    let mut b1 = [0u8; 32];
    let mut previous_buffer = &mut b0;
    let mut next_buffer = &mut b1;

    lcd.set_display_control(DisplayControl::default()).await.ok();
    loop {
        ticker.next().await;
        next_buffer.fill(0);
        //lcd_first_line(next_buffer).await;
        lcd_second_line(&mut next_buffer[16..], &status).await;
        let mut current_address = 0;
        lcd.seek(SeekFrom::Start(current_address)).await.ok();
        for (lcd_addr, array_addr) in (0u64..16).merge((40u64..56)).zip(0usize..32) {
            if previous_buffer[array_addr] != next_buffer[array_addr] {
                if lcd_addr != current_address {
                    lcd.seek(SeekFrom::Start(lcd_addr)).await.ok();
                    current_address = lcd_addr;
                }
                current_address += 1;
                lcd.write_char(next_buffer[array_addr]).await.ok();
            }
        }
        (previous_buffer, next_buffer) = (next_buffer, previous_buffer);
    }
}

const START: Instant = Instant::from_ticks(0);

// async fn lcd_first_line(mut buffer: &mut [u8]) {
//     let hour = Duration::from_secs(60 * 64).as_millis();
//     let minute = Duration::from_secs(60).as_millis();
//     let second = Duration::from_secs(1).as_millis();
//     let mut now = Instant::now().duration_since(START).as_millis();
//     let hours = now / hour;
//     now %= hour;
//     let minutes = now / minute;
//     now %= minute;
//     let seconds = now / second;
//     now %= second;
//     write!(buffer, "{hours:02}:{minutes:02}:{seconds:02}.{now:03}").ok();
// }

#[inline(always)]
async fn lcd_second_line(mut buffer: &mut [u8], status: &Mutex<NoopRawMutex, (u16, u16)>) {
    let (polling_rate, drop_count) = *status.lock().await;
    if polling_rate != 0 {
        let drop_percentage = (drop_count as u32) * 100 / polling_rate as u32;
        write!(buffer, "x: {polling_rate:04} y: {drop_percentage:04}").unwrap();
        //write!(buffer, "{polling_rate:03}Hz : {drop_percentage:03}%").unwrap();
    } else {
        write!(buffer, "NO CONN!").ok();
    }
}

#[embassy_executor::main]
async fn main(spawner: Spawner) {
    // Initialize the allocator BEFORE you use it
    unsafe { HEAP.init(HEAP_MEM.as_ptr() as usize, HEAP_SIZE) }
    let mut dp = arduino_hal::Peripherals::take().unwrap();
    let pins: Pins = arduino_hal::pins!(dp);
    dp.CPU.smcr.write(|w| w.se().set_bit());
    init_system_time(&mut dp.TC1);
    let mut serial1 = UsartDriver::new(dp.USART0, Baudrate::<DefaultClock>::new(57600));
    //spawner.spawn(blink(pins.d13.into_output())).unwrap();
    let lcd = Lcd::new(
        LcdPinConfiguration {
            en: pins.d11.into_output(),
            rs: pins.d12.into_output(),
            bus: HalfWidthBus {
                d4: pins.d10.into_output(),
                d5: pins.d9.into_output(),
                d6: pins.d8.into_output(),
                d7: pins.d7.into_output(),
            },
        },
        EmbassyDelayNs,
        Lines::TwoLines,
        EntryMode::default()
    ).await.unwrap();

    let status = Rc::new(Mutex::new((0, 0)));
    spawner.spawn(lcd_update(lcd, status.clone())).unwrap();
    //let stop_channel: Rc<Channel<NoopRawMutex, bool, 1>> = Rc::new(Channel::new());
    spawner.spawn(serial(serial1, pins.d13.into_output().downgrade(), status)).unwrap();

}

pub struct EmbassyDelayNs;

impl DelayNs for EmbassyDelayNs {
    async fn delay_ns(&mut self, ns: u32) {
        embassy_time::Timer::after_micros(ns.div_ceil(1000) as u64).await;
    }

    async fn delay_us(&mut self, us: u32) {
        embassy_time::Timer::after_micros(us as u64).await;
    }

    async fn delay_ms(&mut self, ms: u32) {
        embassy_time::Timer::after_millis(ms as u64).await;
    }
}


#[panic_handler]
fn panic(info: &PanicInfo) -> ! {
    avr_device::interrupt::disable();
    let dp = unsafe { arduino_hal::Peripherals::steal() };
    let pins = arduino_hal::pins!(dp);
    let serial = arduino_hal::default_serial!(dp, pins, 57600);
    let mut writer = Writer {
        u_write: serial
    };
    writeln!(&mut writer, "{}", info).ok();
    loop {

    }
}

struct Writer<T> {
    u_write: T
}

impl<T: uWrite> Write for Writer<T> {
    fn write_str(&mut self, s: &str) -> core::fmt::Result {
        self.u_write.write_str(s).map_err(|x| match x {
            _ => core::fmt::Error
        })
    }

    fn write_char(&mut self, c: char) -> core::fmt::Result {
        self.u_write.write_char(c).map_err(|x| match x {
            _ => core::fmt::Error
        })
    }
}