path: root/src/poem/elements/verse.rs

use super::rune::Rune;
use super::stanza::Stanza;
use super::word::Word;
use crate::poem::Poem;
use libc::{waitpid, WNOHANG};
use std::fs::OpenOptions;
use std::io::{self, Read, Write};
use std::os::unix::process::CommandExt;
use std::path::Path;
use std::process::{Command, Output, Stdio};
use std::sync::{Arc, Mutex};

/// A [Stanza] and it's [meter](Rune)
///
/// In addition to a [Stanza] and a [meter](Rune), this also holds a [bool]
/// value called `couplet`, indicating that it needs to accept input on `STDIN`
/// from the previous [Verse].
#[derive(Debug, Clone)]
pub struct Verse {
    pub stanza: Stanza,
    pub couplet: u8,
    pub io: Vec<Rune>,
    pub ip: Stanza,
    pub op: Stanza,
    pub ep: Stanza,
    pub poems: Vec<Poem>,
    pub meter: Rune,
}

impl Verse {
    /// Create a new [Verse]
    ///
    /// Returns a new [Verse], with an empty [Stanza], a meter of [Rune::None],
    /// and `couplet` set to 0.
    ///
    /// # Fields
    /// stanza - The command (a `verb()` and a `clause()`)
    /// couplet - Indicates couplet status
    ///           0: Not a couplet (`cat Cargo.toml`)
    ///           1: Left side of a couplet (`cat Cargo.toml | ...`)
    ///           2: Right side of a couplet (`... | lolcat`)
    ///           3: Sandwiched between couplets (`... | grep Ca | ...`)
    /// io - A list of IO operations ([Rune::Read], [Rune::Write], etc.)
    /// ip - A list of filenames for reading into STDIN when:
    ///      [Rune::Read]
    ///      is specified
    /// op - A list of filenames for redirecting STDOUT to when:
    ///      [Rune::Write],
    ///      [Rune::WriteAll],
    ///      [Rune::Addendum],
    ///      and/or [Rune::AddendumAll]
    ///      is specified
    /// ep - A list of filenames for redirecting STDERR to when:
    ///      [Rune::Write2],
    ///      [Rune::WriteAll],
    ///      [Rune::Addendum2],
    ///      and/or [Rune::AddendumAll]
    ///      is specified
    /// poems - Internal commands to run before the [Verse] is recited
    ///         ([Rune::Poem]).
    /// meter - Determines how the verse is recited in relation to the other
    ///         verses in the [Poem].
    ///         [Rune::None] -> Run the command and print the output
    ///         [Rune::Couplet] -> Pipe the output of this verse into the next
    ///         [Rune::Quiet] -> Run in the background
    ///         [Rune::And] -> Run the next verse only if this verse succeeds
    ///         [Rune::Continue] -> Run the next verse, regardless of whether
    ///                             or not this verse succeeds
    pub fn new() -> Self {
        Verse {
            stanza: Stanza::new(),
            couplet: 0,
            io: Vec::new(),
            ip: Stanza::new(),
            op: Stanza::new(),
            ep: Stanza::new(),
            poems: Vec::new(),
            meter: Rune::None,
        }
    }

    /// Get the [Verse]'s verb
    ///
    /// Return the program to be forked
    pub fn verb(&self) -> String {
        self.stanza[0].clone()
    }

    /// Get the [Verse]'s clause
    ///
    /// Return program arguments, if they exist
    pub fn clause(&self) -> Option<Vec<String>> {
        match self.stanza.len() {
            0 => None,
            1 => None,
            _ => Some(self.stanza[1..].to_vec()),
        }
    }

    /// Alias to [Verse].stanza.push()
    pub fn push(&mut self, word: String) {
        self.stanza.push(word);
    }

    /// Alias to [Verse].stanza.is_empty()
    pub fn is_empty(&self) -> bool {
        self.stanza.is_empty()
    }

    /// Alias to [Verse].stanza.clear()
    pub fn clear(&mut self) {
        self.stanza.clear();
        self.io.clear();
        self.poems.clear();
    }

    /// Check if the [Verse] contains any internal poems
    pub fn poems(&self) -> bool {
        if self.poems.len() > 0 {
            return true;
        }
        false
    }

