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

#![no_std]
#![no_main]

use cmdr_keyboard_42::hardware::{self, timers};
use cmdr_keyboard_42::{ButtonMatrix, KeyboardState, StatusLed, StatusMode};
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::Pins,
    pac,
    pio::PIOExt,
    timer::Timer,
    watchdog::Watchdog,
};
use usb_device::class_prelude::*;
use usb_device::prelude::*;
use usb_device::device::UsbDeviceState;
use usbd_human_interface_device::prelude::*;

#[unsafe(link_section = ".boot2")]
#[unsafe(no_mangle)]
#[used]
pub static BOOT2_FIRMWARE: [u8; 256] = rp2040_boot2::BOOT_LOADER_W25Q080;

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

    let mut watchdog = Watchdog::new(pac.WATCHDOG);

    let clocks = init_clocks_and_plls(
        hardware::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();
    let sio = Sio::new(pac.SIO);

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

    let button_matrix_row_pins: &mut [&mut dyn InputPin<Error = core::convert::Infallible>;
             hardware::KEY_ROWS] = &mut [
        &mut cmdr_keyboard_42::get_pin!(pins, button_row_0).into_pull_up_input(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_row_1).into_pull_up_input(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_row_2).into_pull_up_input(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_row_3).into_pull_up_input(),
    ];

    let button_matrix_col_pins: &mut [&mut dyn OutputPin<Error = core::convert::Infallible>;
             hardware::KEY_COLS] = &mut [
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_0).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_1).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_2).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_3).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_4).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_5).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_6).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_7).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_8).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_9).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_10).into_push_pull_output(),
        &mut cmdr_keyboard_42::get_pin!(pins, button_col_11).into_push_pull_output(),
    ];

    let mut button_matrix: ButtonMatrix<
        { hardware::KEY_ROWS },
        { hardware::KEY_COLS },
        { hardware::NUMBER_OF_KEYS },
    > = ButtonMatrix::new(
        button_matrix_row_pins,
        button_matrix_col_pins,
        hardware::MATRIX_DEBOUNCE_SCANS_PRESS,
        hardware::MATRIX_DEBOUNCE_SCANS_RELEASE,
    );

    let mut keyboard_state = KeyboardState::new();

    let (mut pio, sm0, _, _, _) = pac.PIO0.split(&mut pac.RESETS);
    let mut status_led = StatusLed::new(
        cmdr_keyboard_42::get_pin!(pins, status_led).into_function(),
        &mut pio,
        sm0,
        clocks.peripheral_clock.freq(),
    );
    status_led.update(StatusMode::Error);

    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_tick_count_down = timer.count_down();
    usb_tick_count_down.start(timers::USB_TICK_INTERVAL_US.micros());

    let mut status_led_count_down = timer.count_down();
    status_led_count_down.start(timers::STATUS_LED_INTERVAL_MS.millis());
    let mut status_time_ms: u32 = 0;
    let mut usb_initialized = false;
    let mut usb_suspended = false;
    let mut wake_on_input = false;
    let mut last_activity_ms: u32 = 0;

    button_matrix.init_pins();

    for _ in 0..hardware::INITIAL_SCAN_PASSES {
        button_matrix.scan_matrix(&mut delay);
    }

    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);
    }

    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(hardware::usb::VID, hardware::usb::PID))
            .strings(&[StringDescriptors::default()
                .manufacturer(hardware::usb::MANUFACTURER)
                .product(hardware::usb::PRODUCT)
                .serial_number(hardware::usb::SERIAL_NUMBER)])
            .unwrap()
            .build();

    loop {
        if status_led_count_down.wait().is_ok() {
            status_time_ms = status_time_ms.saturating_add(timers::STATUS_LED_INTERVAL_MS);
            let idle_elapsed = status_time_ms.saturating_sub(last_activity_ms);
            let idle_mode = usb_initialized && idle_elapsed >= timers::IDLE_TIMEOUT_MS;
            let usb_active = usb_initialized && !idle_mode;
            status_led.apply_summary(
                keyboard_state.status_summary(usb_initialized, usb_active, usb_suspended, idle_mode),
                status_time_ms,
            );
        }

        // Skip high-frequency scanning when suspended to save power
        // Only scan periodically to detect wake-up inputs
        let should_scan = if usb_suspended {
            // When suspended, reduce scan frequency by factor of 20 (every ~5ms instead of 250μs)
            static mut SUSPENDED_SCAN_COUNTER: u8 = 0;
            unsafe {
                SUSPENDED_SCAN_COUNTER = (SUSPENDED_SCAN_COUNTER + 1) % 20;
                SUSPENDED_SCAN_COUNTER == 0
            }
        } else {
            true
        };

        if usb_tick_count_down.wait().is_ok() && should_scan {
            button_matrix.scan_matrix(&mut delay);
            let pressed_keys = button_matrix.buttons_pressed();

            // Check for input activity
            if pressed_keys.iter().any(|pressed| *pressed) {
                last_activity_ms = status_time_ms;

                // Wake from USB suspend if input detected
                if wake_on_input && usb_suspended {
                    // TODO: Implement remote wakeup if supported by host
                    wake_on_input = false;
                }
            }

            let keyboard_report = keyboard_state.process_scan(pressed_keys);

            // Only transmit USB reports when not suspended
            if !usb_suspended {
                match keyboard.device().write_report(keyboard_report) {
                    Err(UsbHidError::WouldBlock) | Err(UsbHidError::Duplicate) => {}
                    Ok(_) => {
                        usb_initialized = true;
                    }
                    Err(_) => {
                        keyboard_state.mark_stopped();
                        usb_initialized = false;
                    }
                };
            }

            match keyboard.tick() {
                Err(UsbHidError::WouldBlock) | Ok(_) => {}
                Err(_) => {
                    keyboard_state.mark_stopped();
                    usb_initialized = false;
                }
            };
        }

        if usb_dev.poll(&mut [&mut keyboard]) {
            match keyboard.device().read_report() {
                Err(UsbError::WouldBlock) => {}
                Err(_) => {
                    keyboard_state.mark_stopped();
                    usb_initialized = false;
                }
                Ok(leds) => {
                    keyboard_state.update_caps_lock(leds.caps_lock);
                    keyboard_state.mark_started();
                    usb_initialized = true;
                    last_activity_ms = status_time_ms;
                }
            }
        }

        // Check USB device state for suspend/resume handling
        let usb_state = usb_dev.state();
        let was_suspended = usb_suspended;
        usb_suspended = usb_state == UsbDeviceState::Suspend;

        // Handle USB resume transition
        if was_suspended && !usb_suspended {
            // Device was suspended and is now resumed
            last_activity_ms = status_time_ms;
            wake_on_input = false;
        }

        // Handle USB suspend transition
        if !was_suspended && usb_suspended {
            // Device has just been suspended - enter power saving mode
            wake_on_input = true;
        }
    }
}