//! Project: CMtec CMDR Keyboard 51
//! Date:    2025-03-09
//! Author:  Christoffer Martinsson
//! Email:   cm@cmtec.se
//! License: Please refer to LICENSE in root directory

#![no_std]
#![no_main]

mod button_matrix;
mod layout;
mod status_led;

use button_matrix::ButtonMatrix;
use core::convert::Infallible;
use cortex_m::delay::Delay;
use embedded_hal::digital::{InputPin, OutputPin};
use embedded_hal_0_2::timer::CountDown;
use fugit::ExtU32;
use panic_halt as _;
use rp2040_hal::{
    Sio,
    clocks::{Clock, init_clocks_and_plls},
    gpio::{AnyPin, Pins},
    pac,
    pio::{PIOExt, StateMachineIndex},
    timer::Timer,
    watchdog::Watchdog,
};
use status_led::{StatusMode, Ws2812StatusLed};
use usb_device::class_prelude::*;
use usb_device::prelude::*;
use usbd_human_interface_device::page::Keyboard;
use usbd_human_interface_device::prelude::*;

// The linker will place this boot block at the start of our program image. We
/// need this to help the ROM bootloader get our code up and running.
#[unsafe(link_section = ".boot2")]
#[unsafe(no_mangle)]
#[used]
pub static BOOT2_FIRMWARE: [u8; 256] = rp2040_boot2::BOOT_LOADER_W25Q080;

const XTAL_FREQ_HZ: u32 = 12_000_000u32;

// Public constants
pub const KEY_ROWS: usize = 5;
pub const KEY_COLS: usize = 12;
pub const NUMBER_OF_KEYS: usize = KEY_ROWS * KEY_COLS;

// Public types
#[derive(Copy, Clone, Default)]
pub struct KeyboardButton {
    pub pressed: bool,
    pub previous_pressed: bool,
    pub fn_mode: u8,
}

