#![no_std]
#![no_main]

use panic_halt as _; // you can put a breakpoint on `rust_begin_unwind` to catch panics

use atsamd_hal::{
    clock::{ClockGenId, ClockSource, GenericClockController},
    pac::{interrupt, CorePeripherals, Peripherals},
    prelude::_atsamd_hal_embedded_hal_digital_v2_OutputPin,
    rtc::{Count32Mode, Rtc},
    timer::TimerCounter,
    usb::UsbBus,
};
use common::{handle_input, Handshake, UnsafeSync};
use core::cell::RefCell;
use cortex_m::{
    asm::delay as cycle_delay,
    interrupt::Mutex,
    peripheral::{NVIC, SCB},
};
use once_cell::unsync::OnceCell;
use smart_leds::{SmartLedsWrite, RGB8};
use usb_device::{bus::UsbBusAllocator, prelude::*};
use usbd_serial::{SerialPort, USB_CLASS_CDC};

mod bsp;
mod buttons;

use bsp::{entry, Dotstar, Pins};
use buttons::Buttons;

static USB_ALLOCATOR: UnsafeSync<OnceCell<UsbBusAllocator<UsbBus>>> =
    UnsafeSync::new(OnceCell::new());
static RTC_COUNT: Mutex<RefCell<Option<Rtc<Count32Mode>>>> = Mutex::new(RefCell::new(None));

static USB_BUS: Mutex<RefCell<Option<UsbDevice<UsbBus>>>> = Mutex::new(RefCell::new(None));
static USB_SERIAL: Mutex<RefCell<Option<SerialPort<UsbBus>>>> = Mutex::new(RefCell::new(None));

static HANDSHAKE: Mutex<RefCell<Handshake>> = Mutex::new(RefCell::new(Handshake::new()));

fn get_current_time() -> u32 {
    cortex_m::interrupt::free(|cs| {
        RTC_COUNT
            .borrow(cs)
            .borrow()
            .as_ref()
            .map(|rtc| rtc.count32())
            .unwrap_or(0)
    })
}

fn is_handshake_complete() -> bool {
    cortex_m::interrupt::free(|cs| HANDSHAKE.borrow(cs).borrow().is_complete())
}

#[entry]
fn main() -> ! {
    let mut peripherals = Peripherals::take().unwrap();
    let mut core = CorePeripherals::take().unwrap();
    // initializes gclk0 at 48MHz and gclk1 at 32KHz
    let mut clocks = GenericClockController::with_internal_32kosc(
        peripherals.GCLK,
        &mut peripherals.PM,
        &mut peripherals.SYSCTRL,
        &mut peripherals.NVMCTRL,
    );
    let pins = Pins::new(peripherals.PORT);

    let mut red_led = pins.d13.into_push_pull_output();
    red_led.set_low().unwrap();

    let usb_allocator = bsp::usb_allocator(
        peripherals.USB,
        &mut clocks,
        &mut peripherals.PM,
        pins.usb_dm,
        pins.usb_dp,
    )
    .unwrap();

    // SAFETY: No interrupt will access USB_ALLOCATOR before this point, since it is only
    // referenced from inside the USB_SERIAL mutex and the USB_BUS mutex, which have not been set
    // yet.
    //
    // SAFETY: This is the only line of the program that modifies USB_ALLOCATOR
    let allocator = unsafe {
        USB_ALLOCATOR.get().set(usb_allocator).ok().unwrap();
        USB_ALLOCATOR.get().get().unwrap()
    };

    cortex_m::interrupt::free(|cs| {
        let serial = USB_SERIAL.borrow(cs);
        let mut serial = serial.borrow_mut();
        *serial = Some(SerialPort::new(allocator));
    });

    cortex_m::interrupt::free(|cs| {
        let bus = USB_BUS.borrow(cs);
        let mut bus = bus.borrow_mut();
        *bus = Some(
            UsbDeviceBuilder::new(allocator, UsbVidPid(0x16c0, 0x27dd))
                .manufacturer("Aode Lion")
                .product("Trinket Streamdeck")
                .serial_number("AAAA")
                .device_class(USB_CLASS_CDC)
                .build(),
        );
    });

    unsafe {
        core.NVIC.set_priority(interrupt::USB, 1);
        NVIC::unmask(interrupt::USB);
    }

    let mut buttons = Buttons {
        d0: pins.d0,
        d1: pins.d1,
        d2: pins.d2,
        d3: pins.d3,
        // d4: pins.d4,
    }
    .init();

    // Use 32KHz clock for Rtc, can't use gclk1 for some reason
    let rtc_gclk = clocks
        .configure_gclk_divider_and_source(ClockGenId::GCLK3, 32, ClockSource::OSC32K, true)
        .unwrap();
    let rtc_clock = clocks.rtc(&rtc_gclk).unwrap();
    let mut counter = Rtc::count32_mode(peripherals.RTC, rtc_clock.freq(), &mut peripherals.PM);
    counter.set_count32(0);

    cortex_m::interrupt::free(|cs| {
        let rtc = RTC_COUNT.borrow(cs);
        *rtc.borrow_mut() = Some(counter);
    });

    // Use 48MHz clock for dotstar
    let dotstar_gclk = clocks.gclk0();
    let tc4_tc5 = &clocks.tc4_tc5(&dotstar_gclk).unwrap();
    let timer = TimerCounter::tc4_(tc4_tc5, peripherals.TC4, &mut peripherals.PM);

    let mut updated_at = 0;

    const DOTSTAR_DELAY: u32 = 1000;
    const DOTSTAR_OFF: RGB8 = RGB8 { r: 0, g: 0, b: 0 };

    let mut rgb = Dotstar {
        ci: pins.dotstar_ci,
        di: pins.dotstar_di,
        nc: pins.dotstar_nc,
    }
    .init(timer);
    rgb.write([DOTSTAR_OFF].iter().cloned()).unwrap();

    let mut dotstar_off = true;

    let mut output = [0; 8];

    loop {
        let current_time = get_current_time();

        if !is_handshake_complete() {
            cycle_delay(1024 * 15);
            continue;
        }

        if !dotstar_off && updated_at + DOTSTAR_DELAY < current_time {
            let _ = rgb.write([DOTSTAR_OFF].iter().cloned());
            dotstar_off = true;
        }

        if let Some(len) = buttons.tick(current_time, &mut output) {
            updated_at = current_time;
            let _ = write_serial(&output[0..len]);

            let byte = output[1..len]
                .iter()
                .fold(0u8, |byte, key| byte | (1 << key));

            dotstar_off = false;
            let color = color_from_byte(byte);
            let _ = rgb.write([color].iter().cloned());
        }
    }
}

