oxo-flow-cli 0.12.0

use anyhow::{Context, Result};
use colored::Colorize;
use oxo_flow_core::backend::ExecutorBackend;
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;

/// Emit the `oxo_submit` shell helper: submits a script and echoes the bare
/// scheduler job id.
///
/// Only PBS's `qsub` prints an id a dependency flag can consume directly.
/// SLURM needs `--parsable`; SGE and LSF print sentences. Chaining raw
/// submit output into `--dependency=afterok:` produced
/// `afterok:Submitted batch job 12345` and broke silently (issue #74
/// phase-1 item 1). The patterns mirror `parse_job_id` in
/// `oxo_flow_core::backend::cluster` so the wrapper and the live submit
/// path agree on what a job id is.
fn generate_submit_helper(backend: &ClusterBackend) -> String {
    let body = match backend {
        // `--parsable` prints a bare id, or `<id>;<cluster>` on federated
        // clusters — keep the id.
        ClusterBackend::Slurm => {
            "  out=$(sbatch --parsable \"$@\") || return $?\n  id=${out%%;*}\n"
        }
        // qsub already prints a bare id (`12345.headnode` is accepted
        // verbatim by `-W depend=`). Exactly one token, or fall through to
        // the error path — mashing multi-token output together would feed a
        // garbage id to the next job's `-W depend=`.
        ClusterBackend::Pbs => {
            "  out=$(qsub \"$@\") || return $?\n  set -- $out\n  if [ $# -eq 1 ]; then id=\"$1\"; else id=\"\"; fi\n"
        }
        // "Your job 12345 (\"align\") has been submitted"
        ClusterBackend::Sge => {
            "  out=$(qsub \"$@\") || return $?\n  id=$(printf '%s' \"$out\" | grep -oE 'Your job(-array)? [0-9]+' | grep -oE '[0-9]+$')\n"
        }
        // "Job <12345> is submitted to queue <normal>."
        ClusterBackend::Lsf => {
            "  out=$(bsub \"$@\") || return $?\n  id=$(printf '%s' \"$out\" | grep -oE 'Job <[0-9]+>' | grep -oE '[0-9]+')\n"
        }
    };

    let mut helper = String::new();
    helper.push_str("# Submit one script and echo its scheduler job id.\n");
    helper.push_str("oxo_submit() {\n");
    helper.push_str("  local out id\n");
    helper.push_str(body);
    helper.push_str("  if [ -z \"$id\" ]; then\n");
    helper.push_str("    echo \"oxo-flow: cannot parse job id from: $out\" >&2\n");
    helper.push_str("    return 1\n");
    helper.push_str("  fi\n");
    helper.push_str("  printf '%s' \"$id\"\n");
    helper.push_str("}\n\n");
    helper
}

/// 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 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");
    script.push_str(&generate_submit_helper(backend));

    // 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[{}]=$(oxo_submit --dependency=afterok:{} {})\n",
                        rule_name,
                        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[{}]=$(oxo_submit -W depend=afterok:{} {})\n",
                        rule_name,
                        dep_str,
                        script_path.display()
                    ));
                }
                ClusterBackend::Sge => {
                    // SGE takes one comma-separated -hold_jid list; a repeated
                    // flag would keep only the last dependency.
                    let hold_jid = dep_job_refs.join(",");
                    script.push_str(&format!(
                        "JOB_IDS[{}]=$(oxo_submit -hold_jid {} {})\n",
                        rule_name,
                        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[{}]=$(oxo_submit -w '{}' {})\n",
                        rule_name,
                        dep_str,
                        script_path.display()
                    ));
                }
            }
        } else {
            // No dependencies
            script.push_str(&format!(
                "JOB_IDS[{}]=$(oxo_submit {})\n",
                rule_name,
                script_path.display()
            ));
        }

        // Double quotes: the id has to expand, not print literally.
        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,
            walltime,
            extra_args,
            output,
            target,
            dry_run,
            with_dependencies,
        } => {
            let mut config = WorkflowConfig::from_file(&workflow)
                .with_context(|| format!("failed to parse {}", workflow.display()))?;

            // Expand wildcards before the DAG is built, exactly as `run`
            // does (issue #74 phase 1). Without this a scatter rule stays a
            // single template and the generated script submits a literal
            // `{sample}` to the scheduler.
            config.apply_defaults();
            config
                .expand_wildcards()
                .context("failed to expand wildcard rules")?;

            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: walltime.clone(),
                extra_args: extra_args.clone(),
            };

            if dry_run {
                eprintln!(
                    "{} (dry-run) would generate {} job scripts for {} rule instances",
                    "Cluster:".bold().yellow(),
                    backend,
                    order.len()
                );
                return Ok(());
            }

            std::fs::create_dir_all(&output)?;

            // The rendered scripts declare `#SBATCH --output=logs/<rule>.out`
            // and slurmd opens that file at job launch — before the script
            // body's `mkdir -p logs` runs. Create the directory now (issue #74
            // phase-1 note 2).
            if let Some(wf_dir) = workflow.parent() {
                std::fs::create_dir_all(wf_dir.join("logs"))?;
            }

            eprintln!(
                "{} Generating {} job scripts for {} rule instances",
                "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;
                    }
                };

                // Render through the ExecutorBackend trait (issue #78): the
                // command stays a thin render layer over the same directive
                // generator the live submit path uses.
                let wrapped_cmd = env_resolver
                    .wrap_command(
                        &shell_cmd,
                        &rule.environment,
                        Some(&rule.resources),
                        Path::new("."),
                    )
                    .map_err(|e| anyhow::anyhow!("environment wrapping failed: {}", e))?;
                let scheduled = oxo_flow_core::backend::ScheduledRule {
                    rule: rule.clone(),
                    shell_cmd: wrapped_cmd,
                    workdir: std::path::PathBuf::from("."),
                    dependencies: dag.dependencies(rule_name).unwrap_or_default(),
                    wildcard_values: wildcard_values.clone(),
                };
                let executor = oxo_flow_core::backend::cluster::ClusterExecutor::new(
                    cluster_backend,
                    cluster_config.clone(),
                );
                let script = executor.render_script(&scheduled)?;

                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 } => {
            // Issue #67 §4: the last CLI stub. SLURM fetches a precise
            // accounting record (`sacct --format=JobID,State,ExitCode,
            // Elapsed,MaxRSS`); PBS/SGE/LSF stay best-effort (qstat -f /
            // qacct / bacct) — the same per-scheduler contract the
            // BackendDriver uses.
            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 executor = oxo_flow_core::backend::cluster::ClusterExecutor::new(
                cluster_backend,
                oxo_flow_core::cluster::ClusterJobConfig {
                    backend: cluster_backend,
                    queue: None,
                    account: None,
                    walltime: None,
                    extra_args: Vec::new(),
                },
            );
            let logs = executor
                .logs(&job_id)
                .await
                .context("failed to fetch job logs")?;
            if logs.trim().is_empty() {
                eprintln!(
                    "{} No accounting records found for job ID {}",
                    "Warning:".bold().yellow(),
                    job_id
                );
            } else {
                println!("{logs}");
            }
        }
    }
    Ok(())
}