#[rp2040_hal::entry]
fn main() -> ! {
    // Grab our singleton objects
    let mut pac = pac::Peripherals::take().unwrap();

    // Set up the watchdog driver - needed by the clock setup code
    let mut watchdog = Watchdog::new(pac.WATCHDOG);

    // Configure clocks and PLLs
    let clocks = init_clocks_and_plls(
        XTAL_FREQ_HZ,
        pac.XOSC,
        pac.CLOCKS,
        pac.PLL_SYS,
        pac.PLL_USB,
        &mut pac.RESETS,
        &mut watchdog,
    )
    .ok()
    .unwrap();

    let core = pac::CorePeripherals::take().unwrap();

    // The single-cycle I/O block controls our GPIO pins
    let sio = Sio::new(pac.SIO);

    // Set the pins to their default state
    let pins = Pins::new(
        pac.IO_BANK0,
        pac.PADS_BANK0,
        sio.gpio_bank0,
        &mut pac.RESETS,
    );

    // Setting up array with pins connected to button rows
    let button_matrix_row_pins: &mut [&mut dyn InputPin<Error = Infallible>; KEY_ROWS] = &mut [
        &mut pins.gpio29.into_pull_up_input(),
        &mut pins.gpio26.into_pull_up_input(),
        &mut pins.gpio15.into_pull_up_input(),
        &mut pins.gpio10.into_pull_up_input(),
        &mut pins.gpio9.into_pull_up_input(),
    ];

    // Setting up array with pins connected to button columns
    let button_matrix_col_pins: &mut [&mut dyn OutputPin<Error = Infallible>; KEY_COLS] = &mut [
        &mut pins.gpio11.into_push_pull_output(),
        &mut pins.gpio12.into_push_pull_output(),
        &mut pins.gpio13.into_push_pull_output(),
        &mut pins.gpio14.into_push_pull_output(),
        &mut pins.gpio27.into_push_pull_output(),
        &mut pins.gpio28.into_push_pull_output(),
        &mut pins.gpio2.into_push_pull_output(),
        &mut pins.gpio3.into_push_pull_output(),
        &mut pins.gpio4.into_push_pull_output(),
        &mut pins.gpio5.into_push_pull_output(),
        &mut pins.gpio6.into_push_pull_output(),
        &mut pins.gpio7.into_push_pull_output(),
    ];

    // Create button matrix object that scans all the PCB buttons
    let mut button_matrix: ButtonMatrix<KEY_ROWS, KEY_COLS, NUMBER_OF_KEYS> =
        ButtonMatrix::new(button_matrix_row_pins, button_matrix_col_pins, 5);

    // Create status LED
    let (mut pio, sm0, _, _, _) = pac.PIO0.split(&mut pac.RESETS);
    let mut status_led = Ws2812StatusLed::new(
        pins.gpio16.into_function(),
        &mut pio,
        sm0,
        clocks.peripheral_clock.freq(),
    );

    // Create keyboard button array
    let mut buttons: [KeyboardButton; NUMBER_OF_KEYS] = [KeyboardButton::default(); NUMBER_OF_KEYS];

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

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

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

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

    let mut start_count_down = timer.count_down();
    start_count_down.start(5000.millis());

    // Create variables to track caps lock and fn mode
    let mut caps_lock_active: bool = false;
    let mut fn_mode: u8;
    let mut sticky_state: u8 = 0;
    let mut sticky_key: Keyboard = Keyboard::NoEventIndicated;
    let mut started: bool = false;

    // Initialize button matrix
    button_matrix.init_pins();

    // Scan matrix to get initial state
    for _ in 0..10 {
        button_matrix.scan_matrix(&mut delay);
    }

    // Check if esc key is pressed while power on. If yes then enter bootloader
    if button_matrix.buttons_pressed()[0] {
        status_led.update(StatusMode::Bootloader);
        let gpio_activity_pin_mask: u32 = 0;
        let disable_interface_mask: u32 = 0;
        rp2040_hal::rom_data::reset_to_usb_boot(gpio_activity_pin_mask, disable_interface_mask);
    }

    // Configure USB
    let usb_bus = UsbBusAllocator::new(rp2040_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, 0x0003))
        .strings(&[StringDescriptors::default()
            .manufacturer("CMtec")
            .product("CMDR keyboard 51")
            .serial_number("0002")])
        .unwrap()
        .build();

    loop {
        if status_led_count_down.wait().is_ok() {
            update_status_led(&mut status_led, &caps_lock_active, &sticky_state, &started);
        }

        if start_count_down.wait().is_ok() && !started {
            started = true;
        }

        if usb_hid_report_count_down.wait().is_ok() {
            let pressed_keys = button_matrix.buttons_pressed();

            fn_mode = get_fn_mode(pressed_keys);

            if !caps_lock_active && sticky_state != 2 {
                update_status_led(&mut status_led, &caps_lock_active, &sticky_state, &started);
            }

            for (index, key) in pressed_keys.iter().enumerate() {
                buttons[index].pressed = *key;
            }

            let keyboard_report =
                get_keyboard_report(&mut buttons, fn_mode, &mut sticky_state, &mut sticky_key);

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

        if usb_tick_count_down.wait().is_ok() {
            button_matrix.scan_matrix(&mut delay);

            match keyboard.tick() {
                Err(UsbHidError::WouldBlock) => {}
                Ok(_) => {}
                Err(e) => {
                    status_led.update(StatusMode::Error);
                    core::panic!("Failed to process keyboard tick: {:?}", e)
                }
            };
        }

        if usb_dev.poll(&mut [&mut keyboard]) {
            match keyboard.device().read_report() {
                Err(UsbError::WouldBlock) => {}
                Err(e) => {
                    status_led.update(StatusMode::Error);
                    core::panic!("Failed to read keyboard report: {:?}", e)
                }
                Ok(leds) => {
                    caps_lock_active = leds.caps_lock;
                }
            }
        }
    }
}

