bux 0.9.0

//! Async host-side client for communicating with a bux guest agent.
//!
//! Each operation opens a **dedicated connection** to the guest agent.
//! This eliminates contention — multiple execs, file transfers, and control
//! operations can proceed concurrently without any locking.

use std::io;
use std::path::PathBuf;

use bux_proto::{
    ControlReq, ControlResp, ExecIn, ExecOut, ExecStart, Hello, HelloAck, PROTOCOL_VERSION,
    STREAM_CHUNK_SIZE, UploadResult,
};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite};
use tokio::net::UnixStream;
use tokio::net::unix::{OwnedReadHalf, OwnedWriteHalf};

/// Output captured from a completed exec.
#[derive(Debug)]
#[non_exhaustive]
pub struct ExecOutput {
    /// Unique execution identifier assigned by the guest.
    pub exec_id: String,
    /// Child process ID inside the guest.
    pub pid: i32,
    /// Captured stdout bytes.
    pub stdout: Vec<u8>,
    /// Captured stderr bytes (empty in TTY mode).
    pub stderr: Vec<u8>,
    /// Process exit code.
    pub code: i32,
    /// Signal that terminated the process, if any.
    pub signal: Option<i32>,
    /// Whether the exec was killed due to timeout.
    pub timed_out: bool,
    /// Wall-clock duration in milliseconds.
    pub duration_ms: u64,
    /// Error message from the guest agent, if any.
    pub error_message: Option<String>,
}

/// Information returned by a successful ping.
#[derive(Debug)]
#[non_exhaustive]
pub struct PongInfo {
    /// Guest agent version string.
    pub version: String,
    /// Guest uptime in milliseconds.
    pub uptime_ms: u64,
}

/// Handle to a running exec with a dedicated connection.
///
/// The connection is split into read/write halves so stdin writes and
/// stdout/stderr reads proceed concurrently without deadlock.
#[derive(Debug)]
pub struct ExecHandle {
    /// Unique execution identifier assigned by the guest.
    exec_id: String,
    /// Child process ID inside the guest.
    pid: i32,
    /// Read half — receives [`ExecOut`] messages from the guest.
    reader: OwnedReadHalf,
    /// Write half — sends [`ExecIn`] messages to the guest.
    writer: OwnedWriteHalf,
}

