FoundryOS/kernel/src/arch/x86_64/drivers/keyboard.rs

use core::{
    pin::Pin,
    task::{Context, Poll},
};

use crate::println;
use crossbeam::queue::ArrayQueue;
use futures_util::{Stream, StreamExt, task::AtomicWaker};
use pc_keyboard::{
    DecodedKey, HandleControl, KeyCode, Keyboard, ScancodeSet1,
    layouts::Uk105Key,
};
use spin::{Lazy, Mutex, Once};

static KBD_QUEUE: Once<ArrayQueue<u8>> = Once::new();
static WAKER: AtomicWaker = AtomicWaker::new();

pub static KEYBOARD: Lazy<Mutex<Keyboard<Uk105Key, ScancodeSet1>>> =
    Lazy::new(|| {
        Mutex::new(Keyboard::new(
            ScancodeSet1::new(),
            // TODO: Expose an API to change the default KB layout.
            Uk105Key,
            HandleControl::Ignore,
        ))
    });
pub static SCANCODE_STREAM: Lazy<Mutex<ScancodeStream>> =
    Lazy::new(|| Mutex::new(ScancodeStream::new()));

pub fn add_scancode(scancode: u8) {
    if let Some(queue) = KBD_QUEUE.get() {
        if queue.push(scancode).is_err() {
            println!("WARNING: scancode queue full; dropping keyboard input");
        } else {
            WAKER.wake();
        }
    } else {
        println!("WARNING: scancode queue not initialized");
    }
}

pub struct ScancodeStream {
    _private: (),
}

impl ScancodeStream {
    pub fn new() -> Self {
        KBD_QUEUE.call_once(|| ArrayQueue::new(5));
        Self { _private: () }
    }

    pub fn try_next(&mut self) -> Option<u8> {
        KBD_QUEUE.get().and_then(|queue| queue.pop())
    }
}

impl Default for ScancodeStream {
    fn default() -> Self {
        Self::new()
    }
}

impl Stream for ScancodeStream {
    type Item = u8;

    fn poll_next(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        let queue = KBD_QUEUE.get().unwrap();

        if let Some(scancode) = queue.pop() {
            return Poll::Ready(Some(scancode));
        }

        WAKER.register(cx.waker());
        WAKER.register(cx.waker());

        queue.pop().map_or(Poll::Pending, |scancode| {
            WAKER.take();
            Poll::Ready(Some(scancode))
        })
    }
}

pub async fn get_keystroke_async() -> KeyStroke {
    loop {
        if let Some(scancode) = SCANCODE_STREAM.lock().next().await {
            if let Ok(keystroke) = KeyStroke::try_from(scancode) {
                return keystroke;
            }
        }
    }
}

pub fn get_keystroke_optional() -> Option<KeyStroke> {
    if let Some(scancode) = SCANCODE_STREAM.lock().try_next() {
        if let Ok(keystroke) = KeyStroke::try_from(scancode) {
            return Some(keystroke);
        }
    }
    None
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum KeyStroke {
    Char(char),
    Ctrl,
    RCtrl,
    Alt,
    RAlt,
    Shift,
    RShift,
    Meta,
    RMeta,
    Backspace,
    Left,
    Right,
    Up,
    Down,
    None,
    Enter,
    Escape,
    Del,
}

impl KeyStroke {
    pub const fn from_keycode(key: KeyCode) -> Self {
        match key {
            KeyCode::LControl => Self::Ctrl,
            KeyCode::RControl => Self::RCtrl,
            KeyCode::LAlt => Self::Alt,
            KeyCode::RAlt2 => Self::RAlt,
            KeyCode::LShift => Self::Shift,
            KeyCode::RShift => Self::RShift,
            KeyCode::LWin => Self::Meta,
            KeyCode::RWin => Self::RMeta,
            KeyCode::Backspace => Self::Backspace,
            KeyCode::ArrowLeft => Self::Left,
            KeyCode::ArrowRight => Self::Right,
            KeyCode::ArrowUp => Self::Up,
            KeyCode::ArrowDown => Self::Down,
            KeyCode::Return => Self::Enter,
            KeyCode::Escape => Self::Escape,
            KeyCode::Delete => Self::Del,
            _ => Self::None,
        }
    }
}

impl TryFrom<u8> for KeyStroke {
    type Error = ();

    fn try_from(code: u8) -> Result<Self, Self::Error> {
        let mut keyboard = KEYBOARD.lock();

        let key = match keyboard.add_byte(code) {
            Ok(Some(event)) => match keyboard.process_keyevent(event) {
                Some(key) => key,
                _ => return Err(()),
            },
            _ => return Err(()),
        };

        match key {
            DecodedKey::Unicode(ch) => Ok(Self::Char(ch)),
            DecodedKey::RawKey(key) => match Self::from_keycode(key) {
                Self::None => Err(()),
                key => Ok(key),
            },
        }
    }
}

impl TryInto<char> for KeyStroke {
    type Error = ();

    fn try_into(self) -> Result<char, Self::Error> {
        match self {
            Self::Char(c) => Ok(c),
            _ => Err(()),
        }
    }
}

impl core::fmt::Display for KeyStroke {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::Char(c) => write!(f, "{}", c),
            Self::Ctrl => write!(f, "CTRL"),
            Self::RCtrl => write!(f, "RCtrl"),
            Self::Alt => write!(f, "ALT"),
            Self::RAlt => write!(f, "RAlt"),
            Self::Shift => write!(f, "SHIFT"),
            Self::RShift => write!(f, "RShift"),
            Self::Meta => write!(f, "META"),
            Self::RMeta => write!(f, "RMeta"),
            Self::Backspace => write!(f, "BACKSPACE"),
            Self::Left => write!(f, "LEFT"),
            Self::Right => write!(f, "RIGHT"),
            Self::Up => write!(f, "UP"),
            Self::Down => write!(f, "DOWN"),
            Self::Enter => write!(f, "ENTER"),
            Self::Escape => write!(f, "ESCAPE"),
            Self::None => write!(f, "NONE"),
            Self::Del => write!(f, "DEL"),
        }
    }
}