//! Pseudo-terminal devices

// TODO handle werase key
use core::{
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll},
};

use alloc::{boxed::Box, sync::Arc};
use async_trait::async_trait;
use libk_util::{ring::BlockingRingQueue, sync::spin_rwlock::IrqSafeRwLock};
use yggdrasil_abi::{
    error::Error,
    io::{TerminalOptions, TerminalSize},
};

use super::{
    file::TerminalHalf,
    terminal::{self, Terminal, TerminalInput, TerminalOutput},
};

const CAPACITY: usize = 32768;

struct PtyOutput {
    ring: BlockingRingQueue<u8>,
    shutdown: AtomicBool,
    size: IrqSafeRwLock<TerminalSize>,
}

impl TerminalOutput for PtyOutput {
    fn write(&self, byte: u8) -> Result<(), Error> {
        block!(self.ring.write(byte).await)
    }

    fn write_multiple(&self, bytes: &[u8]) -> Result<usize, Error> {
        block!(self.ring.write_all(bytes).await)?;
        Ok(bytes.len())
    }

    fn notify_readers(&self) {
        self.ring.notify_readers();
    }

    fn size(&self) -> TerminalSize {
        *self.size.read()
    }

    fn set_size(&self, size: TerminalSize) -> Result<(), Error> {
        // TODO checks
        *self.size.write() = size;
        Ok(())
    }
}

impl PtyOutput {
    pub fn with_capacity(size: TerminalSize, capacity: usize) -> Result<Self, Error> {
        Ok(Self {
            ring: BlockingRingQueue::try_with_capacity(capacity)?,
            shutdown: AtomicBool::new(false),
            size: IrqSafeRwLock::new(size),
        })
    }

    pub fn poll_read(&self, cx: &mut Context<'_>) -> Poll<()> {
        if self.shutdown.load(Ordering::Acquire) || self.ring.poll_readable(cx).is_ready() {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    pub async fn read(&self, buffer: &mut [u8]) -> Result<usize, Error> {
        if self.shutdown.load(Ordering::Acquire) {
            return Ok(0);
        }

        let mut lock = self.ring.read_lock().await;
        let len = terminal::read_all(&mut lock, buffer, None);
        self.ring.notify_writers();
        Ok(len)
    }

    pub fn read_nonblocking(&self, buffer: &mut [u8]) -> Result<usize, Error> {
        if self.shutdown.load(Ordering::Acquire) {
            return Ok(0);
        }

        let mut lock = self.ring.try_read_lock().ok_or(Error::WouldBlock)?;
        let len = terminal::read_all(&mut lock, buffer, None);
        self.ring.notify_writers();
        Ok(len)
    }
}

/// Slave part of a PTY device
#[derive(Clone)]
pub struct PseudoTerminalSlave(Arc<Terminal<PtyOutput>>);

/// Master part of a PTY device
#[derive(Clone)]
pub struct PseudoTerminalMaster(Arc<Terminal<PtyOutput>>);

#[async_trait]
impl TerminalHalf for PseudoTerminalSlave {
    async fn read(&self, buf: &mut [u8]) -> Result<usize, Error> {
        self.0.read_from_input(buf).await
    }

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

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

    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.0.input().poll_read(cx).map(Ok)
    }

    fn device_request(&self, option: u32, buffer: &mut [u8], len: usize) -> Result<usize, Error> {
        self.0.handle_device_request(option, buffer, len)
    }
}

#[async_trait]
impl TerminalHalf for PseudoTerminalMaster {
    /// Reads from the slave-to-master half of the PTY
    async fn read(&self, buf: &mut [u8]) -> Result<usize, Error> {
        self.0.output().read(buf).await
    }

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

    fn write(&self, buf: &[u8]) -> Result<usize, Error> {
        for &byte in buf {
            self.0.write_to_input(byte);
        }
        Ok(buf.len())
    }

    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.0.output().poll_read(cx).map(Ok)
    }

    fn device_request(&self, option: u32, buffer: &mut [u8], len: usize) -> Result<usize, Error> {
        self.0.handle_device_request(option, buffer, len)
    }
}

impl Drop for PseudoTerminalMaster {
    fn drop(&mut self) {
        let config = self.0.config();
        self.0.write_to_input(config.chars.eof);
    }
}

impl Drop for PseudoTerminalSlave {
    fn drop(&mut self) {
        let output = self.0.output();
        output.shutdown.store(true, Ordering::Release);
        output.ring.notify_writers();
        output.ring.notify_readers();
    }
}

pub fn create(
    config: TerminalOptions,
    size: TerminalSize,
) -> Result<(PseudoTerminalMaster, PseudoTerminalSlave), Error> {
    let input = TerminalInput::with_capacity(CAPACITY)?;
    let output = PtyOutput::with_capacity(size, CAPACITY)?;
    let pty = Arc::new(Terminal::from_parts(config, input, output));

    let master = PseudoTerminalMaster(pty.clone());
    let slave = PseudoTerminalSlave(pty);

    Ok((master, slave))
}
