// TODO: Add license header
// TODO: Add documentation header
// TODO: GUI lock button support

#![no_std]
#![no_main]

mod layout;

use core::convert::Infallible;
use cortex_m::delay::Delay;
use embedded_hal::digital::v2::*;
use embedded_hal::timer::CountDown;
use fugit::ExtU32;
use panic_halt as _;
use smart_leds::{SmartLedsWrite, RGB8};
use usb_device::class_prelude::*;
use usb_device::prelude::*;
use usbd_human_interface_device::page::Keyboard;
use usbd_human_interface_device::prelude::*;
use waveshare_rp2040_zero::entry;
use waveshare_rp2040_zero::{
    hal::{
        clocks::{init_clocks_and_plls, Clock},
        pac,
        pio::PIOExt,
        timer::Timer,
        watchdog::Watchdog,
        Sio,
    },
    Pins, XOSC_CRYSTAL_FREQ,
};
use ws2812_pio::Ws2812;

pub const KEY_ROWS: usize = 4;
pub const KEY_COLS: usize = 12;
pub const NUMBER_OF_KEYS: usize = 42;

#[derive(Copy, Clone)]
struct ButtonMatrix {
    current_state: bool,
    previous_state: bool,
    fn_mode: u8,
}

impl ButtonMatrix {
    fn default() -> Self {
        Self {
            current_state: false,
            previous_state: false,
            fn_mode: 0,
        }
    }
}

#[entry]
fn main() -> ! {
    let mut pac = pac::Peripherals::take().unwrap();
    let core = pac::CorePeripherals::take().unwrap();
    let mut watchdog = Watchdog::new(pac.WATCHDOG);

    let clocks = init_clocks_and_plls(
        XOSC_CRYSTAL_FREQ,
        pac.XOSC,
        pac.CLOCKS,
        pac.PLL_SYS,
        pac.PLL_USB,
        &mut pac.RESETS,
        &mut watchdog,
    )
    .ok()
    .unwrap();

    let sio = Sio::new(pac.SIO);
    let pins = Pins::new(
        pac.IO_BANK0,
        pac.PADS_BANK0,
        sio.gpio_bank0,
        &mut pac.RESETS,
    );

    let usb_bus = UsbBusAllocator::new(waveshare_rp2040_zero::hal::usb::UsbBus::new(
        pac.USBCTRL_REGS,
        pac.USBCTRL_DPRAM,
        clocks.usb_clock,
        true,
        &mut pac.RESETS,
    ));

    let mut keyboard = UsbHidClassBuilder::new()
        .add_device(
            usbd_human_interface_device::device::keyboard::NKROBootKeyboardConfig::default(),
        )
        .build(&usb_bus);

    let mut usb_dev = UsbDeviceBuilder::new(&usb_bus, UsbVidPid(0x1209, 0x0001))
        .manufacturer("usbd-human-interface-device")
        .product("CMDR keyboard")
        .serial_number("0001")
        .build();

    let timer = Timer::new(pac.TIMER, &mut pac.RESETS);
    let mut delay = cortex_m::delay::Delay::new(core.SYST, clocks.system_clock.freq().to_Hz());

    // Configure the status LED
    let (mut pio, sm0, _, _, _) = pac.PIO0.split(&mut pac.RESETS);
    let mut status_led = Ws2812::new(
        pins.neopixel.into_mode(),
        &mut pio,
        sm0,
        clocks.peripheral_clock.freq(),
        timer.count_down(),
    );

    let mut button_matrix: [ButtonMatrix; NUMBER_OF_KEYS] =
        [ButtonMatrix::default(); NUMBER_OF_KEYS];

    let button_matrix_row_pins: &[&dyn InputPin<Error = Infallible>] = &[
        &pins.gp0.into_pull_up_input(),
        &pins.gp1.into_pull_up_input(),
        &pins.gp29.into_pull_up_input(),
        &pins.gp28.into_pull_up_input(),
    ];

    let button_matrix_col_pins: &mut [&mut dyn OutputPin<Error = Infallible>] = &mut [
        &mut pins.gp12.into_push_pull_output(),
        &mut pins.gp13.into_push_pull_output(),
        &mut pins.gp14.into_push_pull_output(),
        &mut pins.gp15.into_push_pull_output(),
        &mut pins.gp26.into_push_pull_output(),
        &mut pins.gp27.into_push_pull_output(),
        &mut pins.gp7.into_push_pull_output(),
        &mut pins.gp8.into_push_pull_output(),
        &mut pins.gp6.into_push_pull_output(),
        &mut pins.gp9.into_push_pull_output(),
        &mut pins.gp10.into_push_pull_output(),
        &mut pins.gp11.into_push_pull_output(),
    ];

    let mut report_count_down = timer.count_down();
    report_count_down.start(10.millis());

    let mut tick_count_down = timer.count_down();
    tick_count_down.start(1.millis());

    let mut indicator_count_down = timer.count_down();
    indicator_count_down.start(250.millis());

    let mut status_led_onoff: bool = false;
    let status_led_color: [RGB8; 5] = [
        (0, 0, 0).into(),   // Off
        (10, 7, 0).into(),  // Green
        (10, 4, 10).into(), // Blue
        (5, 10, 0).into(),  // Orange
        (2, 20, 0).into(),  // Red
    ];

    let mut caps_lock_active: bool = false;
    let mut fn_mode: u8 = 0;

    // Set all column pins to output and high
    init_button_matrix_pins(button_matrix_col_pins, &mut delay);

    loop {
        if indicator_count_down.wait().is_ok() {
            // Set status LED colour based on function layer and capslock
            // 0 = green, 1 = blue, 2 = orange, capslock active = flashing.
            if caps_lock_active == true && status_led_onoff == true {
                status_led
                    .write([status_led_color[0]].iter().copied())
                    .unwrap();
                status_led_onoff = false;
            } else {
                status_led
                    .write([status_led_color[usize::from(fn_mode) + 1]].iter().copied())
                    .unwrap();
                status_led_onoff = true;
            }
        }

        if report_count_down.wait().is_ok() {
            // Scan keyboard matrix
            let pressed_keys =
                get_pressed_buttons(button_matrix_row_pins, button_matrix_col_pins, &mut delay);

            // Get current function layer
            fn_mode = get_fn_mode(pressed_keys);

            // Set status LED colour based on function layer and capslock
            // 0 = green, 1 = blue, 2 = orange, capslock active = flashing.
            if caps_lock_active == false {
                status_led
                    .write([status_led_color[usize::from(fn_mode) + 1]].iter().copied())
                    .unwrap();
                status_led_onoff = true;
            }
            // Copy result from scanned keys to matrix struct
            for (index, key) in pressed_keys.iter().enumerate() {
                button_matrix[index].current_state = *key;
            }

            // Generate keyboard report
            let keyboard_report = get_keyboard_report(&mut button_matrix, layout::MAP, fn_mode);

            match keyboard.device().write_report(keyboard_report) {
                Err(UsbHidError::WouldBlock) => {}
                Err(UsbHidError::Duplicate) => {}
                Ok(_) => {}
                Err(e) => {
                    core::panic!("Failed to write keyboard report: {:?}", e)
                }
            };
        }

        //Tick once per ms
        if tick_count_down.wait().is_ok() {
            match keyboard.tick() {
                Err(UsbHidError::WouldBlock) => {}
                Ok(_) => {}
                Err(e) => {
                    core::panic!("Failed to process keyboard tick: {:?}", e)
                }
            };
        }

        if usb_dev.poll(&mut [&mut keyboard]) {
            match keyboard.device().read_report() {
                Err(UsbError::WouldBlock) => {
                    //do nothing
                }
                Err(e) => {
                    core::panic!("Failed to read keyboard report: {:?}", e)
                }
                Ok(leds) => {
                    if leds.caps_lock == true {
                        caps_lock_active = true;
                    } else {
                        caps_lock_active = false;
                    }
                }
            }
        }
    }
}

