#![no_std]
#![no_main]

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 MatrixKey {
    current_state: bool,
    previous_state: bool,
    fn_mode: u8,
}

impl MatrixKey {
    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 matrix_keys: [MatrixKey; NUMBER_OF_KEYS] = [MatrixKey::default(); NUMBER_OF_KEYS];

    let matrix_rows: &[&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 matrix_cols: &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(),
    ];

    // -------------------------------------      -------------------------------------
    // |  0  |  1  |  2  |  3  |  4  |  5  |      |  6  |  7  |  8  |  9  | 10  | 11  |
    // | 12  | 13  | 14  | 15  | 16  | 17  |      | 18  | 19  | 20  | 21  | 22  | 23  |
    // | 24  | 25  | 26  | 27  | 28  | 29  |      | 30  | 31  | 32  | 33  | 34  | 35  |
    // ------------------| 36  | 37  | 38  |      | 39  | 40  | 41  |------------------
    //                   -------------------      -------------------
    // GRAVE = §
    // SEMICOLON = ö
    // APOSTROPHE = ä
    // LEFT_BRACE = å
    // FORWARDSLASH = -
    // NON_US_BACKSLASH = <
    // EQUAL = ´
    // BACKSLASH = '
    // RIGHT_BRACE = ^
    // MINUS = +
    // LEFT_ALT = Alt
    // RIGHT_ALT = AltGr

    let layout: [[Keyboard; NUMBER_OF_KEYS]; 3] = [
        [
            // Function layer 0
            Keyboard::Tab,
            Keyboard::Q,
            Keyboard::W,
            Keyboard::E,
            Keyboard::R,
            Keyboard::T,
            Keyboard::Y,
            Keyboard::U,
            Keyboard::I,
            Keyboard::O,
            Keyboard::P,
            Keyboard::LeftBrace,
            Keyboard::LeftControl,
            Keyboard::A,
            Keyboard::S,
            Keyboard::D,
            Keyboard::F,
            Keyboard::G,
            Keyboard::H,
            Keyboard::J,
            Keyboard::K,
            Keyboard::L,
            Keyboard::Semicolon,
            Keyboard::Apostrophe,
            Keyboard::LeftShift,
            Keyboard::Z,
            Keyboard::X,
            Keyboard::C,
            Keyboard::V,
            Keyboard::B,
            Keyboard::N,
            Keyboard::M,
            Keyboard::Comma,
            Keyboard::Dot,
            Keyboard::ForwardSlash,
            Keyboard::RightShift,
            Keyboard::LeftAlt,
            Keyboard::NoEventIndicated,
            Keyboard::Space,
            Keyboard::Space,
            Keyboard::NoEventIndicated,
            Keyboard::RightAlt,
        ],
        [
            // Function layer 1
            Keyboard::Escape,
            Keyboard::F1,
            Keyboard::F2,
            Keyboard::F3,
            Keyboard::F4,
            Keyboard::F5,
            Keyboard::F6,
            Keyboard::F7,
            Keyboard::F8,
            Keyboard::F9,
            Keyboard::F10,
            Keyboard::DeleteBackspace,
            Keyboard::LeftControl,
            Keyboard::Keyboard1,
            Keyboard::Keyboard2,
            Keyboard::Keyboard3,
            Keyboard::Keyboard4,
            Keyboard::Keyboard5,
            Keyboard::Keyboard6,
            Keyboard::Keyboard7,
            Keyboard::Keyboard8,
            Keyboard::Keyboard9,
            Keyboard::Keyboard0,
            Keyboard::ReturnEnter,
            Keyboard::LeftShift,
            Keyboard::Keyboard6,
            Keyboard::Keyboard7,
            Keyboard::Keyboard8,
            Keyboard::Keyboard9,
            Keyboard::Keyboard0,
            Keyboard::NonUSBackslash,
            Keyboard::Equal,
            Keyboard::Backslash,
            Keyboard::RightBrace,
            Keyboard::Minus,
            Keyboard::RightShift,
            Keyboard::LeftAlt,
            Keyboard::NoEventIndicated,
            Keyboard::DeleteBackspace,
            Keyboard::DeleteBackspace,
            Keyboard::NoEventIndicated,
            Keyboard::RightAlt,
        ],
        [
            // Function layer 2
            Keyboard::F11,
            Keyboard::F12,
            Keyboard::F13,
            Keyboard::F14,
            Keyboard::F15,
            Keyboard::F16,
            Keyboard::Grave,
            Keyboard::NoEventIndicated,
            Keyboard::LeftGUI,
            Keyboard::NoEventIndicated,
            Keyboard::CapsLock,
            Keyboard::DeleteBackspace,
            Keyboard::LeftControl,
            Keyboard::NoEventIndicated,
            Keyboard::NoEventIndicated,
            Keyboard::F17,
            Keyboard::F18,
            Keyboard::F19,
            Keyboard::LeftArrow,
            Keyboard::DownArrow,
            Keyboard::UpArrow,
            Keyboard::RightArrow,
            Keyboard::DeleteForward,
            Keyboard::ReturnEnter,
            Keyboard::LeftShift,
            Keyboard::F20,
            Keyboard::F21,
            Keyboard::F22,
            Keyboard::F23,
            Keyboard::F24,
            Keyboard::Home,
            Keyboard::PageDown,
            Keyboard::PageUp,
            Keyboard::End,
            Keyboard::Insert,
            Keyboard::RightShift,
            Keyboard::LeftAlt,
            Keyboard::NoEventIndicated,
            Keyboard::LeftGUI,
            Keyboard::DeleteBackspace,
            Keyboard::NoEventIndicated,
            Keyboard::RightAlt,
        ],
    ];

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

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

    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 = false;

    init_keyboard_matrix_pins(matrix_cols, &mut delay);

    // Infinite colour wheel loop
    loop {
        if input_count_down.wait().is_ok() {
            // Scan keyboard matrix
            let pressed_keys = get_pressed_keys(matrix_rows, matrix_cols, &mut delay);

            // Get current function layer
            let 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 = red.
            if caps_lock_active == true {
                status_led
                    .write([status_led_color[4]].iter().copied())
                    .unwrap();
            } else {
                status_led
                    .write([status_led_color[usize::from(fn_mode) + 1]].iter().copied())
                    .unwrap();
            }

            // Copy result from scanned keys to matrix struct
            for (index, key) in pressed_keys.iter().enumerate() {
                matrix_keys[index].current_state = *key;
            }

            // Generate keyboard report
            let keyboard_report = get_keyboard_report(matrix_keys, layout, 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_keyboard_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_keys(
    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;
    if pressed_keys[37] == true && pressed_keys[40] == true {
        fn_mode = 2;
    } else if pressed_keys[37] == true || pressed_keys[40] == true {
        fn_mode = 1;
    }
    fn_mode
}

// Generate keyboard report based on pressed keys and Fn mode (0, 1 or 2)
fn get_keyboard_report(
    mut matrix_keys: [MatrixKey; NUMBER_OF_KEYS],
    layout: [[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.previous_state = key.current_state;
            key.fn_mode = fn_mode;
        }
        if key.current_state == true {
            keyboard_report[index] = layout[usize::from(key.fn_mode)][index];
        }
    }
    // Return report
    keyboard_report
}