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

use core::{
    mem::MaybeUninit,
    task::{Context, Poll},
};

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

use crate::task::process::Process;

use super::{CharDevice, 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 {
    fn write(&self, byte: u8) -> Result<(), Error>;

    fn notify_readers(&self) {}

    fn size(&self) -> TerminalSize {
        TerminalSize {
            rows: 80,
            columns: 24,
        }
    }
    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() {
                    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(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),
        }
    }
}

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(&'static self, buf: &[u8]) -> Result<usize, Error> {
        self.write_to_output(buf)
    }

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

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

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

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

    fn is_writable(&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)
    }
}

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
}