impl ExecHandle {
    #[allow(
        clippy::missing_docs_in_private_items,
        reason = "private helper with clear signature"
    )]
    async fn collect_output(
        exec_id: String,
        pid: i32,
        reader: &mut OwnedReadHalf,
        mut on: impl FnMut(&ExecOut),
    ) -> io::Result<ExecOutput> {
        let mut stdout = Vec::new();
        let mut stderr = Vec::new();
        loop {
            let msg = bux_proto::recv(reader).await?;
            on(&msg);
            match msg {
                ExecOut::Stdout(d) => stdout.extend(d),
                ExecOut::Stderr(d) => stderr.extend(d),
                ExecOut::Exit {
                    code,
                    signal,
                    timed_out,
                    duration_ms,
                    error_message,
                } => {
                    return Ok(ExecOutput {
                        exec_id,
                        pid,
                        stdout,
                        stderr,
                        code,
                        signal,
                        timed_out,
                        duration_ms,
                        error_message,
                    });
                }
                ExecOut::Error(e) => return Err(io::Error::other(e)),
                _ => {}
            }
        }
    }

    /// Unique execution identifier.
    #[must_use]
    pub fn exec_id(&self) -> &str {
        &self.exec_id
    }

    /// Process ID inside the guest.
    #[must_use]
    pub const fn pid(&self) -> i32 {
        self.pid
    }

    /// Writes data to the process's stdin.
    ///
    /// # Errors
    ///
    /// Returns an error if sending to the guest fails.
    pub async fn write_stdin(&mut self, data: &[u8]) -> io::Result<()> {
        bux_proto::send(&mut self.writer, &ExecIn::Stdin(data.to_vec())).await
    }

    /// Closes the process's stdin (sends EOF).
    ///
    /// # Errors
    ///
    /// Returns an error if sending to the guest fails.
    pub async fn close_stdin(&mut self) -> io::Result<()> {
        bux_proto::send(&mut self.writer, &ExecIn::StdinClose).await
    }

    /// Sends a POSIX signal to the process.
    ///
    /// # Errors
    ///
    /// Returns an error if sending to the guest fails.
    pub async fn signal(&mut self, sig: i32) -> io::Result<()> {
        bux_proto::send(&mut self.writer, &ExecIn::Signal(sig)).await
    }

    /// Resizes the PTY window (only for TTY sessions).
    ///
    /// # Errors
    ///
    /// Returns an error if sending to the guest fails.
    pub async fn resize_tty(
        &mut self,
        rows: u16,
        cols: u16,
        x_pixels: u16,
        y_pixels: u16,
    ) -> io::Result<()> {
        bux_proto::send(
            &mut self.writer,
            &ExecIn::ResizeTty(bux_proto::TtyConfig::new(rows, cols, x_pixels, y_pixels)),
        )
        .await
    }

    /// Reads the next output event from the guest.
    ///
    /// Returns `None` when the connection closes unexpectedly.
    ///
    /// # Errors
    ///
    /// Returns an error if receiving from the guest fails.
    pub async fn next_output(&mut self) -> io::Result<ExecOut> {
        bux_proto::recv(&mut self.reader).await
    }

    /// Waits for the process to exit, collecting all output.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection or protocol exchange fails.
    pub async fn wait_with_output(self) -> io::Result<ExecOutput> {
        let Self {
            exec_id,
            pid,
            mut reader,
            writer: _,
        } = self;
        Self::collect_output(exec_id, pid, &mut reader, |_| {}).await
    }

    /// Streams output via callback, returns collected output.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection or protocol exchange fails.
    pub async fn stream(self, on: impl FnMut(&ExecOut)) -> io::Result<ExecOutput> {
        let Self {
            exec_id,
            pid,
            mut reader,
            writer: _,
        } = self;
        Self::collect_output(exec_id, pid, &mut reader, on).await
    }

    #[allow(missing_docs, reason = "API pending stabilization")]
    /// # Errors
    ///
    /// Returns an error if the connection or protocol exchange fails.
    ///
    /// # Panics
    ///
    /// Should not panic in practice; the internal `expect` is guarded
    /// by the read length.
    pub async fn stream_with_input<R>(
        self,
        mut input: R,
        on: impl FnMut(&ExecOut),
    ) -> io::Result<ExecOutput>
    where
        R: AsyncRead + Unpin + Send + 'static,
    {
        let Self {
            exec_id,
            pid,
            mut reader,
            mut writer,
        } = self;
        #[allow(
            clippy::excessive_nesting,
            reason = "inherent in async spawn + loop pattern"
        )]
        let stdin_task = tokio::spawn(async move {
            let mut buf = [0_u8; 8192];
            loop {
                let n = input.read(&mut buf).await?;
                if n == 0 {
                    return bux_proto::send(&mut writer, &ExecIn::StdinClose).await;
                }
                #[allow(clippy::expect_used, reason = "n is bounded by buf.len()")]
                {
                    bux_proto::send(
                        &mut writer,
                        &ExecIn::Stdin(buf.get(..n).expect("n <= buf.len()").to_vec()),
                    )
                    .await?;
                }
            }
        });

        let output = Self::collect_output(exec_id, pid, &mut reader, on).await;
        stdin_task.abort();
        match stdin_task.await {
            Ok(Err(err))
                if matches!(
                    err.kind(),
                    io::ErrorKind::BrokenPipe
                        | io::ErrorKind::ConnectionAborted
                        | io::ErrorKind::ConnectionReset
                        | io::ErrorKind::UnexpectedEof
                ) => {}
            Ok(Err(err)) if output.is_ok() => return Err(err),
            Err(join_err) if join_err.is_cancelled() => {}
            Err(join_err) if output.is_ok() => return Err(io::Error::other(join_err)),
            Ok(Ok(()) | Err(_)) | Err(_) => {}
        }
        output
    }
}

/// Stateless connection factory to a running guest agent.
///
/// Each method opens a **dedicated connection**, sends a [`Hello`] message
/// to identify the operation, and processes the response on that connection.
/// Multiple operations can run concurrently without contention.
#[derive(Debug, Clone)]
pub(crate) struct Client {
    /// Socket path (Unix socket mapped from vsock by libkrun).
    socket_path: PathBuf,
}

impl Client {
    /// Creates a new client targeting the given Unix socket path.
    ///
    /// Does **not** connect immediately — connections are opened per-operation.
    pub(crate) fn new(path: impl Into<PathBuf>) -> Self {
        Self {
            socket_path: path.into(),
        }
    }

