embassy-stm32 = { git = "https://github.com/embassy-rs/embassy.git", features = ["time-driver-any", "stm32f303cb", "unstable-pac", "memory-x"] }

embassy-stm32 = { git = "https://github.com/embassy-rs/embassy.git", features = ["time-driver-any", "stm32f303cb", "unstable-pac", "memory-x", "nightly", "unstable-traits"] }
embassy-embedded-hal = { git = "https://github.com/embassy-rs/embassy.git" }
embedded-hal-async = "0.1.0-alpha.1"

group_imports = "StdExternalCrate"
imports_granularity = "Crate"

use embassy_stm32::Peripherals;
pub trait Keyboard {
    fn init(p: Peripherals) -> Self;
    fn run(self) -> !;
}
pub mod moonlander;

#![no_main]
#![no_std]
#![feature(type_alias_impl_trait)]
use ashtoret as _; // global logger + panicking-behavior + memory layout
use defmt::*;
use embassy::executor::Spawner;
use embassy::time::{Duration, Timer};
use embassy_stm32::{
    gpio::{Level, Output, Speed},
    Peripherals,
};
use ashtoret::keyboards::{moonlander::Moonlander, Keyboard};
#[embassy::main]
async fn main(_spawner: Spawner, p: Peripherals) {
    let keyboard = Moonlander::init(p);
    keyboard.run()
}

