yggdrasil/kernel/libk/src/vfs/terminal.rs

use core::{
    mem::MaybeUninit,
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll},
};

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use device_api::device::Device;
use libk_util::{
    ring::{LossyRingQueue, RingBuffer},
    sync::{spin_rwlock::IrqSafeRwLock, IrqSafeSpinlock},
};
use yggdrasil_abi::{
    error::Error,
    io::{
        DeviceRequest, KeyboardKey, KeyboardKeyEvent, TerminalInputOptions, TerminalLineOptions,
        TerminalOptions, TerminalOutputOptions, TerminalSize,
    },
    process::{ProcessGroupId, Signal},
};

use crate::{device::char::CharDevice, task::process::Process};

use super::FileReadiness;

pub struct Terminal<O: TerminalOutput> {
    input: TerminalInput,
    output: O,
    config: IrqSafeRwLock<TerminalOptions>,
}

// Keyboard, PTY slave, terminal, etc. -> program
pub struct TerminalInput {
    buffer: IrqSafeSpinlock<InputBuffer>,
    ready_ring: LossyRingQueue<u8>,
    signal_pgroup: IrqSafeRwLock<Option<ProcessGroupId>>,
}

// Program -> screen, PTY master, terminal, etc.
pub trait TerminalOutput: Sync + Send {
    fn write(&self, byte: u8) -> Result<(), Error>;

    fn notify_readers(&self) {}

    fn size(&self) -> TerminalSize {
        TerminalSize {
            rows: 24,
            columns: 80,
        }
    }
    fn set_size(&self, size: TerminalSize) -> Result<(), Error> {
        let _ = size;
        Err(Error::InvalidOperation)
    }

    fn write_multiple(&self, bytes: &[u8]) -> Result<usize, Error> {
        let mut written = 0;
        for &byte in bytes {
            self.write(byte)?;
            written += 1;
        }
        Ok(written)
    }
}

struct InputBuffer {
    pending: Box<[MaybeUninit<u8>]>,
    position: usize,
}

impl<O: TerminalOutput> Terminal<O> {
    pub fn from_parts(config: TerminalOptions, input: TerminalInput, output: O) -> Self {
        Self {
            input,
            output,
            config: IrqSafeRwLock::new(config),
        }
    }

    pub fn input(&self) -> &TerminalInput {
        &self.input
    }

    pub fn output(&self) -> &O {
        &self.output
    }

    pub fn config(&self) -> TerminalOptions {
        *self.config.read()
    }

    #[inline]
    pub async fn read_from_input(&self, buffer: &mut [u8]) -> Result<usize, Error> {
        let eof = {
            let config = self.config.read();
            config.is_canonical().then_some(config.chars.eof)
        };
        self.input.read(buffer, eof).await
    }

    #[inline]
    pub fn read_from_input_nonblocking(&self, buffer: &mut [u8]) -> Result<usize, Error> {
        let eof = {
            let config = self.config.read();
            config.is_canonical().then_some(config.chars.eof)
        };
        self.input.read_nonblocking(buffer, eof)
    }

    pub fn putc_to_output(&self, byte: u8) -> Result<(), Error> {
        self.output.write(byte)
    }

    pub fn write_to_output(&self, buffer: &[u8]) -> Result<usize, Error> {
        // TODO handle options
        self.output.write_multiple(buffer)
    }

