heroforge-core 0.2.2

Pure Rust core library for reading and writing Fossil SCM repositories
Documentation
//! Socket Client for Heroforge
//!
//! This module provides a client for communicating with the Heroforge daemon
//! via Unix socket.

use std::io::Write;
use std::path::{Path, PathBuf};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::UnixStream;

use super::error::{RhaiError, RhaiResult};
use super::socket_server::default_socket_path;

/// Check if a marker line indicates end of response.
fn is_response_end(line: &str) -> bool {
    let trimmed = line.trim();
    trimmed.len() >= 2 && trimmed.chars().all(|c| c == '=')
}

/// Client for communicating with the Heroforge daemon.
pub struct SocketClient {
    socket_path: PathBuf,
}

impl SocketClient {
    /// Create a new client with the specified socket path.
    pub fn new(socket_path: PathBuf) -> Self {
        Self { socket_path }
    }

    /// Create with default socket path.
    pub fn with_default_path() -> Self {
        Self::new(default_socket_path())
    }

    /// Get the socket path.
    pub fn socket_path(&self) -> &Path {
        &self.socket_path
    }

    /// Check if the daemon is running.
    pub async fn ping(&self) -> bool {
        super::socket_server::ping(&self.socket_path).await
    }

    /// Send a script to the daemon and stream output to stdout.
    pub async fn send_script_streaming(&self, script: &str) -> RhaiResult<()> {
        let stream = UnixStream::connect(&self.socket_path).await.map_err(|e| {
            RhaiError::SocketError(format!(
                "Failed to connect to {}: {}",
                self.socket_path.display(),
                e
            ))
        })?;

        let (reader, mut writer) = stream.into_split();
        let mut reader = BufReader::new(reader);

        // Send script with delimiter
        writer
            .write_all(script.as_bytes())
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to send script: {}", e)))?;

        if !script.ends_with('\n') {
            writer
                .write_all(b"\n")
                .await
                .map_err(|e| RhaiError::SocketError(format!("Failed to send newline: {}", e)))?;
        }

        writer
            .write_all(b"===\n")
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to send delimiter: {}", e)))?;

        writer
            .flush()
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to flush: {}", e)))?;

        // Read status line
        let mut status_line = String::new();
        reader
            .read_line(&mut status_line)
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to read status: {}", e)))?;

        let status = status_line.trim();
        if status != "ok" {
            return Err(RhaiError::SocketError(format!(
                "Unexpected status: {}",
                status
            )));
        }

        // Stream output to stdout in real-time
        let mut line = String::new();
        let mut has_error = false;
        let mut error_msg = String::new();

        loop {
            line.clear();
            let bytes_read = reader
                .read_line(&mut line)
                .await
                .map_err(|e| RhaiError::SocketError(format!("Failed to read output: {}", e)))?;

            if bytes_read == 0 {
                break;
            }

            if is_response_end(&line) {
                break;
            }

            if line.starts_with("ERROR: ") {
                has_error = true;
                error_msg = line[7..].trim().to_string();
            }

            // Print immediately to stdout
            print!("{}", line);
            std::io::stdout().flush().ok();
        }

        if has_error {
            Err(RhaiError::ScriptError(error_msg))
        } else {
            Ok(())
        }
    }

    /// Send a script and collect all output.
    pub async fn send_script(&self, script: &str) -> RhaiResult<String> {
        let stream = UnixStream::connect(&self.socket_path).await.map_err(|e| {
            RhaiError::SocketError(format!(
                "Failed to connect to {}: {}",
                self.socket_path.display(),
                e
            ))
        })?;

        let (reader, mut writer) = stream.into_split();
        let mut reader = BufReader::new(reader);

        // Send script with delimiter
        writer
            .write_all(script.as_bytes())
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to send script: {}", e)))?;

        if !script.ends_with('\n') {
            writer
                .write_all(b"\n")
                .await
                .map_err(|e| RhaiError::SocketError(format!("Failed to send newline: {}", e)))?;
        }

        writer
            .write_all(b"===\n")
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to send delimiter: {}", e)))?;

        writer
            .flush()
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to flush: {}", e)))?;

        // Read status line
        let mut status_line = String::new();
        reader
            .read_line(&mut status_line)
            .await
            .map_err(|e| RhaiError::SocketError(format!("Failed to read status: {}", e)))?;

        let status = status_line.trim();
        if status != "ok" {
            return Err(RhaiError::SocketError(format!(
                "Unexpected status: {}",
                status
            )));
        }

        // Collect all output
        let mut output = String::new();
        let mut line = String::new();
        let mut has_error = false;
        let mut error_msg = String::new();

        loop {
            line.clear();
            let bytes_read = reader
                .read_line(&mut line)
                .await
                .map_err(|e| RhaiError::SocketError(format!("Failed to read output: {}", e)))?;

            if bytes_read == 0 {
                break;
            }

            if is_response_end(&line) {
                break;
            }

            if line.starts_with("ERROR: ") {
                has_error = true;
                error_msg = line[7..].trim().to_string();
            }

            output.push_str(&line);
        }

        if has_error {
            Err(RhaiError::ScriptError(error_msg))
        } else {
            Ok(output)
        }
    }
}

/// Send a script to the daemon (convenience function).
pub async fn send_script(socket_path: &Path, script: &str) -> RhaiResult<String> {
    SocketClient::new(socket_path.to_path_buf())
        .send_script(script)
        .await
}

/// Send a script and stream output (convenience function).
pub async fn send_script_streaming(socket_path: &Path, script: &str) -> RhaiResult<()> {
    SocketClient::new(socket_path.to_path_buf())
        .send_script_streaming(script)
        .await
}

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

    #[test]
    fn test_is_response_end() {
        assert!(is_response_end("====="));
        assert!(is_response_end("==="));
        assert!(is_response_end("  =====  "));
        assert!(!is_response_end("="));
        assert!(!is_response_end("abc"));
    }
}