use alloc::sync::Arc;
use embassy::{blocking_mutex::raw::ThreadModeRawMutex, mutex::Mutex};
use embassy_stm32::gpio::{AnyPin, Output};
pub struct Leds {
    pub left: [Output<'static, AnyPin>; 3],
    pub right: Arc<Mutex<ThreadModeRawMutex, [bool; 3]>>,
}
impl Leds {
    pub async fn set(&mut self, idx: usize, val: bool) {
        match idx {
            0..=2 => {
                let pin = &mut self.left[idx];
                if val {
                    pin.set_high();
                } else {
                    pin.set_low();
                }
            }
            3..=5 => {
                self.right.lock().await[idx - 3] = val;
            }
            _ => {}
        }
    }
    pub async fn set_all(&mut self, val: u8) {
        let mut lock = self.right.lock().await;
        for (idx, pin) in self.left.iter_mut().enumerate() {
            if val & (0x01 << idx) == (0x01 << idx) {
                pin.set_high();
            } else {
                pin.set_low();
            }
            let sh = idx + 3;
            lock[idx] = val & (0x01 << sh) == (0x01 << sh);
        }
    }
    pub async fn clear(&mut self) {
        self.set_all(0).await;
    }
}

#![no_main]
#![no_std]
#![feature(type_alias_impl_trait)]
#![feature(generic_associated_types)]
mod leds;
mod matrix;
extern crate alloc;
use alloc::{sync::Arc, vec::Vec};
use ashtoret as _; // global logger + panicking-behavior + memory layout
use ashtoret::{
    drivers::{matrix::Matrix, mcp23018::Mcp23018},
    init_alloc, inputs, outputs,
};
use embassy::{
    blocking_mutex::raw::ThreadModeRawMutex,
    executor::Spawner,
    mutex::Mutex,
    time::{Duration, Timer},
    util::Forever,
};
use embassy_embedded_hal::shared_bus::i2c::I2cBusDevice;
use embassy_stm32::{
    gpio::{AnyPin, Output},
    i2c::I2c,
    interrupt,
    peripherals::{DMA1_CH6, DMA1_CH7, I2C1},
    time::U32Ext,
    Peripherals,
};
use crate::{leds::Leds, matrix::MoonlanderMatrix};
#[embassy::main]
async fn main(_spawner: Spawner, p: Peripherals) {
    init_alloc();
    let mut leds = Leds {
        left: outputs![p.PB5, p.PB4, p.PB3],
        right: Arc::new(Mutex::new([false; 3])),
    };
    static I2C_BUS: Forever<Mutex<ThreadModeRawMutex, I2c<I2C1, DMA1_CH6, DMA1_CH7>>> =
        Forever::new();
    let i2c_bus = I2c::new(
        p.I2C1,
        p.PB6,
        p.PB7,
        interrupt::take!(I2C1_EV),
        p.DMA1_CH6,
        p.DMA1_CH7,
        100000.hz(),
    );
    let i2c_bus = I2C_BUS.put(Mutex::new(i2c_bus));
    let i2c_dev1 = I2cBusDevice::new(i2c_bus);
    let mcp23018 = Mcp23018::new(i2c_dev1);
    let mut matrix = Matrix::<7, 14, _>::new(MoonlanderMatrix {
        rows: outputs![p.PB10, p.PB11, p.PB12, p.PB13, p.PB14, p.PB15],
        cols: inputs![p.PA0, p.PA1, p.PA2, p.PA3, p.PA6, p.PA7, p.PB0],
        ext: mcp23018,
        ext_init: false,
        leds: Arc::clone(&leds.right),
    })
    .with_debounce(Duration::from_millis(5));
    // placeholder, binary counting on both halves
    let mut counter = 0;
    loop {
        leds.set_all(counter).await;
        counter += 1;
        Timer::after(Duration::from_millis(250)).await;
        matrix.update().await;
    }
}

use alloc::sync::Arc;
use core::future::Future;
use ashtoret::drivers::{
    matrix::{MatrixArray, MatrixScanner},
    mcp23018::{Mcp23018, Port},
};
use defmt::error;
use embassy::{
    blocking_mutex::raw::ThreadModeRawMutex,
    mutex::Mutex,
    time::{block_for, Duration},
};
use embassy_stm32::gpio::{AnyPin, Input, Output};
use embedded_hal_async::i2c::I2c;
pub struct MoonlanderMatrix<I2C> {
    pub rows: [Output<'static, AnyPin>; 6],
    pub cols: [Input<'static, AnyPin>; 7],
    pub ext: Mcp23018<I2C>,
    pub ext_init: bool,
    pub leds: Arc<Mutex<ThreadModeRawMutex, [bool; 3]>>,
}
impl<I2C: 'static, E> MoonlanderMatrix<I2C>
where
    I2C: I2c<Error = E>,
{
    async fn init_ext(&mut self) -> Result<(), E> {
        self.ext
            .config_port(Port::A, 0b00000000, 0b10000000)
            .await?;
        self.ext.config_port(Port::B, 0b00111111, 0b11111111).await
    }
    async fn scan_ext(&mut self, matrix: &mut MatrixArray<7, 14>) -> Result<bool, E> {
        let mut has_changed = false;
        let led_lock = self.leds.lock().await;
        let led_mask = [
            (!led_lock[2] as u8) << 7,
            (!led_lock[1] as u8) << 6 | (!led_lock[0] as u8) << 7,
        ];
        for (y, row) in (0..6).zip(matrix.iter_mut()) {
            self.ext
                .set_all([0x01 << y | led_mask[0], led_mask[1]])
                .await?;
            let data = self.ext.read_port(Port::B).await?;
            for (x, col) in row.iter_mut().enumerate() {
                let val = data & (0x01 << x) == 0x01 << x;
                if *col != val {
                    has_changed = true;
                    *col = val;
                }
            }
        }
        Ok(has_changed)
    }
}
impl<I2C: 'static, E> MatrixScanner<7, 14> for MoonlanderMatrix<I2C>
where
    I2C: I2c<Error = E>,
{
    type ScanFuture<'a> = impl Future<Output = bool>;
    fn scan<'a>(&'a mut self, matrix: &'a mut MatrixArray<7, 14>) -> Self::ScanFuture<'a> {
        async move {
            let mut has_changed = false;
            for (y, row_pin) in self.rows.iter_mut().enumerate() {
                row_pin.set_high();
                // Note: QMK orders the pin state changes such that it can skip this delay if it already had to wait for I2C
                // TODO: Replicate this while keeping the code readable
                block_for(Duration::from_micros(10));
                for (x, col) in self.cols.iter().enumerate() {
                    let val = col.is_high();
                    if matrix[x][y] != val {
                        has_changed = true;
                        matrix[y][x] = val;
                    }
                }
                row_pin.set_low();
            }
            if !self.ext_init && self.init_ext().await.is_ok() {
                self.ext_init = true;
            }
            if self.ext_init {
                let res = self.scan_ext(matrix).await;
                match res {
                    Ok(e) => has_changed = e,
                    Err(_) => {
                        self.ext_init = false;
                        error!("Failed to scan right half of Moonlander despite it being initialized. This error is expected if you just unplugged it.");
                    }
                }
            }
            has_changed
        }
    }
}

use core::future::Future;

type MatrixArray<const COLS: usize, const ROWS: usize> = [[bool; COLS]; ROWS];

pub type MatrixArray<const CS: usize, const RS: usize> = [[bool; CS]; RS];

pub enum ScanKind<const COLS: usize, const ROWS: usize, T> {
    Direct([[Input<'static, AnyPin>; COLS]; ROWS]),
    Row2Col {
        row_pins: [Output<'static, AnyPin>; ROWS],
        col_pins: [Input<'static, AnyPin>; COLS],
    },
    Custom(T),

// Note: Common scanning types like simple row2col/col2row should be defined here
pub trait MatrixScanner<const CS: usize, const RS: usize> {
    type ScanFuture<'a>: Future<Output = bool>
    where
        Self: 'a;
    fn scan<'a>(&'a mut self, matrix: &'a mut MatrixArray<CS, RS>) -> Self::ScanFuture<'a>;

pub struct Matrix<const COLS: usize, const ROWS: usize, T> {
    state: MatrixArray<COLS, ROWS>,
    prev_state: MatrixArray<COLS, ROWS>,
    pins: ScanKind<COLS, ROWS, T>,

pub struct Matrix<const CS: usize, const RS: usize, SC> {
    state: MatrixArray<CS, RS>,
    prev_state: MatrixArray<CS, RS>,
    scanner: SC,

    scan_fn: Option<fn(&ScanKind<COLS, ROWS, T>) -> MatrixArray<COLS, ROWS>>,

impl<const COLS: usize, const ROWS: usize, T> Matrix<{ COLS }, { ROWS }, T> {
    pub fn new(pins: ScanKind<COLS, ROWS, T>) -> Self {

impl<const CS: usize, const RS: usize, SC> Matrix<{ CS }, { RS }, SC>
where
    SC: MatrixScanner<CS, RS>,
{
    pub fn new(scanner: SC) -> Self {

            state: [[false; COLS]; ROWS],
            prev_state: [[false; COLS]; ROWS],
            pins,

            state: [[false; CS]; RS],
            prev_state: [[false; CS]; RS],
            scanner,

            scan_fn: None,

        self
    }
    pub fn with_custom_scan(
        mut self,
        scan_fn: fn(&ScanKind<COLS, ROWS, T>) -> MatrixArray<COLS, ROWS>,
    ) -> Self {
        self.scan_fn = Some(scan_fn);

    fn raw_update(&mut self) -> tinyvec::TinyVec<[KeyEvent; 6]> {

    async fn raw_update(&mut self) -> Option<tinyvec::TinyVec<[KeyEvent; 6]>> {

        self.state = if let Some(custom_scan) = self.scan_fn {
            custom_scan(&self.pins)
        } else {
            // Note: Not needed for the Moonlander, as it implements a custom scanning routine
            todo!()
        };
        let mut vec = tinyvec::TinyVec::new();

        if self.scanner.scan(&mut self.state).await {
            let mut vec = tinyvec::TinyVec::new();

        for (idx, (key, prev_key)) in self
            .state
            .iter()
            .flatten()
            .zip(self.prev_state.iter().flatten())
            .enumerate()
        {
            if let Some(event) = if *key && !prev_key {
                Some(KeyEvent::Pressed(idx as u16))
            } else if !key && *prev_key {
                Some(KeyEvent::Released(idx as u16))
            } else {
                None
            } {
                vec.push(event)

            for (idx, (key, prev_key)) in self
                .state
                .iter()
                .flatten()
                .zip(self.prev_state.iter().flatten())
                .enumerate()
            {
                if let Some(event) = if *key && !prev_key {
                    Some(KeyEvent::Pressed(idx as u16))
                } else if !key && *prev_key {
                    Some(KeyEvent::Released(idx as u16))
                } else {
                    None
                } {
                    vec.push(event)
                }

        }

        vec

            Some(vec)
        } else {
            None
        }

        let changes = self.raw_update();
        if !changes.is_empty() {
            if let Some(debounce) = self.debounce {
                Timer::after(debounce).await;
                if self.raw_update().is_empty() {
                    Some(changes)
                } else {
                    None
                }
            } else {

        let changes = self.raw_update().await?;
        if let Some(debounce) = self.debounce {
            Timer::after(debounce).await;
            if self.raw_update().await.is_none() {

            } else {
                None

            None

            Some(changes)

use crate::drivers::SharedBus;
use embassy_stm32::i2c::{self, I2c};

// https://ww1.microchip.com/downloads/en/devicedoc/22103a.pdf
// Note: Default config is BANK = 0; SEQOP = 0

pub struct Mcp23018<T: i2c::Instance> {
    i2c: SharedBus<I2c<'static, T>>,

use embedded_hal_async::i2c::I2c;
pub const MCP23018_DEFAULT_ADDR: u8 = 0b0100000;
#[allow(dead_code)]
#[allow(non_camel_case_types)]
#[derive(Copy, Clone)]
#[repr(u8)]
pub enum Port {
    A = 0x0,
    B = 0x1,
}
#[allow(dead_code)]
#[allow(non_camel_case_types)]
#[derive(Copy, Clone)]
#[repr(u8)]
pub enum Register {
    IODIR = 0x00,
    IPOL = 0x02,
    GPINTEN = 0x04,
    DEFVAL = 0x06,
    INTCON = 0x08,
    IOCON = 0x0A,
    GPPU = 0x0C,
    INTF = 0x0E,
    INTCAP = 0x10,
    GPIO = 0x12,
    OLAT = 0x14,
}
pub struct Mcp23018<I2C> {
    i2c: I2C,
    addr: u8,

impl<T: i2c::Instance> Mcp23018<T> {
    pub fn new() -> Self {
        todo!()

#[allow(dead_code)]
impl<E, I2C: I2c<Error = E>> Mcp23018<I2C> {
    pub fn new(i2c: I2C) -> Self {
        Self {
            i2c,
            addr: MCP23018_DEFAULT_ADDR,
        }
    }
    pub fn with_addr(mut self, addr: u8) -> Self {
        self.addr = addr & !0x01;
        self
    }
    // 0 -> Output, no pullup
    // 1 -> Input, pullup
    pub async fn config_port(&mut self, port: Port, io_dir: u8, pullup: u8) -> Result<(), E> {
        self.write_reg(Register::IODIR, port, io_dir).await?;
        self.write_reg(Register::GPPU, port, pullup).await
    }
    pub async fn set_port(&mut self, port: Port, data: u8) -> Result<(), E> {
        self.write_reg(Register::GPIO, port, data).await
    }
    pub async fn set_all(&mut self, data: [u8; 2]) -> Result<(), E> {
        self.write_reg_seq(Register::GPIO, data).await
    }
    pub async fn read_port(&mut self, port: Port) -> Result<u8, E> {
        self.read_reg(Register::GPIO, port).await
    }
    async fn write_reg(&mut self, reg: Register, port: Port, data: u8) -> Result<(), E> {
        self.i2c
            .write(self.addr, &[reg as u8 | port as u8, data])
            .await
    }
    async fn write_reg_seq(&mut self, reg: Register, data: [u8; 2]) -> Result<(), E> {
        self.i2c
            .write(self.addr, &[reg as u8, data[0], data[1]])
            .await
    }
    async fn read_reg(&mut self, reg: Register, port: Port) -> Result<u8, E> {
        let mut buf = [0];
        self.i2c
            .write_read(self.addr, &[reg as u8], &mut buf)
            .await?;
        Ok(buf[0])

use embassy::blocking_mutex::NoopMutex;

pub type SharedBus<T> = &'static NoopMutex<T>;

// Equivalent to QMK's keycodes
#[derive(Clone, PartialEq, Debug, defmt::Format)]
pub enum Keycode<T: Clone> {
    Noop,
    Trans,
    A,
    B,
    C,
    D,
    E,
    F,
    G,
    H,
    I,
    J,
    K,
    L,
    M,
    N,
    O,
    P,
    Q,
    R,
    S,
    T,
    U,
    V,
    W,
    X,
    Y,
    Z,
    Enter,
    Escape,
    Backspace,
    Tab,
    Space,
    Minus,
    Equal,
    LeftBracket,
    RightBracket,
    Backslash,
    NonusHash,
    Custom(T),
}

use crate::keycodes::Keycode;
use alloc::vec::Vec;
type Layer<const CS: usize, const RS: usize, T> = [[Keycode<T>; CS]; RS];
pub struct Layout<const CS: usize, const RS: usize, T: Clone> {
    layers: Vec<Layer<CS, RS, T>>,
}
impl<const CS: usize, const RS: usize, T: Clone> Layout<CS, RS, T> {
    pub fn new() -> Layout<CS, RS, T> {
        Self { layers: Vec::new() }
    }
    pub fn push_layers(&mut self, layers: impl IntoIterator<Item = Layer<CS, RS, T>>) {
        self.layers.extend(layers.into_iter());
    }
    pub fn get_kc(&self, layer: usize, row: usize, col: usize) -> Option<Keycode<T>> {
        self.layers.get(layer)?.get(row)?.get(col).cloned()
    }
}

#![feature(generic_associated_types)]
#![feature(associated_type_defaults)]

extern crate alloc;

mod drivers;
pub mod keyboards;

pub mod drivers;
pub mod keycodes;
pub mod layout;

use alloc_cortex_m::CortexMHeap;

use alloc_cortex_m::CortexMHeap;

pub fn init_alloc() {
    use core::mem::MaybeUninit;
    const HEAP_SIZE: usize = 1024;
    static mut HEAP: [MaybeUninit<u8>; HEAP_SIZE] = [MaybeUninit::uninit(); HEAP_SIZE];
    unsafe { ALLOCATOR.init(HEAP.as_ptr() as usize, HEAP_SIZE) }
}