    pub fn write_to_input(&self, mut byte: u8) {
        let config = self.config.read();
        let mut buffer = self.input.buffer.lock();

        if byte == b'\r' && config.input.contains(TerminalInputOptions::CR_TO_NL) {
            byte = b'\n';
        }

        if config.is_canonical() {
            // Canonical line processing

            // Echo back
            if byte == config.chars.erase {
                let echo =
                    buffer.erase_pending() && config.line.contains(TerminalLineOptions::ECHO_ERASE);

                if echo {
                    self.output.write_multiple(b"\x1B[D \x1B[D").ok();
                }

                return;
            }

            if byte == b'\n' {
                // TODO NL_TO_CRNL
                if config.is_echo_newline() {
                    if config.output.contains(TerminalOutputOptions::NL_TO_CRNL) {
                        self.output.write(b'\r').ok();
                    }
                    self.output.write(byte).ok();
                }
            } else if byte.is_ascii_control() {
                if config.line.contains(TerminalLineOptions::ECHO) {
                    self.output.write_multiple(&[b'^', byte + 0x40]).ok();
                }
            } else if config.line.contains(TerminalLineOptions::ECHO) {
                self.output.write(byte).ok();
            }

            if byte == config.chars.interrupt {
                if config.line.contains(TerminalLineOptions::SIGNAL) {
                    self.output.notify_readers();

                    if let Some(group_id) = *self.input.signal_pgroup.read() {
                        Process::signal_group(None, group_id, Signal::Interrupted);
                        self.input.ready_ring.notify_all();
                        return;
                    }
                } else {
                    buffer.write_pending(byte);
                }
            } else {
                buffer.write_pending(byte);
            }

            if byte == b'\n' || byte == config.chars.eof {
                let data = buffer.flush();
                self.input.ready_ring.write_multiple(data);
            }
        } else {
            // Raw line processing
            self.input.ready_ring.write(byte);
        }
    }

    pub fn handle_device_request(&self, req: &mut DeviceRequest) -> Result<(), Error> {
        match req {
            DeviceRequest::SetTerminalGroup(group) => {
                self.input.set_signal_pgroup(*group);
                Ok(())
            }
            DeviceRequest::GetTerminalOptions(options) => {
                options.write(*self.config.read());
                Ok(())
            }
            DeviceRequest::SetTerminalOptions(options) => {
                self.input.flush();
                *self.config.write() = *options;
                Ok(())
            }
            DeviceRequest::SetTerminalSize(size) => {
                // TODO SIGWINCH?
                self.output.set_size(*size)
            }
            DeviceRequest::GetTerminalSize(size) => {
                size.write(self.output.size());
                Ok(())
            }
            _ => Err(Error::InvalidOperation),
        }
    }

    pub async fn consume_keyboard(
        self: Arc<Self>,
        keyboard: Arc<dyn CharDevice>,
        stop: Option<&AtomicBool>,
    ) {
        let mut buf = [0; 4];
        let mut lshift = false;
        let mut rshift = false;
        let mut lctrl = false;
        let mut rctrl = false;

        loop {
            if stop.map(|r| r.load(Ordering::Acquire)).unwrap_or(false) {
                break;
            }

            keyboard.read(&mut buf).await.unwrap();
            let (key, pressed) = KeyboardKeyEvent::from_bytes(buf).split();
            let ctrl = lctrl || rctrl;
            let shift = lshift || rshift;

            match key {
                KeyboardKey::LShift => lshift = pressed,
                KeyboardKey::RShift => rshift = pressed,
                KeyboardKey::LControl => lctrl = pressed,
                KeyboardKey::RControl => rctrl = pressed,
                KeyboardKey::Char(ch) if ctrl && pressed => match ch {
                    b'c' => self.write_to_input(self.config().chars.interrupt),
                    b'd' => self.write_to_input(self.config().chars.eof),
                    _ => (),
                },
                KeyboardKey::Char(ch) if shift && pressed => {
                    let ch = match ch {
                        ch if ch.is_ascii_lowercase() => ch.to_ascii_uppercase(),
                        b'`' => b'~',
                        b'1' => b'!',
                        b'2' => b'@',
                        b'3' => b'#',
                        b'4' => b'$',
                        b'5' => b'%',
                        b'6' => b'^',
                        b'7' => b'&',
                        b'8' => b'*',
                        b'9' => b'(',
                        b'0' => b')',
                        b'-' => b'_',
                        b'=' => b'+',
                        b';' => b':',
                        b'\'' => b'"',
                        b'\\' => b'|',
                        b',' => b'<',
                        b'.' => b'>',
                        b'/' => b'?',
                        _ => continue,
                    };
                    self.write_to_input(ch);
                }
                KeyboardKey::Char(ch) if pressed => {
                    // TODO shift/ctrl
                    self.write_to_input(ch);
                }
                KeyboardKey::Enter if pressed => {
                    self.write_to_input(b'\n');
                }
                KeyboardKey::Backspace if pressed => {
                    self.write_to_input(b'\x7F');
                }
                KeyboardKey::Escape if pressed => {
                    self.write_to_input(b'\x1B');
                }
                _ => (),
            }
        }
    }
}