    /// Verifies connectivity and protocol version by opening a control
    /// connection and performing a handshake.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection fails or versions mismatch.
    pub(crate) async fn handshake(&self) -> io::Result<()> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::Control {
                version: PROTOCOL_VERSION,
            },
        )
        .await?;
        match bux_proto::recv::<HelloAck>(&mut stream).await? {
            HelloAck::Control { version } if version == PROTOCOL_VERSION => Ok(()),
            HelloAck::Control { version } => Err(io::Error::new(
                io::ErrorKind::Unsupported,
                format!("protocol version mismatch: host={PROTOCOL_VERSION}, guest={version}"),
            )),
            HelloAck::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "expected Control ack",
            )),
        }
    }

    /// Requests graceful shutdown of the guest agent.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection or shutdown request fails.
    pub(crate) async fn shutdown(&self) -> io::Result<()> {
        let mut stream = self.open_control().await?;
        bux_proto::send(&mut stream, &ControlReq::Shutdown).await?;
        match bux_proto::recv::<ControlResp>(&mut stream).await? {
            ControlResp::ShutdownOk => Ok(()),
            ControlResp::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "expected ShutdownOk",
            )),
        }
    }

    /// Pings the guest agent and returns agent metadata.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection or ping fails.
    pub(crate) async fn ping(&self) -> io::Result<PongInfo> {
        let mut stream = self.open_control().await?;
        bux_proto::send(&mut stream, &ControlReq::Ping).await?;
        match bux_proto::recv::<ControlResp>(&mut stream).await? {
            ControlResp::Pong { version, uptime_ms } => Ok(PongInfo { version, uptime_ms }),
            ControlResp::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(io::ErrorKind::InvalidData, "expected Pong")),
        }
    }

    /// Freezes all writable guest filesystems (FIFREEZE).
    ///
    /// # Errors
    ///
    /// Returns an error if the quiesce operation fails.
    pub(crate) async fn quiesce(&self) -> io::Result<u32> {
        let mut stream = self.open_control().await?;
        bux_proto::send(&mut stream, &ControlReq::Quiesce).await?;
        match bux_proto::recv::<ControlResp>(&mut stream).await? {
            ControlResp::QuiesceOk { frozen_count } => Ok(frozen_count),
            ControlResp::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "expected QuiesceOk",
            )),
        }
    }

    /// Thaws previously frozen guest filesystems (FITHAW).
    ///
    /// # Errors
    ///
    /// Returns an error if the thaw operation fails.
    pub(crate) async fn thaw(&self) -> io::Result<u32> {
        let mut stream = self.open_control().await?;
        bux_proto::send(&mut stream, &ControlReq::Thaw).await?;
        match bux_proto::recv::<ControlResp>(&mut stream).await? {
            ControlResp::ThawOk { thawed_count } => Ok(thawed_count),
            ControlResp::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "expected ThawOk",
            )),
        }
    }

    /// Starts a command on a dedicated exec connection.
    ///
    /// Returns an [`ExecHandle`] for reading output and writing stdin.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection or exec start fails.
    pub(crate) async fn exec(&self, req: ExecStart) -> io::Result<ExecHandle> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(&mut stream, &Hello::Exec(req)).await?;
        match bux_proto::recv::<HelloAck>(&mut stream).await? {
            HelloAck::ExecStarted { exec_id, pid } => {
                let (reader, writer) = stream.into_split();
                Ok(ExecHandle {
                    exec_id,
                    pid,
                    reader,
                    writer,
                })
            }
            HelloAck::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "expected ExecStarted",
            )),
        }
    }

    /// Executes a command and collects all output.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection or command execution fails.
    pub(crate) async fn exec_output(&self, req: ExecStart) -> io::Result<ExecOutput> {
        self.exec(req).await?.wait_with_output().await
    }

    /// Reads a file from the guest filesystem.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be read.
    pub(crate) async fn read_file(&self, path: &str) -> io::Result<Vec<u8>> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::FileRead {
                path: path.to_owned(),
            },
        )
        .await?;
        Self::expect_ready(&mut stream).await?;
        bux_proto::recv_download(&mut stream).await
    }

    /// Writes a file to the guest filesystem.
    ///
    /// # Errors
    ///
    /// Returns an error if the file cannot be written.
    pub(crate) async fn write_file(&self, path: &str, data: &[u8], mode: u32) -> io::Result<()> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::FileWrite {
                path: path.to_owned(),
                mode,
            },
        )
        .await?;
        Self::expect_ready(&mut stream).await?;
        bux_proto::send_upload(&mut stream, data, STREAM_CHUNK_SIZE).await?;
        Self::expect_upload_ok(&mut stream).await
    }

    /// Copies a tar archive into the guest, unpacking at `dest`.
    ///
    /// # Errors
    ///
    /// Returns an error if the copy operation fails.
    pub(crate) async fn copy_in(&self, dest: &str, tar_data: &[u8]) -> io::Result<()> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::CopyIn {
                dest: dest.to_owned(),
            },
        )
        .await?;
        Self::expect_ready(&mut stream).await?;
        bux_proto::send_upload(&mut stream, tar_data, STREAM_CHUNK_SIZE).await?;
        Self::expect_upload_ok(&mut stream).await
    }

    /// Streams a tar archive from `reader` into the guest, unpacking at `dest`.
    ///
    /// Unlike [`copy_in`](Self::copy_in), this never loads the entire archive
    /// into memory — `O(chunk_size)` regardless of total size.
    ///
    /// # Errors
    ///
    /// Returns an error if the streaming copy fails.
    pub(crate) async fn copy_in_from_reader(
        &self,
        dest: &str,
        reader: &mut (impl AsyncRead + Unpin + Send),
    ) -> io::Result<()> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::CopyIn {
                dest: dest.to_owned(),
            },
        )
        .await?;
        Self::expect_ready(&mut stream).await?;
        bux_proto::send_upload_from_reader(&mut stream, reader, STREAM_CHUNK_SIZE).await?;
        Self::expect_upload_ok(&mut stream).await
    }

    /// Copies a path from the guest as a tar archive.
    ///
    /// # Errors
    ///
    /// Returns an error if the copy operation fails.
    pub(crate) async fn copy_out(&self, path: &str) -> io::Result<Vec<u8>> {
        self.copy_out_opts(path, false).await
    }

    /// Copies a path from the guest as a tar archive with options.
    ///
    /// # Errors
    ///
    /// Returns an error if the copy operation fails.
    pub(crate) async fn copy_out_opts(
        &self,
        path: &str,
        follow_symlinks: bool,
    ) -> io::Result<Vec<u8>> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::CopyOut {
                path: path.to_owned(),
                follow_symlinks,
            },
        )
        .await?;
        Self::expect_ready(&mut stream).await?;
        bux_proto::recv_download(&mut stream).await
    }

    /// Streams a path from the guest as a tar archive directly to `writer`.
    ///
    /// Unlike [`copy_out`](Self::copy_out), this never loads the entire archive
    /// into memory — `O(chunk_size)` regardless of total size.
    ///
    /// # Errors
    ///
    /// Returns an error if the streaming copy fails.
    pub(crate) async fn copy_out_to_writer(
        &self,
        path: &str,
        follow_symlinks: bool,
        writer: &mut (impl AsyncWrite + Unpin + Send),
    ) -> io::Result<u64> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::CopyOut {
                path: path.to_owned(),
                follow_symlinks,
            },
        )
        .await?;
        Self::expect_ready(&mut stream).await?;
        bux_proto::recv_download_to_writer(&mut stream, writer).await
    }

    /// Opens a raw Unix socket connection to the guest agent.
    async fn connect_raw(&self) -> io::Result<UnixStream> {
        UnixStream::connect(&self.socket_path).await
    }

    /// Opens a control connection (`Hello::Control` + `HelloAck::Control`).
    async fn open_control(&self) -> io::Result<UnixStream> {
        let mut stream = self.connect_raw().await?;
        bux_proto::send(
            &mut stream,
            &Hello::Control {
                version: PROTOCOL_VERSION,
            },
        )
        .await?;
        match bux_proto::recv::<HelloAck>(&mut stream).await? {
            HelloAck::Control { version } if version == PROTOCOL_VERSION => Ok(stream),
            HelloAck::Control { version } => Err(io::Error::new(
                io::ErrorKind::Unsupported,
                format!("protocol version mismatch: host={PROTOCOL_VERSION}, guest={version}"),
            )),
            HelloAck::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "expected Control ack",
            )),
        }
    }

    /// Expects a `HelloAck::Ready` response.
    async fn expect_ready(
        stream: &mut (impl AsyncRead + AsyncWrite + Unpin + Send),
    ) -> io::Result<()> {
        match bux_proto::recv::<HelloAck>(stream).await? {
            HelloAck::Ready => Ok(()),
            HelloAck::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "expected Ready ack",
            )),
        }
    }

    /// Expects an `UploadResult::Ok` response.
    async fn expect_upload_ok(
        stream: &mut (impl AsyncRead + AsyncWrite + Unpin + Send),
    ) -> io::Result<()> {
        match bux_proto::recv::<UploadResult>(stream).await? {
            UploadResult::Ok => Ok(()),
            UploadResult::Error(e) => Err(io::Error::other(e)),
            _ => Err(io::Error::other("unexpected upload result")),
        }
    }
}