oxo-flow-cli 0.29.1

use anyhow::{Context, Result};
use colored::Colorize;
use indicatif::{ProgressBar, ProgressStyle};
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::Semaphore;

use crate::commands::{
    collect_batch_items, expand_batch_template, print_banner, run_batch_command, wrap_batch_command,
};

/// Result for a single batch item.
#[derive(Debug, Clone, serde::Serialize)]
pub struct BatchResult {
    pub item: String,
    pub command: String,
    pub exit_code: Option<i32>,
    pub success: bool,
    pub error: Option<String>,
}

#[allow(clippy::too_many_arguments)]
pub async fn batch_command(
    template: String,
    items: Vec<String>,
    jobs: usize,
    stop_on_error: bool,
    file: Option<PathBuf>,
    json_output: bool,
    dry_run: bool,
    workdir: Option<PathBuf>,
    environment: Option<String>,
    generate_workflow: bool,
    output: Option<PathBuf>,
) -> Result<()> {
    print_banner();

    // `-j 0` means "no explicit concurrency limit": clamp once at the
    // boundary so the semaphore below is never a zero-permit gate that
    // hangs the first item forever (`run` clamps identically, issue #136
    // fix 1; audit finding: `batch -j 0` hung).
    let jobs = jobs.max(1);

    // Collect items from command line or file (validates non-empty)
    let items = collect_batch_items(&items, file.as_ref())?;

    // Generate workflow instead of executing
    if generate_workflow {
        return generate_batch_workflow(&template, &items, output.as_ref(), environment.as_ref());
    }
    if output.is_some() {
        // The flag is consumed only by the --generate-workflow branch —
        // silently ignoring it made callers wait on a results file that is
        // never written (audit #670).
        anyhow::bail!(
            "-o/--output only applies to --generate-workflow; batch results print to stdout"
        );
    }

    // Dry-run: print expanded commands. The preview IS the result in
    // dry-run — with --json-output (or the global --json) it is emitted
    // as JSON on stdout so scripts can consume the previewed commands
    // without executing anything (issue #864). The human-readable preview
    // stays on stderr in both modes, so piping stdout never mixes in
    // non-JSON bytes.
    if dry_run {
        let expanded: Vec<BatchResult> = items
            .iter()
            .enumerate()
            .map(|(nr, item)| {
                let cmd = expand_batch_template(&template, item, nr + 1);
                BatchResult {
                    item: item.clone(),
                    command: cmd,
                    exit_code: None,
                    success: true,
                    error: None,
                }
            })
            .collect();
        eprintln!(
            "{} {} items with template: {}",
            "Batch (dry-run):".bold().yellow(),
            items.len(),
            template
        );
        for (nr, r) in expanded.iter().enumerate() {
            eprintln!("  [{}] {}", nr + 1, r.command);
        }
        if json_output {
            let output = serde_json::json!({
                "tool": "oxo-flow batch",
                "dry_run": true,
                "template": template,
                "total": expanded.len(),
                "failed": 0,
                "success": true,
                "results": expanded,
            });
            println!("{}", serde_json::to_string_pretty(&output)?);
        }
        return Ok(());
    }

    let semaphore = Arc::new(Semaphore::new(jobs));
    let results = Arc::new(Mutex::new(Vec::new()));
    let stop_flag = Arc::new(AtomicBool::new(false));

    let progress = ProgressBar::new(items.len() as u64);
    let style = ProgressStyle::default_bar()
        .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} ({msg})")?
        .progress_chars("#>-");
    progress.set_style(style);

    let mut tasks = Vec::new();

    let env_clone = environment.clone();
    let workdir_clone = workdir
        .clone()
        .unwrap_or_else(|| std::env::current_dir().unwrap_or_default());

    for (nr, item) in items.iter().enumerate() {
        if stop_flag.load(Ordering::Relaxed) {
            break;
        }

        let permit = semaphore
            .clone()
            .acquire_owned()
            .await
            // Safe: the semaphore is never closed; acquire only fails on a
            // closed semaphore, and we hold Arc clones for the whole run.
            .expect("batch semaphore closed unexpectedly");
        let cmd = expand_batch_template(&template, item, nr + 1);

        let item_clone = item.clone();
        let cmd_clone = cmd.clone();
        let results_clone = results.clone();
        let stop_flag_clone = stop_flag.clone();
        let progress_clone = progress.clone();
        let env_clone = env_clone.clone();
        let workdir_clone = workdir_clone.clone();

        progress.set_message(format!("item {}", nr + 1));

        if stop_flag_clone.load(Ordering::Relaxed) {
            progress_clone.inc(1);
            if stop_on_error {
                stop_flag_clone.store(true, Ordering::Relaxed);
            }
            // Record the skip so the summary reports the full item count
            // (a skipped item is neither done nor a failure).
            if let Ok(mut results_lock) = results_clone.lock() {
                results_lock.push(BatchResult {
                    item: item_clone.clone(),
                    command: cmd_clone,
                    exit_code: None,
                    success: false,
                    error: Some("skipped (stop-on-error)".to_string()),
                });
            }
            // Release permit and continue
            drop(permit);
            continue;
        }

        // Wrap with environment if specified
        let final_cmd = if let Some(env) = &env_clone {
            wrap_batch_command(&cmd, env)
        } else {
            cmd.clone()
        };

        let task = tokio::spawn(async move {
            // Clone for spawn_blocking (needs 'static lifetime)
            let final_cmd_for_blocking = final_cmd.clone();
            let workdir_for_blocking = workdir_clone.clone();

            // Execute command in blocking context
            let result = tokio::task::spawn_blocking(move || {
                run_batch_command(&final_cmd_for_blocking, &workdir_for_blocking)
            })
            .await;

            // Join failure means the blocking thread died (panic/pressure) —
            // record it as this item's error instead of panicking the task.
            let exit_result = match result {
                Ok(r) => r,
                Err(join_err) => Err(anyhow::anyhow!("batch worker crashed: {join_err}")),
            };

            // Record result
            if let Ok(mut results_lock) = results_clone.lock() {
                match exit_result {
                    Ok(exit_code) => {
                        results_lock.push(BatchResult {
                            item: item_clone.clone(),
                            command: cmd_clone,
                            exit_code: Some(exit_code),
                            success: exit_code == 0,
                            error: None,
                        });
                        if exit_code != 0 && stop_on_error {
                            stop_flag_clone.store(true, Ordering::Relaxed);
                        }
                    }
                    Err(e) => {
                        results_lock.push(BatchResult {
                            item: item_clone.clone(),
                            command: cmd_clone,
                            exit_code: None,
                            success: false,
                            error: Some(e.to_string()),
                        });
                        if stop_on_error {
                            stop_flag_clone.store(true, Ordering::Relaxed);
                        }
                    }
                }
            }
            progress_clone.inc(1);

            // Release permit
            drop(permit);
        });

        tasks.push(task);
    }

    // Wait for all tasks to complete
    for task in tasks {
        if let Err(e) = task.await {
            tracing::error!("task failed: {}", e);
        }
    }

    progress.finish_and_clear();

    // Calculate summary. Items skipped by stop-on-error are recorded, but
    // items past the loop-top `break` are not — `total_results` is the
    // number of *recorded* items (executed + skipped), which may be smaller
    // than `items.len()` when the batch aborted early.
    let results_lock = results
        .lock()
        .map_err(|_| anyhow::anyhow!("results mutex poisoned"))?;
    let total_results = results_lock.len();
    let failed = results_lock.iter().filter(|r| !r.success).count();
    let skipped = results_lock
        .iter()
        .filter(|r| r.error.as_deref() == Some("skipped (stop-on-error)"))
        .count();
    let done = total_results - failed;

    // Output summary
    if json_output {
        let output = serde_json::json!({
            "tool": "oxo-flow batch",
            "template": template,
            "total": total_results,
            "failed": failed - skipped,
            "skipped": skipped,
            "success": failed == 0,
            "results": results_lock.clone()
        });
        println!("{}", serde_json::to_string_pretty(&output)?);
    } else {
        eprintln!(
            "\n{} {}/{} items succeeded, {} failed, {} skipped",
            "Done:".bold(),
            done,
            total_results,
            failed - skipped,
            skipped
        );
        if failed > 0 {
            eprintln!("\n{}", "Failures:".bold().red());
            for r in results_lock
                .iter()
                .filter(|r| !r.success && r.error.as_deref() != Some("skipped (stop-on-error)"))
            {
                eprintln!(
                    "  {} — {}",
                    r.item.bold(),
                    r.error.as_deref().unwrap_or("exit code non-zero")
                );
            }
        }
    }

    // #540: any failed item makes the batch fail — matching every
    // job-runner analog (xargs -P, make -k, Snakemake): a failed item must
    // be visible to scripts through the exit code, not just the printed
    // JSON. --stop-on-error stays the abort-on-first-failure switch, not
    // the only failure signal.
    if failed > 0 {
        let failed_only = failed - skipped;
        let skipped_note = if skipped > 0 {
            format!(", {skipped} skipped after stop-on-error")
        } else {
            String::new()
        };
        return Err(anyhow::anyhow!(
            "batch execution failed: {failed_only} item(s) did not succeed{skipped_note}"
        ));
    }

    Ok(())
}