    /// Push a word to the [Verse]'s [Stanza]
    ///
    /// Push a word to the [Verse]'s [Stanza], or one of its IO channels. If
    /// `word` is empty, this function will simply return without performing
    /// any operations. A channel of [None] will push to the [Verse]'s
    /// [Stanza], while IO [Rune]s will determine which IO channel a word will
    /// get pushed onto.
    ///
    /// # Arguments
    /// word - The word to push onto the [Stanza]/channel
    /// channel - Specifiy which channel to use
    ///
    /// # Examples
    /// ```
    /// word.push("c");
    /// word.push("a");
    /// word.push("t");
    /// verse.add(&mut word, None); // Pushes onto [Stanza]
    /// verse.add(&mut word, Some(Rune::Write)); // Pushes onto [op]
    /// verse.add(&mut word, Some(Rune::WriteAll)); // Pushes onto [op], [ep]
    /// ```
    pub fn add(&mut self, word: &mut Word, channel: Option<Rune>) {
        // Do nothing if the stack is empty
        if word.is_empty() {
            return;
        }

        // Push the word
        match channel {
            Some(Rune::Read) => self.ip.push(word.iter().collect()),
            Some(Rune::Write) | Some(Rune::Addendum) => self.op.push(word.iter().collect()),
            Some(Rune::Write2) | Some(Rune::Addendum2) => self.ep.push(word.iter().collect()),
            Some(Rune::WriteAll) | Some(Rune::AddendumAll) => {
                self.op.push(word.iter().collect());
                self.ep.push(word.iter().collect());
            }
            Some(_) | None => self.push(word.iter().collect()),
        }

        // Clear the stack
        word.clear();
    }

    /// Check if the `verb()` exists in the `$PATH`
    ///
    /// First checks if the `verb()` is a relative or full path. If it is,
    /// check whether or not it exists. If it does exist, return true,
    /// otherwise seeif the `verb()` is cached in our list of binaries. Search is
    /// done in $PATH order.
    ///
    /// # Examples
    /// ```
    /// let bins = vec!["cargo", "ruby", "cat"]
    ///            .into_iter()
    ///            .map(String::from)
    ///            .collect<Vec<String>>();
    ///
    /// let command_success = vec!["cargo", "build", "--release"]
    ///                       .into_iter()
    ///                       .map(String::from)
    ///                       .collect<Vec<String>>();
    ///
    /// let command_fail = vec!["make", "-j8"]
    ///                    .into_iter()
    ///                    .map(String::from)
    ///                    .collect<Vec<String>>();
    ///
    /// let stanza_success = Stanza::new(command_success);
    /// let stanza_fail = Stanza::new(command_fail);
    ///
    /// stanza_success.spellcheck(bins) // -> true
    /// stanza_fail.spellcheck(bins) // -> false
    /// ```
    pub fn spellcheck(&self, bins: &Vec<String>) -> bool {
        // An empty verb (i.e. the empty string) cannot be a program, so
        // return false
        // Thanks to the parsing in Poem::read, however, it's
        // unlikely for this to happen
        if self.verb().is_empty() {
            return false;
        }

        // Only search the $PATH if a full or relative path was not given, or
        // if the path given does not exist
        if !Path::new(self.verb().as_str()).exists() {
            // Try to find a binary in our path with the same name as the verb
            // Searches in $PATH order
            match bins
                .iter()
                .find(|bin| bin.split('/').last().unwrap() == self.verb())
            {
                Some(_) => return true,
                None => return false,
            }
        }

        // Return true if the full path or relative path exists
        true
    }

