//! Async keyboard input.
//!
//! The interrupt handler does the bare minimum: read the scancode and push it
//! into a lock-free queue (`add_scancode`), then wake whoever is waiting.
//! Decoding and printing happen later in a normal async task
//! (`print_keypresses`), outside interrupt context. That keeps the handler
//! short and means it never takes a lock (no deadlock risk).

use crate::{print, serial_print, vga_buffer};
use conquer_once::spin::OnceCell;
use core::{
    pin::Pin,
    task::{Context, Poll},
};
use crossbeam_queue::ArrayQueue;
use futures_util::{
    stream::{Stream, StreamExt},
    task::AtomicWaker,
};
use pc_keyboard::{DecodedKey, HandleControl, KeyCode, Keyboard, ScancodeSet1, layouts};

/// Scancodes waiting to be decoded. `OnceCell` because ArrayQueue::new
/// allocates, so it can only be created after the heap exists.
static SCANCODE_QUEUE: OnceCell<ArrayQueue<u8>> = OnceCell::uninit();
/// The waker of the task waiting for scancodes (if any).
static WAKER: AtomicWaker = AtomicWaker::new();

/// Called by the keyboard interrupt handler. Must not block or allocate.
pub(crate) fn add_scancode(scancode: u8) {
    if let Ok(queue) = SCANCODE_QUEUE.try_get() {
        if queue.push(scancode).is_err() {
            serial_print!("WARNING: scancode queue full; dropping keyboard input\n");
        } else {
            WAKER.wake();
        }
    }
    // Queue not created yet (no keyboard task): just drop the key.
}

/// A `Stream` (async iterator) of scancodes. Only one may exist.
pub struct ScancodeStream {
    _private: (), // prevent construction without `new`
}

impl ScancodeStream {
    pub fn new() -> Self {
        SCANCODE_QUEUE
            .try_init_once(|| ArrayQueue::new(100))
            .expect("ScancodeStream::new should only be called once");
        ScancodeStream { _private: () }
    }
}

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<u8>> {
        let queue = SCANCODE_QUEUE.try_get().expect("scancode queue not initialized");

        // Fast path: something is already there.
        if let Some(scancode) = queue.pop() {
            return Poll::Ready(Some(scancode));
        }

        // Register our waker *then* check again: a key could arrive between
        // the first pop and the registration, and we'd miss its wake-up.
        WAKER.register(cx.waker());
        match queue.pop() {
            Some(scancode) => {
                WAKER.take();
                Poll::Ready(Some(scancode))
            }
            None => Poll::Pending,
        }
    }
}

/// The keyboard task: decode scancodes and echo characters forever.
pub async fn print_keypresses() {
    let mut scancodes = ScancodeStream::new();
    let mut keyboard = Keyboard::new(ScancodeSet1::new(), layouts::Us104Key, HandleControl::Ignore);

    while let Some(scancode) = scancodes.next().await {
        if let Ok(Some(key_event)) = keyboard.add_byte(scancode)
            && let Some(key) = keyboard.process_keyevent(key_event)
        {
            match key {
                // Backspace: erase the previous character on screen. On
                // serial, the uart_16550 crate already turns a 0x08 byte into
                // "\x08 \x08" (back, overwrite with space, back) for us.
                DecodedKey::Unicode('\u{8}') => {
                    vga_buffer::backspace();
                    serial_print!("\x08");
                }
                DecodedKey::Unicode(character) => {
                    print!("{}", character);
                    serial_print!("{}", character);
                }
                // Modifiers only change how other keys decode: print nothing.
                DecodedKey::RawKey(
                    KeyCode::LShift | KeyCode::RShift | KeyCode::LControl | KeyCode::RControl
                    | KeyCode::LAlt | KeyCode::RAltGr | KeyCode::CapsLock,
                ) => {}
                DecodedKey::RawKey(key) => {
                    print!("<{:?}>", key);
                    serial_print!("<{:?}>", key);
                }
            }
        }
    }
}