//! Project: CMtec CMDR Joystick
//! 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_joystick::buttons::SpecialAction;
use cmdr_joystick::hardware::{self, timers};
use cmdr_joystick::status::StatusMode;
use cmdr_joystick::usb_joystick_device::JoystickConfig;
use cmdr_joystick::{bootloader, Board, BoardParts, JoystickState};
use embedded_hal_0_2::adc::OneShot;
use embedded_hal_0_2::timer::CountDown;
use fugit::ExtU32;
use panic_halt as _;
use usb_device::prelude::*;
use usbd_human_interface_device::prelude::{UsbHidClassBuilder, UsbHidError};

// Embed the boot2 image for the W25Q080 flash; required for RP2040 to boot from external flash.
#[link_section = ".boot2"]
#[no_mangle]
#[used]
pub static BOOT2_FIRMWARE: [u8; 256] = rp2040_boot2::BOOT_LOADER_W25Q080;

#[rp2040_hal::entry]
fn main() -> ! {
    // Firmware entry point initializes hardware before handing off to RTIC.
    let BoardParts {
        mut button_matrix,
        mut status_led,
        mut delay,
        timer,
        mut adc,
        mut axis_pins,
        mut left_extra_button,
        mut right_extra_button,
        mut eeprom,
        usb_bus,
    } = Board::new().into_parts();

    // Build the HID joystick class on the shared USB bus.
    let mut usb_hid_joystick = UsbHidClassBuilder::new()
        .add_device(JoystickConfig::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();

    let mut state = JoystickState::new();
    status_led.update(StatusMode::Error);

    button_matrix.prime(&mut delay, 10);
    let initial_pressed = button_matrix.buttons_pressed();
    if bootloader::startup_requested(&initial_pressed) {
        bootloader::enter(&mut status_led);
    }

    {
        // Load persisted calibration values from EEPROM if available.
        let mut read_fn = |addr: u32| eeprom.read_byte(addr).map_err(|_| ());
        state.load_calibration(&mut read_fn);
    }

    // Set up periodic timers for scanning, status updates, and USB activity.
    let mut scan_tick = timer.count_down();
    scan_tick.start(timers::SCAN_INTERVAL_US.micros());

    let mut status_tick = timer.count_down();
    status_tick.start(timers::STATUS_LED_INTERVAL_MS.millis());

    let mut usb_tick = timer.count_down();
    usb_tick.start(timers::USB_UPDATE_INTERVAL_MS.millis());

    let mut status_time_ms: u32 = 0;
    let mut suspended_scan_counter: u8 = 0;

    // Main control loop: service USB, process inputs, and emit reports.
    loop {
        // Service the USB stack and HID class when data is pending.
        if usb_dev.poll(&mut [&mut usb_hid_joystick]) {
            state.usb_state().on_poll();
        }

        let usb_state = usb_dev.state();
        state.usb_state().on_suspend_change(usb_state);

        // Periodically refresh the status LED animation.
        if status_tick.wait().is_ok() {
            status_time_ms = status_time_ms.saturating_add(timers::STATUS_LED_INTERVAL_MS);
            status_led.update_from_system_state(state.system_state(), status_time_ms);
        }

        // Slow the scan cadence when USB is suspended to save power.
        let should_scan = if state.usb_state().suspended {
            suspended_scan_counter = suspended_scan_counter.wrapping_add(1) % 10;
            suspended_scan_counter == 0
        } else {
            suspended_scan_counter = 0;
            true
        };

        if should_scan && scan_tick.wait().is_ok() {
            // Scan buttons, read analog axes, and update state machines.
            button_matrix.scan_matrix(&mut delay);

            let mut raw_values = [
                adc.read(&mut axis_pins.left_x).unwrap(),
                adc.read(&mut axis_pins.left_y).unwrap(),
                adc.read(&mut axis_pins.right_x).unwrap(),
                adc.read(&mut axis_pins.right_y).unwrap(),
            ];
            state.tick_smoothers(&mut raw_values);

            state.update_button_states(
                &mut button_matrix,
                &mut left_extra_button,
                &mut right_extra_button,
            );

            // Evaluate special button combinations (bootloader, calibration, etc.).
            let action = state.check_special_action();
            if matches!(action, SpecialAction::Bootloader) {
                if !state.usb_state().suspended {
                    let clear_report = JoystickState::empty_report();
                    for _ in 0..3 {
                        match usb_hid_joystick.device().write_report(&clear_report) {
                            Ok(_) => break,
                            Err(UsbHidError::WouldBlock) => {
                                let _ = usb_hid_joystick.tick();
                            }
                            Err(_) => break,
                        }
                    }
                }
                bootloader::enter(&mut status_led);
            } else if !matches!(action, SpecialAction::None) {
                let mut write_page =
                    |page: u32, data: &[u8]| eeprom.write_page(page, data).map_err(|_| ());
                state.handle_special_action(action, &mut write_page);
            }

            // Track calibration extrema and process axis/virtual/button logic.
            state.update_calibration_tracking();

            if state.process_axes() {
                state.usb_state().handle_input_activity();
            }

            if state.update_virtual_axes() {
                state.usb_state().handle_input_activity();
            }

            if state.finalize_button_logic(&timer) {
                state.usb_state().handle_input_activity();
            }
        }

        // Advance USB idle timers and decide when to send reports.
        let usb_tick_elapsed = usb_tick.wait().is_ok();
        if usb_tick_elapsed {
            state
                .usb_state()
                .advance_idle_timer(timers::USB_UPDATE_INTERVAL_MS);
        }

        // Emit a new HID report when activity is pending and USB is ready.
        if state.usb_state().activity
            && (usb_tick_elapsed || state.usb_state().send_pending)
            && !state.usb_state().suspended
        {
            let report = state.build_report();
            match usb_hid_joystick.device().write_report(&report) {
                Err(UsbHidError::WouldBlock) => {}
                Ok(_) => {
                    state.usb_state().acknowledge_report();
                }
                Err(error) => {
                    status_led.update(StatusMode::Error);
                    panic!("Failed to write joystick report: {:?}", error);
                }
            }
        } else if usb_tick_elapsed && state.usb_state().active && !state.usb_state().suspended {
            state.usb_state().idle_mode = true;
        }
    }
}