impl TerminalInput {
    pub fn with_capacity(capacity: usize) -> Result<Self, Error> {
        Ok(Self {
            buffer: IrqSafeSpinlock::new(InputBuffer::with_capacity(capacity)),
            ready_ring: LossyRingQueue::try_with_capacity(capacity)?,
            signal_pgroup: IrqSafeRwLock::new(None),
        })
    }

    pub async fn read(&self, buffer: &mut [u8], eof: Option<u8>) -> Result<usize, Error> {
        let mut lock = self.ready_ring.read_lock().await;
        Ok(read_all(&mut lock, buffer, eof))
    }

    pub fn read_nonblocking(&self, buffer: &mut [u8], eof: Option<u8>) -> Result<usize, Error> {
        let mut lock = self.ready_ring.try_read_lock().ok_or(Error::WouldBlock)?;
        Ok(read_all(&mut lock, buffer, eof))
    }

    pub fn poll_read(&self, cx: &mut Context<'_>) -> Poll<()> {
        self.ready_ring.poll_readable(cx)
    }

    pub fn flush(&self) {
        let mut lock = self.buffer.lock();
        let data = lock.flush();
        self.ready_ring.write_multiple(data);
    }

    pub fn set_signal_pgroup(&self, value: ProcessGroupId) {
        self.signal_pgroup.write().replace(value);
    }
}

impl InputBuffer {
    pub fn with_capacity(capacity: usize) -> Self {
        Self {
            pending: Box::new_uninit_slice(capacity),
            position: 0,
        }
    }

    pub fn write_pending(&mut self, byte: u8) {
        if self.position == self.pending.len() {
            todo!();
        }

        self.pending[self.position].write(byte);
        self.position += 1;
    }

    pub fn erase_pending(&mut self) -> bool {
        if self.position != 0 {
            self.position -= 1;
            true
        } else {
            false
        }
    }

    pub fn flush(&mut self) -> &[u8] {
        let data = unsafe { MaybeUninit::slice_assume_init_ref(&self.pending[..self.position]) };
        self.position = 0;
        data
    }
}

#[async_trait]
impl<O: TerminalOutput> CharDevice for Terminal<O> {
    async fn write(&self, buf: &[u8]) -> Result<usize, Error> {
        self.write_to_output(buf)
    }

    fn write_nonblocking(&self, buf: &[u8]) -> Result<usize, Error> {
        self.write_to_output(buf)
    }

    async fn read(&self, buf: &mut [u8]) -> Result<usize, Error> {
        self.read_from_input(buf).await
    }

    fn read_nonblocking(&self, buf: &mut [u8]) -> Result<usize, Error> {
        self.read_from_input_nonblocking(buf)
    }

    fn is_readable(&self) -> bool {
        true
    }

    fn is_writeable(&self) -> bool {
        true
    }

    fn is_terminal(&self) -> bool {
        true
    }

    fn device_request(&self, req: &mut DeviceRequest) -> Result<(), Error> {
        self.handle_device_request(req)
    }
}

impl<O: TerminalOutput> FileReadiness for Terminal<O> {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.input().poll_read(cx).map(Ok)
    }
}

impl<O: TerminalOutput> Device for Terminal<O> {
    fn display_name(&self) -> &str {
        "Terminal"
    }
}

pub fn read_all(source: &mut RingBuffer<u8>, target: &mut [u8], eof: Option<u8>) -> usize {
    let mut pos = 0;
    while pos < target.len()
        && let Some(ch) = source.try_read()
    {
        if eof.map(|eof| eof == ch).unwrap_or(false) {
            break;
        }
        target[pos] = ch;
        pos += 1;
    }
    pos
}