use std::io::{Read, Write};
use anyhow::Result;
use portable_pty::{CommandBuilder, MasterPty, PtySize, native_pty_system};
use tokio::sync::mpsc;

/// A spawned PTY with its child process. Output chunks arrive on `output`.
pub struct PtyHandle {
    master: Box<dyn MasterPty + Send>,
    writer: Box<dyn Write + Send>,
    child: Box<dyn portable_pty::Child + Send + Sync>,
    /// Raw output chunks read off the PTY master (already on the async side).
    pub output: mpsc::Receiver<Vec<u8>>,
}

/// Parameters for spawning a surface's process.
pub struct SpawnSpec {
    pub command: String,
    pub args: Vec<String>,
    pub cwd: std::path::PathBuf,
    pub cols: u16,
    pub rows: u16,
    /// Extra environment variables (e.g. SPACESH_SURFACE_ID).
    pub env: Vec<(String, String)>,
}

impl PtyHandle {
    pub fn spawn(spec: SpawnSpec) -> Result<Self> {
        let pty_system = native_pty_system();
        let pair = pty_system.openpty(PtySize {
            rows: spec.rows,
            cols: spec.cols,
            pixel_width: 0,
            pixel_height: 0,
        })?;

        let mut cmd = CommandBuilder::new(&spec.command);
        for a in &spec.args {
            cmd.arg(a);
        }
        cmd.cwd(&spec.cwd);
        for (k, v) in &spec.env {
            cmd.env(k, v);
        }

        let child = pair.slave.spawn_command(cmd)?;
        // The slave handle must be dropped so the child is the only holder; otherwise
        // EOF is never observed on the master after the child exits.
        drop(pair.slave);

        let writer = pair.master.take_writer()?;
        let mut reader = pair.master.try_clone_reader()?;

        let (tx, rx) = mpsc::channel::<Vec<u8>>(256);
        std::thread::spawn(move || {
            let mut buf = [0u8; 8192];
            loop {
                match reader.read(&mut buf) {
                    Ok(0) => break, // EOF: child closed the pty
                    Ok(n) => {
                        if tx.blocking_send(buf[..n].to_vec()).is_err() {
                            break; // receiver gone
                        }
                    }
                    Err(_) => break,
                }
            }
        });

        Ok(Self {
            master: pair.master,
            writer,
            child,
            output: rx,
        })
    }

    pub fn write_input(&mut self, bytes: &[u8]) -> Result<()> {
        self.writer.write_all(bytes)?;
        self.writer.flush()?;
        Ok(())
    }

    pub fn resize(&self, cols: u16, rows: u16) -> Result<()> {
        self.master.resize(PtySize { rows, cols, pixel_width: 0, pixel_height: 0 })?;
        Ok(())
    }

    /// Best-effort wait for the child's exit code (blocking).
    pub fn wait(&mut self) -> i32 {
        match self.child.wait() {
            Ok(status) => status.exit_code() as i32,
            Err(_) => -1,
        }
    }

    pub fn kill(&mut self) {
        let _ = self.child.kill();
    }
}

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

    fn shell_spec(script: &str) -> SpawnSpec {
        SpawnSpec {
            command: "/bin/sh".into(),
            args: vec!["-c".into(), script.into()],
            cwd: std::env::temp_dir(),
            cols: 80,
            rows: 24,
            env: vec![("SPACESH_SURFACE_ID".into(), "s_test".into())],
        }
    }

    #[tokio::test]
    async fn spawn_echo_produces_output() {
        let mut handle = PtyHandle::spawn(shell_spec("printf SPACESH_OK")).unwrap();
        let mut collected = Vec::new();
        // Drain until EOF (channel closes when the reader thread sees EOF).
        while let Some(chunk) = handle.output.recv().await {
            collected.extend_from_slice(&chunk);
        }
        let text = String::from_utf8_lossy(&collected);
        assert!(text.contains("SPACESH_OK"), "got: {text:?}");
    }

    #[tokio::test]
    async fn resize_does_not_error() {
        let handle = PtyHandle::spawn(shell_spec("sleep 0.2")).unwrap();
        handle.resize(120, 40).unwrap();
    }

    #[tokio::test]
    async fn input_is_echoed_back() {
        // `cat` echoes stdin back to stdout on a pty.
        let mut handle = PtyHandle::spawn(shell_spec("cat")).unwrap();
        handle.write_input(b"hello\n").unwrap();
        let mut collected = Vec::new();
        // Read a few chunks then kill cat to end the stream.
        if let Some(chunk) = handle.output.recv().await {
            collected.extend_from_slice(&chunk);
        }
        handle.kill();
        let text = String::from_utf8_lossy(&collected);
        assert!(text.contains("hello"), "got: {text:?}");
    }
}