use std::fmt;

pub trait PtySession {
    fn write(&mut self, bytes: &[u8]) -> Result<(), PtyError>;
    fn resize(&mut self, cols: u16, rows: u16) -> Result<(), PtyError>;
    fn try_read(&mut self) -> Result<Vec<u8>, PtyError>;
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PtyError {
    InvalidSize,
    Closed,
    Io(String),
}

impl fmt::Display for PtyError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidSize => write!(f, "invalid PTY size"),
            Self::Closed => write!(f, "PTY session is closed"),
            Self::Io(message) => write!(f, "PTY I/O error: {message}"),
        }
    }
}

impl std::error::Error for PtyError {}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PtySize {
    pub cols: u16,
    pub rows: u16,
}

impl PtySize {
    pub fn new(cols: u16, rows: u16) -> Result<Self, PtyError> {
        if cols == 0 || rows == 0 {
            return Err(PtyError::InvalidSize);
        }
        Ok(Self { cols, rows })
    }
}

#[derive(Debug, Clone)]
pub struct FakePtySession {
    written: Vec<u8>,
    output: Vec<u8>,
    size: PtySize,
    closed: bool,
}

impl FakePtySession {
    pub fn new(cols: u16, rows: u16) -> Self {
        Self {
            written: Vec::new(),
            output: Vec::new(),
            size: PtySize::new(cols, rows).expect("fake PTY size is valid"),
            closed: false,
        }
    }

    pub fn queue_output(&mut self, bytes: impl AsRef<[u8]>) {
        self.output.extend_from_slice(bytes.as_ref());
    }

    pub fn written(&self) -> &[u8] {
        &self.written
    }

    pub fn size(&self) -> &PtySize {
        &self.size
    }
}

impl PtySession for FakePtySession {
    fn write(&mut self, bytes: &[u8]) -> Result<(), PtyError> {
        if self.closed {
            return Err(PtyError::Closed);
        }
        self.written.extend_from_slice(bytes);
        Ok(())
    }

    fn resize(&mut self, cols: u16, rows: u16) -> Result<(), PtyError> {
        self.size = PtySize::new(cols, rows)?;
        Ok(())
    }

    fn try_read(&mut self) -> Result<Vec<u8>, PtyError> {
        if self.closed {
            return Err(PtyError::Closed);
        }
        Ok(std::mem::take(&mut self.output))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn fake_pty_records_writes_and_resizes() {
        let mut pty = FakePtySession::new(80, 24);

        pty.write(b"echo ok\n").unwrap();
        pty.resize(120, 40).unwrap();

        assert_eq!(pty.written(), b"echo ok\n");
        assert_eq!(
            pty.size(),
            &PtySize {
                cols: 120,
                rows: 40
            }
        );
    }

    #[test]
    fn fake_pty_reads_queued_output_once() {
        let mut pty = FakePtySession::new(80, 24);
        pty.queue_output(b"ok");

        assert_eq!(pty.try_read().unwrap(), b"ok");
        assert!(pty.try_read().unwrap().is_empty());
    }

    #[test]
    fn rejects_zero_sized_pty() {
        assert_eq!(PtySize::new(0, 24).unwrap_err(), PtyError::InvalidSize);
        assert_eq!(PtySize::new(80, 0).unwrap_err(), PtyError::InvalidSize);
    }
}