/// Render `shell = <cmd>` as a TOML table line (with a trailing blank line).
/// Serializing through the toml crate handles backslashes, quotes and
/// newlines correctly, unlike manual escaping which produced invalid TOML for
/// commands containing literal `\n` or trailing `\`.
fn format_shell_line(cmd: &str) -> Result<String> {
    let mut table = toml::map::Map::new();
    table.insert("shell".to_string(), toml::Value::String(cmd.to_string()));
    let rendered = toml::to_string(&toml::Value::Table(table))
        .with_context(|| format!("serialize shell command {cmd:?} as TOML"))?;
    // rendered ends with one newline from the serializer.
    Ok(format!("{rendered}\n"))
}

/// Map a `-e/--environment` spec (`conda: envs/x.yaml`, `docker: image`)
/// onto the `[environment]` TABLE form `EnvironmentSpec` actually
/// deserializes. The generator used to emit the string form verbatim, which
/// `validate` rejected (`invalid type: string, expected struct
/// EnvironmentSpec`) — every generated workflow was dead on arrival.
/// Multi-token specs (`docker: --gpus all image`) and bare environment
/// names have no faithful table form, so they are refused with a pointer
/// at the manual edit instead of writing a broken file.
fn environment_table(env: &str) -> Result<(&str, String)> {
    let Some((kind, spec)) = env.split_once(':') else {
        return Err(anyhow::anyhow!(
            "bare environment name '{env}' cannot map to a workflow [environment] table — \
             point at an env file instead (e.g. -e \"conda: envs/env.yaml\")"
        ));
    };
    let kind = match kind.trim() {
        "conda" | "mamba" | "pixi" | "docker" | "singularity" | "venv" => kind.trim(),
        other => {
            return Err(anyhow::anyhow!(
                "unsupported environment kind '{other}:' — expected conda/mamba/pixi/docker/singularity/venv"
            ));
        }
    };
    let spec = spec.trim();
    if spec.is_empty() || spec.contains(char::is_whitespace) {
        return Err(anyhow::anyhow!(
            "environment spec '{env}' carries flags or extra tokens that have no [environment] \
             table form — generate the workflow and add the flags by hand \
             (e.g. environment = {{ docker = \"image\", gpus = \"all\" }})"
        ));
    }
    Ok((kind, spec.to_string()))
}

