use anyhow::{Context, Result};
use colored::Colorize;
use oxo_flow_core::config::WorkflowConfig;
use oxo_flow_core::dag::WorkflowDag;
use std::collections::HashMap;
use crate::ClusterAction;
use crate::commands::print_banner;
pub async fn cluster_command(action: ClusterAction) -> Result<()> {
print_banner();
match action {
ClusterAction::Submit {
workflow,
backend,
queue,
account,
output,
target: _,
dry_run,
} => {
let config = WorkflowConfig::from_file(&workflow)
.with_context(|| format!("failed to parse {}", workflow.display()))?;
let dag =
WorkflowDag::from_rules(&config.rules).context("failed to build workflow DAG")?;
let order = dag.execution_order()?;
let cluster_backend = match backend.as_str() {
"pbs" => oxo_flow_core::cluster::ClusterBackend::Pbs,
"sge" => oxo_flow_core::cluster::ClusterBackend::Sge,
"lsf" => oxo_flow_core::cluster::ClusterBackend::Lsf,
_ => oxo_flow_core::cluster::ClusterBackend::Slurm,
};
let cluster_config = oxo_flow_core::cluster::ClusterJobConfig {
backend: cluster_backend,
queue: queue.clone(),
account: account.clone(),
walltime: None,
extra_args: vec![],
};
if dry_run {
eprintln!(
"{} (dry-run) would generate {} job scripts for {} rules",
"Cluster:".bold().yellow(),
backend,
order.len()
);
return Ok(());
}
std::fs::create_dir_all(&output)?;
eprintln!(
"{} Generating {} job scripts for {} rules",
"Cluster:".bold().cyan(),
backend,
order.len()
);
let env_resolver = oxo_flow_core::environment::EnvironmentResolver::new();
let mut wildcard_values: HashMap<String, String> = HashMap::new();
for (key, value) in &config.config {
let string_val = match value {
toml::Value::String(s) => s.clone(),
other => other.to_string(),
};
wildcard_values.insert(format!("config.{key}"), string_val);
}
for rule_name in &order {
let rule = config.get_rule(rule_name).unwrap();
let shell_cmd = match oxo_flow_core::executor::process::build_execution_command(
rule,
&wildcard_values,
&config.workflow.interpreter_map,
) {
Some(cmd) => cmd,
None => {
eprintln!(
" {} {} — no shell command or script, skipping",
"⊘".yellow(),
rule_name
);
continue;
}
};
let script = oxo_flow_core::cluster::generate_submit_script_with_env(
&cluster_backend,
rule,
&shell_cmd,
&cluster_config,
&env_resolver,
)
.map_err(|e| anyhow::anyhow!("environment wrapping failed: {}", e))?;
let script_path = output.join(format!("{rule_name}.sh"));
std::fs::write(&script_path, &script)?;
eprintln!(" {} {}", "✓".green(), script_path.display());
}
eprintln!(
"\n{} {} scripts written to {}",
"Done:".bold(),
order.len(),
output.display()
);
eprintln!(
" Submit with: {} {}/*.sh",
oxo_flow_core::cluster::submit_command(&cluster_backend),
output.display()
);
}
ClusterAction::Status { backend, job_ids } => {
let cluster_backend = match backend.as_str() {
"pbs" => oxo_flow_core::cluster::ClusterBackend::Pbs,
"sge" => oxo_flow_core::cluster::ClusterBackend::Sge,
"lsf" => oxo_flow_core::cluster::ClusterBackend::Lsf,
_ => oxo_flow_core::cluster::ClusterBackend::Slurm,
};
let status_cmd = oxo_flow_core::cluster::status_command(&cluster_backend);
eprintln!("{} Executing '{}'...", "Cluster:".bold().cyan(), status_cmd);
let mut parts = status_cmd.split_whitespace();
let program = parts.next().unwrap_or(status_cmd);
let mut args: Vec<&str> = parts.collect();
for id in &job_ids {
args.push(id);
}
match std::process::Command::new(program).args(&args).status() {
Ok(status) => {
if !status.success() {
anyhow::bail!(
"Command failed with exit code: {}",
status.code().unwrap_or(-1)
);
}
}
Err(e) => {
eprintln!(" Is {} installed on this system?", program);
anyhow::bail!("Failed to execute status command: {}", e);
}
}
}
ClusterAction::Cancel { backend, job_ids } => {
let cancel_cmd = match backend.as_str() {
"pbs" => "qdel",
"sge" => "qdel",
"lsf" => "bkill",
_ => "scancel",
};
if job_ids.is_empty() {
eprintln!(
"{} No job IDs provided. Usage: oxo-flow cluster cancel <JOB_ID>...",
"Warning:".bold().yellow()
);
} else {
eprintln!(
"{} Canceling {} job(s)...",
"Cluster:".bold().cyan(),
job_ids.len()
);
match std::process::Command::new(cancel_cmd)
.args(&job_ids)
.status()
{
Ok(status) => {
if status.success() {
eprintln!("{} Successfully canceled jobs.", "✓".green());
} else {
anyhow::bail!(
"Command failed with exit code: {}",
status.code().unwrap_or(-1)
);
}
}
Err(e) => {
eprintln!(" Is {} installed on this system?", cancel_cmd);
anyhow::bail!("Failed to execute cancel command: {}", e);
}
}
}
}
ClusterAction::Logs { backend: _, job_id } => {
eprintln!(
"{} Logs for job ID {} not yet implemented",
"⚠️".yellow(),
job_id
);
}
}
Ok(())
}