#[allow(non_snake_case)]
#[interrupt]
fn USB() {
    poll_usb();
}

fn poll_usb() -> Option<()> {
    let ready = cortex_m::interrupt::free(|cs1| {
        cortex_m::interrupt::free(|cs2| {
            let mut usb_dev = USB_BUS.borrow(cs1).borrow_mut();
            let usb_dev = usb_dev.as_mut()?;
            let mut serial = USB_SERIAL.borrow(cs2).borrow_mut();
            let serial = serial.as_mut()?;

            Some(usb_dev.poll(&mut [&mut *serial]))
        })
    })?;

    if !ready {
        return Some(());
    }

    let mut input = [0u8; 32];

    let count = cortex_m::interrupt::free(|cs| {
        USB_SERIAL
            .borrow(cs)
            .borrow_mut()
            .as_mut()?
            .read(&mut input)
            .ok()
    })?;

    let current_time = get_current_time();

    cortex_m::interrupt::free(|cs| {
        let mut handshake = HANDSHAKE.borrow(cs).borrow_mut();

        if count == 32 {
            if let common::HandshakeResponse::Write { bytes } =
                handshake.perform(&input, current_time)?
            {
                write_serial(bytes)?;
            }

            return None;
        }

        Some(())
    })?;

    let mut output_buffer = [0; 256];

    let len = handle_input(&mut Device, &input, &mut output_buffer, count)?;

    if len == 0 {
        return None;
    }

    write_serial(&output_buffer[..len])?;

    Some(())
}

fn write_serial(bytes: &[u8]) -> Option<()> {
    cortex_m::interrupt::free(|cs| {
        let mut serial = USB_SERIAL.borrow(cs).borrow_mut();
        let serial = serial.as_mut()?;
        serial.write(bytes).ok()?;
        serial.flush().ok()
    })
}

struct Device;

impl common::Device<16, 1> for Device {
    // section 10.3.3 of datasheet for SAMD21
    fn serial_number(&mut self) -> [u8; 16] {
        let generator: [*const [u8; 4]; 4] = [
            0x0080A00C as *const [u8; 4],
            0x0080A040 as *const [u8; 4],
            0x0080A044 as *const [u8; 4],
            0x0080A048 as *const [u8; 4],
        ];

        let mut id = [0u8; 16];

        for i in 0..4 {
            id[i * 4] = unsafe { (*generator[i])[3] };
            id[i * 4 + 1] = unsafe { (*generator[i])[2] };
            id[i * 4 + 2] = unsafe { (*generator[i])[1] };
            id[i * 4 + 3] = unsafe { (*generator[i])[0] };
        }

        id
    }

    fn extras() -> [(&'static str, &'static str); 1] {
        [("version", env!("FIRMWARE_VERSION"))]
    }

    // Based on
    // - https://github.com/arduino/ArduinoCore-samd/issues/197
    // - https://github.com/tinygo-org/tinygo/blob/master/src/machine/board_trinket.go
    // - https://github.com/tinygo-org/tinygo/blob/master/src/machine/machine_atsamd21.go
    fn reset(&mut self) {
        cortex_m::interrupt::disable();

        let reset_magic_value: usize = 0xf01669ef;
        let reset_address: *mut usize = 0x20007FFC as *mut usize;

        unsafe {
            *reset_address = reset_magic_value;
        }

        SCB::sys_reset();
    }
}

// Expect byte to be in range 0..32
fn color_from_byte(byte: u8) -> RGB8 {
    let position: u8 = byte * 8; // max value here is 248

    match position {
        0..=85 => RGB8 {
            r: (255 - position * 3),
            g: (position * 3),
            b: 0,
        },
        86..=170 => {
            let position = position - 85;
            RGB8 {
                r: 0,
                g: (255 - position * 3),
                b: (position * 3),
            }
        }
        _ => {
            let position = position - 170;
            RGB8 {
                r: (position * 3),
                g: 0,
                b: (255 - position * 3),
            }
        }
    }
}