pub fn generate_batch_workflow(
    template: &str,
    items: &[String],
    output_path: Option<&PathBuf>,
    environment: Option<&String>,
) -> Result<()> {
    let output_path = output_path
        .cloned()
        .unwrap_or_else(|| PathBuf::from("batch.oxoflow"));

    let env_table = environment
        .as_ref()
        .map(|e| environment_table(e))
        .transpose()?;

    let mut content = String::new();
    content.push_str("[workflow]\n");
    content.push_str("name = \"batch_workflow\"\n");
    content.push_str("version = \"0.1.0\"\n\n");

    for (i, item) in items.iter().enumerate() {
        let cmd = expand_batch_template(template, item, i + 1);
        content.push_str("[[rules]]\n");
        content.push_str(&format!("name = \"batch_item_{}\"\n", i + 1));
        // Inline table form: a bare `[environment]` section header per
        // [[rules]] element is rejected by the parser as a duplicate key,
        // and the inline form keeps each rule self-contained. The value is
        // a single token (environment_table refuses whitespace) rendered
        // as an escaped TOML basic string.
        if let Some((kind, spec)) = &env_table {
            content.push_str(&format!(
                "environment = {{ {} = \"{}\" }}\n",
                kind,
                escape_toml_string(spec)
            ));
        }
        content.push_str(&format_shell_line(&cmd)?);
    }

    // Validate with the real parser before writing — the same guard
    // `oxo-flow import` applies to its transpiler output, so a generator
    // regression can never again produce a file that fails `validate`.
    oxo_flow_core::config::WorkflowConfig::parse(&content)
        .with_context(|| "generated batch workflow failed oxo-flow parse validation")?;

    std::fs::write(&output_path, content)
        .with_context(|| format!("failed to write {}", output_path.display()))?;

    eprintln!(
        "{} Generated batch workflow: {}",
        "✓".green().bold(),
        output_path.display()
    );
    Ok(())
}