    /// Run a command
    ///
    /// The [Poem]::recite() function calls this [Verse::incant] function for
    /// each verse it contains. This function handles the actual setup and
    /// spawning (forking) of a new process specified in the [Verse]. It will
    /// also run IO operations for the verse, and setup appropriate coupling,
    /// as per the [Verse]'s own details, contained throughout its fields.
    pub fn incant(
        &mut self,
        out: &mut Vec<u8>,
        pids: &mut Arc<Mutex<Vec<i32>>>,
    ) -> Result<i32, io::Error> {
        // Read files into 'out' if Rune::Read is present in the verse's IO
        if self.io.contains(&Rune::Read) {
            // Enable piping on stdin
            self.couplet += 2;

            // Read all files specified after '<' into 'out', since there may
            // also be piped output from the last command
            for path in self.ip.iter() {
                let mut file = OpenOptions::new().read(true).open(path)?;
                let mut contents = String::new();
                file.read_to_string(&mut contents)?;
                out.append(&mut contents.as_bytes().to_vec());
            }
        }

        // Build the command
        let mut command = Command::new(self.verb());
        command.args(self.clause().unwrap_or(vec![]));

        // Determine couplet status
        if self.couplet == 1 {
            // Verse is the left half of a couplet
            command.stdout(Stdio::piped());
        } else if self.couplet == 2 {
            // Verse is the right half of a couplet
            command.stdin(Stdio::piped());
        } else if self.couplet == 3 {
            // Verse is taking in and piping out output
            command.stdout(Stdio::piped());
            command.stdin(Stdio::piped());
        }

        // Setup for other IO
        if self.io.contains(&Rune::Write) || self.io.contains(&Rune::Addendum) {
            command.stdout(Stdio::piped());
        }
        if self.io.contains(&Rune::Write2) || self.io.contains(&Rune::Addendum2) {
            command.stderr(Stdio::piped());
        }
        if self.io.contains(&Rune::WriteAll) || self.io.contains(&Rune::AddendumAll) {
            command.stdout(Stdio::piped());
            command.stderr(Stdio::piped());
        }

        // Detach the process group, if in the [Rune::Quiet] meter
        if self.meter == Rune::Quiet {
            command.process_group(0);
        }

        // Spawn the process
        let mut child = command.spawn()?;

        // Pipe in command, if we're the right side of a couplet
        if self.couplet > 1 {
            let stdin = child.stdin.as_mut().ok_or(io::ErrorKind::BrokenPipe)?;
            stdin.write_all(&out)?;
            out.clear();
        }

        // Determine what to do based on the meter
        let mut output: Output;
        let mut err: Vec<u8> = Vec::new();
        match self.meter {
            Rune::None | Rune::And | Rune::Continue => {
                output = child.wait_with_output()?;
                if self.io.contains(&Rune::Write) || self.io.contains(&Rune::Addendum) {
                    out.append(&mut output.stdout);
                }
                if self.io.contains(&Rune::Write2) || self.io.contains(&Rune::Addendum2) {
                    err.append(&mut output.stderr);
                }
                if self.io.contains(&Rune::WriteAll) || self.io.contains(&Rune::AddendumAll) {
                    out.append(&mut output.stdout);
                    err.append(&mut output.stderr);
                }
            }
            Rune::Couplet => {
                output = child.wait_with_output()?;
                out.append(&mut output.stdout);
                if self.io.contains(&Rune::Write2) || self.io.contains(&Rune::Addendum2) {
                    err.append(&mut output.stderr);
                }
                if self.io.contains(&Rune::WriteAll) || self.io.contains(&Rune::AddendumAll) {
                    err.append(&mut output.stderr);
                }
            }
            Rune::Quiet => {
                println!("[&]  {}", child.id());

                pids.lock().unwrap().push(child.id() as i32);
                let stanza = self.stanza.join(" ").to_string();
                let pids = Arc::clone(pids);

                unsafe {
                    signal_hook::low_level::register(signal_hook::consts::SIGCHLD, move || {
                        for pid in pids.lock().unwrap().iter() {
                            let mut pid = *pid;
                            let mut status: i32 = 0;
                            pid = waitpid(pid, &mut status, WNOHANG);
                            if pid > 0 {
                                print!("\n[&]  + done    {}", stanza);
                                io::stdout().flush().unwrap();
                            }
                        }
                    })
                    .unwrap();
                }

                return Ok(0);
            }
            _ => unreachable!(),
        }

        // Perform IO operations
        let mut oi = 0;
        let mut ei = 0;
        self.io.retain(|rune| *rune != Rune::Read);
        for io in self.io.iter() {
            let (f, f2) = match *io {
                Rune::Write => {
                    oi += 1;
                    (
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .write(true)
                                .open(&self.op[oi - 1])?,
                        ),
                        None,
                    )
                }
                Rune::Write2 => {
                    ei += 1;
                    (
                        None,
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .write(true)
                                .open(&self.ep[ei - 1])?,
                        ),
                    )
                }
                Rune::WriteAll => {
                    oi += 1;
                    ei += 1;
                    (
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .write(true)
                                .open(&self.op[oi - 1])?,
                        ),
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .write(true)
                                .open(&self.ep[ei - 1])?,
                        ),
                    )
                }
                Rune::Addendum => {
                    oi += 1;
                    (
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .append(true)
                                .open(&self.op[oi - 1])?,
                        ),
                        None,
                    )
                }
                Rune::Addendum2 => {
                    ei += 1;
                    (
                        None,
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .append(true)
                                .open(&self.ep[ei - 1])?,
                        ),
                    )
                }
                Rune::AddendumAll => {
                    oi += 1;
                    ei += 1;
                    (
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .append(true)
                                .open(&self.op[oi - 1])?,
                        ),
                        Some(
                            OpenOptions::new()
                                .create(true)
                                .append(true)
                                .open(&self.ep[ei - 1])?,
                        ),
                    )
                }
                _ => unreachable!(),
            };

            match f {
                Some(mut file) => file.write(out)?,
                None => 0,
            };

            match f2 {
                Some(mut file) => file.write(&err)?,
                None => 0,
            };
        }

        if !output.status.success() {
            return Ok(output.status.code().unwrap_or(-1));
        }

        err.clear();
        if self.meter != Rune::Couplet {
            out.clear();
        }

        Ok(output.status.code().unwrap_or(0))
    }
}