// Initialise keyboard matrix pins
fn init_button_matrix_pins(cols: &mut [&mut dyn OutputPin<Error = Infallible>], delay: &mut Delay) {
    for col in cols.iter_mut() {
        col.set_high().unwrap();
    }
    delay.delay_us(10);
}

// Scan keyboard matrix for pressed keys and return a bool array
// representing the state of each key (true = pressed)
// TODO: This is a bit of a mess, needs refactoring
fn get_pressed_buttons(
    rows: &[&dyn InputPin<Error = Infallible>],
    cols: &mut [&mut dyn OutputPin<Error = Infallible>],
    delay: &mut Delay,
) -> [bool; NUMBER_OF_KEYS] {
    // Scan keyboard matrix for pressed keys
    let mut pressed_keys: [bool; NUMBER_OF_KEYS] = [false; NUMBER_OF_KEYS];
    for (col_index, col) in cols.iter_mut().enumerate() {
        // Activate column
        col.set_low().unwrap();
        delay.delay_us(10);

        // Read rows
        for (row_index, row) in rows.iter().enumerate() {
            // Do not check unconnected keys in the matrix
            if row_index == 3 && (col_index < 3 || col_index > 8) {
                continue;
            }
            if row_index < 3 && row.is_low().unwrap() {
                pressed_keys[col_index + (row_index * KEY_COLS)] = true;
            } else if row.is_low().unwrap() {
                // Correct index for unconnected keys
                pressed_keys[col_index + (row_index * KEY_COLS) - 3] = true;
            }
        }
        col.set_high().unwrap();
        delay.delay_us(10);
    }
    // Return scan result
    pressed_keys
}

// Get current Fn mode (0, 1 or 2)
fn get_fn_mode(pressed_keys: [bool; NUMBER_OF_KEYS]) -> u8 {
    // Check Fn mode
    let mut fn_mode: u8 = 0;
    for button_id in layout::FN_BUTTONS.iter() {
        if pressed_keys[usize::from(*button_id)] == true {
            fn_mode += 1;
        }
    }
    fn_mode
}

// Generate keyboard report based on pressed keys and Fn mode (0, 1 or 2)
fn get_keyboard_report(
    matrix_keys: &mut [ButtonMatrix; NUMBER_OF_KEYS],
    layout_map: [[Keyboard; NUMBER_OF_KEYS]; 3],
    fn_mode: u8,
) -> [Keyboard; NUMBER_OF_KEYS] {
    // Create default report
    let mut keyboard_report: [Keyboard; NUMBER_OF_KEYS] =
        [Keyboard::NoEventIndicated; NUMBER_OF_KEYS];
    // Filter report based on Fn mode and pressed keys
    for (index, key) in matrix_keys.iter_mut().enumerate() {
        if key.current_state != key.previous_state && key.current_state == true {
            key.fn_mode = fn_mode;
        }
        key.previous_state = key.current_state;

        if key.current_state == true {
            keyboard_report[index] = layout_map[usize::from(key.fn_mode)][index];
        }
    }
    // Return report
    keyboard_report
}