/// Escape a value for a TOML basic string (the generated `environment`
/// values are paths/image refs; quotes and backslashes are the real risks).
fn escape_toml_string(value: &str) -> String {
    value.replace('\\', "\\\\").replace('"', "\\\"")
}

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

    #[test]
    fn shell_line_roundtrips_hostile_commands() {
        for cmd in [
            "echo 'single'",
            "awk '{print $1}'",
            "printf \"tab\\there\"",
            "echo back\\slash",
            "trailing\\",
            "multi\nline\ncmd",
            "quote \" in middle",
        ] {
            let line = format_shell_line(cmd).expect("serialization");
            let doc = format!("[r]\n{line}");
            let parsed: toml::Value =
                toml::from_str(&doc).unwrap_or_else(|e| panic!("invalid TOML for {cmd:?}: {e}"));
            assert_eq!(
                parsed["r"]["shell"].as_str().unwrap(),
                cmd,
                "roundtrip mismatch for {cmd:?}"
            );
        }
    }

    #[test]
    fn generated_workflow_parses_with_hostile_template() {
        let out = std::env::temp_dir().join(format!("batch-test-{}.toml", std::process::id()));
        generate_batch_workflow(
            "echo {}",
            &["we\"ird".to_string(), "back\\slash".to_string()],
            Some(&out),
            None,
        )
        .unwrap();
        let text = std::fs::read_to_string(&out).unwrap();
        let _parsed: toml::Value = toml::from_str(&text)
            .unwrap_or_else(|e| panic!("generated file must be valid TOML: {e}\n{text}"));
        std::fs::remove_file(&out).ok();
    }

    #[test]
    fn generated_workflow_environment_uses_table_form() {
        // The string form (`environment = "conda: …"`) fails
        // `EnvironmentSpec` deserialization — the generated file must carry
        // the table form and parse with the real workflow parser.
        let out = std::env::temp_dir().join(format!("batch-env-{}.toml", std::process::id()));
        generate_batch_workflow(
            "echo {}",
            &["a.fastq".to_string(), "b.fastq".to_string()],
            Some(&out),
            Some(&"conda: envs/tools.yaml".to_string()),
        )
        .unwrap();
        let text = std::fs::read_to_string(&out).unwrap();
        let parsed: toml::Value = toml::from_str(&text)
            .unwrap_or_else(|e| panic!("generated file must be valid TOML: {e}\n{text}"));
        for rule in parsed["rules"].as_array().expect("rules array") {
            assert_eq!(
                rule["environment"]["conda"],
                toml::Value::String("envs/tools.yaml".to_string()),
                "environment must be a table with the conda path\n{text}"
            );
        }
        std::fs::remove_file(&out).ok();
    }

    #[test]
    fn environment_specs_without_table_form_are_refused() {
        // Multi-token docker specs and bare names cannot map onto
        // EnvironmentSpec — refuse at generation time instead of writing a
        // broken file.
        let err = environment_table("docker: --gpus all image").unwrap_err();
        assert!(err.to_string().contains("no [environment] table form"));
        let err = environment_table("myenv").unwrap_err();
        assert!(err.to_string().contains("conda: envs/env.yaml"));
        let err = environment_table("python: x").unwrap_err();
        assert!(err.to_string().contains("unsupported environment kind"));
        // Single-token known kinds map cleanly.
        assert_eq!(
            environment_table("docker: biocontainers/bwa:0.7.17").unwrap(),
            ("docker", "biocontainers/bwa:0.7.17".to_string())
        );
    }
}