#![no_std]
#![no_main]

mod matrix;
mod keymap;
mod display;
mod usb;

use defmt::{info, warn};
use embassy_executor::Spawner;
use embassy_futures::join::join;
use embassy_rp::{
    bind_interrupts,
    i2c::{Config as I2cConfig, I2c, InterruptHandler as I2cInterruptHandler},
    peripherals::USB,
    usb::{Driver, InterruptHandler as UsbInterruptHandler},
};
use embassy_time::{Duration, Timer};
use usbd_hid::descriptor::KeyboardReport;
use crate::keymap::{Action, LayerOperation, get_action};

use {defmt_rtt as _, panic_probe as _};

bind_interrupts!(struct Irqs {
    I2C0_IRQ => I2cInterruptHandler<embassy_rp::peripherals::I2C0>;
    USBCTRL_IRQ => UsbInterruptHandler<USB>;
});


#[embassy_executor::main]
async fn main(_spawner: Spawner) -> () {
    let p = embassy_rp::init(Default::default());

    info!("Firmware starting");

    // Setup USB HID
    let usb_driver = Driver::new(p.USB, Irqs);
    let mut usb_keyboard = usb::setup_usb(usb_driver);
    let usb_fut = usb_keyboard.usb.run();
    info!("USB configured");

    // Setup I2C display
    info!("Initializing display...");
    let i2c_config = I2cConfig::default();
    let bus = I2c::new_async(p.I2C0, p.PIN_17, p.PIN_16, Irqs, i2c_config);
    let mut display = display::init_display(bus);
    match display {
        Some(_) => info!("Display initialized"),
        None => warn!("Can't initialize display"),
    }

    // Setup matrix
    info!("Initializing matrix...");
    let mut matrix = matrix_4x12!(p,
        rows: [PIN_0, PIN_1, PIN_2, PIN_3],
        cols: [PIN_4, PIN_5, PIN_6, PIN_7, PIN_8, PIN_9, PIN_10, PIN_11, PIN_12, PIN_13, PIN_14, PIN_15]
    );
    info!("Matrix initialized");

    let mut writer = usb_keyboard.writer;
    let reader = usb_keyboard.reader;
    let request_handler = usb_keyboard.request_handler;
    info!("Starting main loop...");

    display::clear_display(&mut display);

    // Keyboard scanning task
    let in_fut = async {
        let mut previous_keycodes: [u8; 6] = [0; 6];
        let mut previous_modifier: u8 = 0;
        let mut active_layer: usize = 0;

        loop {
            let scan_result = matrix.scan().await;

            // first pass to check if the layer keys are active
            let mut layer_active = false;

            for (row_idx, row) in scan_result.iter().enumerate() {
                for (col_idx, &is_pressed) in row.iter().enumerate() {
                    if is_pressed {
                        // try to find if any of the pressed keys on the 0th layer is a layer modify key
                        let action = get_action(row_idx, col_idx, 0);
                        if let Action::Layer(LayerOperation::MO(layer)) = action {
                            layer_active = true;
                            active_layer = layer;
                            break;
                        }
                    }
                }
                // break the loop early if layer key is pressed
                if layer_active {
                    break;
                }
            }

            if !layer_active {
                active_layer = 0;
            }

            // Collect all pressed keys and their scancodes
            let mut pressed_keys: heapless::Vec<(usize, usize, keymap::Action), 48> = heapless::Vec::new();

            for (row_idx, row) in scan_result.iter().enumerate() {
                for (col_idx, &is_pressed) in row.iter().enumerate() {
                    if is_pressed {
                        let action = get_action(row_idx, col_idx, active_layer);
                        if !matches!(action, Action::None | Action::Layer(_)) {
                            let _ = pressed_keys.push((row_idx, col_idx, action));
                        }
                    }
                }
            }

            // Build modifier byte and keycodes array
            let mut modifier: u8 = 0;
            let mut keycodes = [0u8; 6];
            let mut keycode_index = 0;

            for (_, _, action) in pressed_keys.iter() {
                if let Action::Key(hid_key) = action {
                    if hid_key.is_modifier() {
                        // Add to modifier byte
                        modifier |= hid_key.modifier_bit();
                    } else if keycode_index < 6 {
                        // Add to keycodes array (max 6 regular keys)
                        keycodes[keycode_index] = hid_key.as_u8();
                        keycode_index += 1;
                    }
                }
            }

            // Update display
            display::update_display(&mut display, &pressed_keys);

            // Send HID report if keycodes or modifiers changed
            if keycodes != previous_keycodes || modifier != previous_modifier {
                if keycodes != [0; 6] || modifier != 0 {
                    info!("Keys: {:?}, Mods: 0x{:02X}", keycodes, modifier);
                } else {
                    info!("All keys released");
                }

                let report = KeyboardReport {
                    keycodes,
                    leds: 0,
                    modifier,
                    reserved: 0,
                };

                match writer.write_serialize(&report).await {
                    Ok(()) => {},
                    Err(e) => warn!("Failed to send report: {:?}", e),
                };

                previous_keycodes = keycodes;
                previous_modifier = modifier;
            }

            Timer::after(Duration::from_millis(10)).await;
        }
    };

    // USB reader task
    let out_fut = async {
        reader.run(false, request_handler).await;
    };

    join(usb_fut, join(in_fut, out_fut)).await;
}