regy 0.1.4

use std::fs::File;
use std::io::{self, Read, Write};
use std::os::fd::{AsRawFd, RawFd};

use nix::libc;
use tokio::io::unix::AsyncFd;

use crate::domain::errors::{AgentError, AgentResult, ErrorCode};
use crate::domain::resize::TerminalSize;

pub(crate) const PTY_READ_CHUNK_BYTES: usize = 16 * 1024;

#[derive(Debug)]
pub struct PtyMaster(AsyncFd<File>);

impl PtyMaster {
    pub fn new(file: File) -> io::Result<Self> {
        set_fd_nonblocking(file.as_raw_fd())?;
        AsyncFd::new(file).map(Self)
    }

    pub fn raw_fd(&self) -> RawFd {
        self.0.get_ref().as_raw_fd()
    }

    pub async fn read_next(&self) -> AgentResult<Vec<u8>> {
        let mut buf = vec![0; PTY_READ_CHUNK_BYTES];
        loop {
            let mut ready = self
                .0
                .readable()
                .await
                .map_err(|e| read_err(e.to_string()))?;
            match ready.try_io(|inner| inner.get_ref().read(&mut buf)) {
                Ok(Ok(0)) => return Ok(Vec::new()),
                Ok(Ok(n)) => {
                    buf.truncate(n);
                    return Ok(buf);
                }
                Ok(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
                Ok(Err(e)) if e.raw_os_error() == Some(libc::EIO) => return Ok(Vec::new()),
                Ok(Err(e)) => return Err(read_err(e.to_string())),
                Err(_) => continue,
            }
        }
    }

    pub async fn write_input(&self, bytes: &[u8]) -> AgentResult<()> {
        let mut pending = bytes;
        while !pending.is_empty() {
            let mut ready = self
                .0
                .writable()
                .await
                .map_err(|e| write_err(e.to_string()))?;
            match ready.try_io(|inner| inner.get_ref().write(pending)) {
                Ok(Ok(0)) => return Err(write_err("short write")),
                Ok(Ok(n)) => pending = &pending[n..],
                Ok(Err(e)) if e.kind() == io::ErrorKind::Interrupted => continue,
                Ok(Err(e)) => return Err(write_err(e.to_string())),
                Err(_) => continue,
            }
        }
        Ok(())
    }

    pub fn resize(&self, size: TerminalSize) -> AgentResult<()> {
        let winsize = libc::winsize {
            ws_col: size.cols,
            ws_row: size.rows,
            ws_xpixel: 0,
            ws_ypixel: 0,
        };
        if unsafe { libc::ioctl(self.raw_fd(), libc::TIOCSWINSZ, &winsize) } == -1 {
            return Err(resize_err(io::Error::last_os_error().to_string()));
        }
        Ok(())
    }
}

pub(crate) fn set_fd_cloexec(fd: RawFd) -> io::Result<()> {
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
    if flags == -1 {
        return Err(io::Error::last_os_error());
    }
    if unsafe { libc::fcntl(fd, libc::F_SETFD, flags | libc::FD_CLOEXEC) } == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

fn set_fd_nonblocking(fd: RawFd) -> io::Result<()> {
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
    if flags == -1 {
        return Err(io::Error::last_os_error());
    }
    if unsafe { libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) } == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(())
}

#[cfg(test)]
pub(crate) fn fd_has_cloexec(fd: RawFd) -> io::Result<bool> {
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
    if flags == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(flags & libc::FD_CLOEXEC != 0)
}

#[cfg(test)]
pub(crate) fn fd_has_nonblocking(fd: RawFd) -> io::Result<bool> {
    let flags = unsafe { libc::fcntl(fd, libc::F_GETFL) };
    if flags == -1 {
        return Err(io::Error::last_os_error());
    }
    Ok(flags & libc::O_NONBLOCK != 0)
}

fn read_err(message: impl Into<String>) -> AgentError {
    AgentError::new(ErrorCode::PtyReadFailed, message)
}

fn write_err(message: impl Into<String>) -> AgentError {
    AgentError::new(ErrorCode::PtyWriteFailed, message)
}

fn resize_err(message: impl Into<String>) -> AgentError {
    AgentError::new(ErrorCode::ResizeFailed, message)
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::File;
    use std::os::fd::AsRawFd;

    #[test]
    fn set_fd_cloexec_marks_duplicate_fd() {
        let file = File::open("/dev/null").unwrap();
        let fd = unsafe { libc::dup(file.as_raw_fd()) };
        assert!(fd >= 0);
        assert!(!fd_has_cloexec(fd).unwrap());
        set_fd_cloexec(fd).unwrap();
        assert!(fd_has_cloexec(fd).unwrap());
        unsafe {
            libc::close(fd);
        }
    }
}