/// Update status LED colour based on function layer and capslock
///
/// Normal = Off (OFF)
/// STICKY lock = blue/falshing blue (ACTIVITY)
/// Capslock active = flashing red (WARNING)
/// Error = steady red (ERROR)
///
/// # Arguments
/// * `status_led` - Reference to status LED
/// * `caps_lock_active` - Is capslock active
fn update_status_led<P, SM, I>(
    status_led: &mut Ws2812StatusLed<P, SM, I>,
    caps_lock_active: &bool,
    sticky_state: &u8,
    started: &bool,
) where
    I: AnyPin<Function = P::PinFunction>,
    P: PIOExt,
    SM: StateMachineIndex,
{
    if *caps_lock_active {
        status_led.update(StatusMode::Warning);
    } else if *sticky_state == 1 {
        status_led.update(StatusMode::Activity);
    } else if *sticky_state == 2 {
        status_led.update(StatusMode::ActivityFlash);
    } else if !(*started) {
        status_led.update(StatusMode::Normal);
    } else {
        status_led.update(StatusMode::Off);
    }
}

/// Get current Fn mode (0, 1 or 2)
/// layout::FN_BUTTONS contains the keycodes for each Fn key
///
/// # Arguments
///
/// * `pressed_keys` - Array of pressed keys
fn get_fn_mode(pressed_keys: [bool; NUMBER_OF_KEYS]) -> u8 {
    // Check how many Fn keys are pressed
    let mut active_fn_keys = layout::FN_BUTTONS
        .iter()
        .filter(|button_id| pressed_keys[**button_id as usize])
        .count() as u8;

    // Limit Fn mode to 2
    if active_fn_keys > 2 {
        active_fn_keys = 2;
    }
    active_fn_keys
}

/// Generate keyboard report based on pressed keys and Fn mode (0, 1 or 2)
/// layout::MAP contains the keycodes for each key in each Fn mode
///
/// # Arguments
///
/// * `matrix_keys` - Array of pressed keys
/// * `fn_mode` - Current function layer
/// * `sticky_state` - Is STICKY lock active
/// * `sticky_key` - the key pressed after STICKY lock was activated
fn get_keyboard_report(
    matrix_keys: &mut [KeyboardButton; NUMBER_OF_KEYS],
    fn_mode: u8,
    sticky_state: &mut u8,
    sticky_key: &mut Keyboard,
) -> [Keyboard; NUMBER_OF_KEYS] {
    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() {
        // Check if STICKY button is pressed (SET STICKY)
        if key.pressed != key.previous_pressed
            && key.pressed
            && index as u8 == layout::STICKY_BUTTON[0]
            && fn_mode == layout::STICKY_BUTTON[1]
            && *sticky_state == 0
        {
            *sticky_state = 1;
        }
        // Check if STICKY button is pressed (CLEAR STICKY)
        else if key.pressed != key.previous_pressed
            && key.pressed
            && index as u8 == layout::STICKY_BUTTON[0]
            && fn_mode == layout::STICKY_BUTTON[1]
            && *sticky_state != 0
        {
            *sticky_state = 0;
            *sticky_key = Keyboard::NoEventIndicated;
        }

        // Set fn mode for the pressed button
        if key.pressed != key.previous_pressed && key.pressed {
            key.fn_mode = fn_mode;
        }
        key.previous_pressed = key.pressed;

        // Skip key if defined as NoEventIndicated
        if layout::MAP[key.fn_mode as usize][index] == Keyboard::NoEventIndicated {
            continue;
        }

        // If STICKY lock is active, hold index key pressed until STICKY lock key is pressed
        // again
        if *sticky_state == 1 && key.pressed {
            *sticky_key = layout::MAP[key.fn_mode as usize][index];
            *sticky_state = 2;
        }

        // Add defined HID key to the report
        if key.pressed {
            keyboard_report[index] = layout::MAP[key.fn_mode as usize][index];
        }
    }

    /// Index of STICKY key in keyboard report
    /// Index 36, 37, 38, 45, 46, 47 are not used by any other keys
    const STICKY_REPORT_INDEX: usize = 46;
    // Add sticky key to the report
    keyboard_report[STICKY_REPORT_INDEX] = *sticky_key;

    keyboard_report
}