oxo-flow-cli 0.9.4

use anyhow::{Context, Result};
use colored::Colorize;
use oxo_flow_core::cluster::ClusterBackend;
use oxo_flow_core::config::WorkflowConfig;
use oxo_flow_core::dag::WorkflowDag;
use std::collections::HashMap;
use std::path::Path;

use crate::ClusterAction;
use crate::commands::print_banner;

/// Generate a submit wrapper script that handles job dependencies.
/// This script tracks job IDs and sets up proper dependency chains.
fn generate_submit_wrapper(
    backend: &ClusterBackend,
    order: &[String],
    dag: &WorkflowDag,
    output_dir: &Path,
) -> Result<String> {
    let submit_cmd = match backend {
        ClusterBackend::Slurm => "sbatch",
        ClusterBackend::Pbs => "qsub",
        ClusterBackend::Sge => "qsub",
        ClusterBackend::Lsf => "bsub",
    };

    let mut script = String::new();
    script.push_str("#!/bin/bash\n");
    script.push_str("# Auto-generated dependency-aware submit script\n");
    script.push_str("# Generated by oxo-flow\n\n");
    script.push_str("set -e\n\n");
    script.push_str("# Track job IDs\ndeclare -A JOB_IDS\n\n");

    // Generate submit commands for each rule in order
    for rule_name in order {
        let script_name = format!("{}.sh", rule_name);
        let script_path = output_dir.join(&script_name);

        // Get dependencies for this rule
        let deps = dag.dependencies(rule_name).unwrap_or_default();
        let dep_job_refs: Vec<String> = deps
            .iter()
            .map(|d| format!("${{JOB_IDS[{}]}}", d))
            .collect();

        script.push_str(&format!("echo 'Submitting {}...'\n", rule_name));

        // Add dependency specification if there are dependencies
        if !dep_job_refs.is_empty() {
            match backend {
                ClusterBackend::Slurm => {
                    let dep_str = dep_job_refs.join(":");
                    script.push_str(&format!(
                        "JOB_IDS[{}]=$({} --dependency=afterok:{} {})\n",
                        rule_name,
                        submit_cmd,
                        dep_str,
                        script_path.display()
                    ));
                }
                ClusterBackend::Pbs => {
                    // PBS uses -W depend=afterok:jobid
                    let dep_str = dep_job_refs.join(":");
                    script.push_str(&format!(
                        "JOB_IDS[{}]=$({} -W depend=afterok:{} {})\n",
                        rule_name,
                        submit_cmd,
                        dep_str,
                        script_path.display()
                    ));
                }
                ClusterBackend::Sge => {
                    // SGE uses -hold_jid jobid
                    let hold_jid = deps
                        .iter()
                        .map(|d| format!("-hold_jid ${{JOB_IDS[{}]}}", d))
                        .collect::<Vec<_>>()
                        .join(" ");
                    script.push_str(&format!(
                        "JOB_IDS[{}]=$({} {} {})\n",
                        rule_name,
                        submit_cmd,
                        hold_jid,
                        script_path.display()
                    ));
                }
                ClusterBackend::Lsf => {
                    // LSF uses -w 'ended(jobid)'
                    let dep_str = dep_job_refs
                        .iter()
                        .map(|d| format!("ended({})", d))
                        .collect::<Vec<_>>()
                        .join(" && ");
                    script.push_str(&format!(
                        "JOB_IDS[{}]=$({} -w '{}') {})\n",
                        rule_name,
                        submit_cmd,
                        dep_str,
                        script_path.display()
                    ));
                }
            }
        } else {
            // No dependencies
            script.push_str(&format!(
                "JOB_IDS[{}]=$({} {})\n",
                rule_name,
                submit_cmd,
                script_path.display()
            ));
        }

        script.push_str(&format!(
            "echo '  Submitted {} as job ID: ${{JOB_IDS[{}]}}'\n\n",
            rule_name, rule_name
        ));
    }

    script.push_str("echo 'All jobs submitted successfully!'\n");
    script.push_str("echo 'Job ID mapping:'\n");
    script.push_str("for name in \"${!JOB_IDS[@]}\"; do\n");
    script.push_str("  echo \"  $name: ${JOB_IDS[$name]}\"\n");
    script.push_str("done\n");

    Ok(script)
}

pub async fn cluster_command(action: ClusterAction) -> Result<()> {
    print_banner();
    match action {
        ClusterAction::Submit {
            workflow,
            backend,
            queue,
            account,
            output,
            target,
            dry_run,
            with_dependencies,
        } => {
            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 = if target.is_empty() {
                dag.execution_order()?
            } else {
                let target_refs: Vec<&str> = target.iter().map(String::as_str).collect();
                dag.execution_order_for_targets(&target_refs)
                    .with_context(|| "failed to resolve target rules")?
            };

            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()
            );

            // Create environment resolver for command wrapping
            let env_resolver = oxo_flow_core::environment::EnvironmentResolver::new();

            // Build config variable map for placeholder expansion
            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)
                    .ok_or_else(|| anyhow::anyhow!("rule '{}' not found in workflow", rule_name))?;

                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;
                    }
                };

                // Generate script with environment wrapping
                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());
            }

            // Generate dependency-aware submit script if requested
            if with_dependencies {
                let submit_script =
                    generate_submit_wrapper(&cluster_backend, &order, &dag, &output)?;
                let submit_path = output.join("submit.sh");
                std::fs::write(&submit_path, submit_script)?;
                eprintln!(
                    "  {} {} (dependency-aware submit script)",
                    "✓".green(),
                    submit_path.display()
                );

                eprintln!(
                    "\n{} {} scripts written to {}",
                    "Done:".bold(),
                    order.len() + 1,
                    output.display()
                );
                eprintln!("  Submit with: bash {}", submit_path.display());
                eprintln!(
                    "  Or manually: {} {}/*.sh",
                    oxo_flow_core::cluster::submit_command(&cluster_backend),
                    output.display()
                );
            } else {
                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(())
}