oxo-flow-cli 0.21.0

use anyhow::{Context, Result};
use colored::Colorize;
use std::path::{Path, PathBuf};

use crate::commands::compute_sha256;

/// Bundle a workflow with its referenced environment files into a verifiable archive.
///
/// Reads the .oxoflow workflow file, follows `[[include]]` references to discover
/// all environment spec files, collects `scripts/` and `bin/` directories, and
/// produces a single `.tar.zst` archive with a complete, checksum-verified manifest.
///
/// With `--with-lockfiles`, generates deterministic conda lockfiles for each
/// conda/mamba environment YAML, ensuring exact reproducibility across time.
pub fn publish_command(
    workflow: PathBuf,
    output: Option<PathBuf>,
    with_lockfiles: bool,
    format: Option<String>,
) -> Result<()> {
    let bundle_format = match format.as_deref() {
        Some(f) => crate::commands::bundle::BundleFormat::parse(f)?,
        None => crate::commands::bundle::BundleFormat::TarZst,
    };
    let workflow_path =
        std::path::absolute(&workflow).context("failed to resolve workflow path")?;
    let workflow_dir = workflow_path.parent().unwrap_or(Path::new("."));

    let workflow_name = workflow_path
        .file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("workflow");

    let output_archive = if let Some(out) = output {
        if out.extension().is_none() {
            PathBuf::from(format!("{}.{}", out.display(), bundle_format.extension()))
        } else {
            out
        }
    } else {
        PathBuf::from(format!(
            "{}-bundle.{}",
            workflow_name,
            bundle_format.extension()
        ))
    };

    // ── Collect all referenced files ──────────────────────────────────────

    let mut referenced_files: Vec<(String, PathBuf)> = Vec::new();
    let mut container_refs: Vec<serde_json::Value> = Vec::new();
    let mut scanned_workflows: std::collections::HashSet<PathBuf> =
        std::collections::HashSet::new();

    // Scan the main workflow and all included sub-workflows recursively.
    scan_workflow_env_files(
        &workflow_path,
        workflow_dir,
        &mut referenced_files,
        &mut container_refs,
        &mut scanned_workflows,
    )?;

    // ── Generate conda lockfiles (if --with-lockfiles) ────────────────────

    // The returned TempDir guard is held by the caller so the generated
    // lockfiles stay on disk until the staging copy below reads them.
    let _lock_staging = if with_lockfiles {
        generate_lockfiles(workflow_dir, &mut referenced_files)
    } else {
        None
    };

    // Collect scripts/ and bin/ directories if they exist (Nextflow-style auto-PATH convention)
    for dir_name in &["scripts", "bin"] {
        let dir_path = workflow_dir.join(dir_name);
        if dir_path.is_dir() {
            collect_directory_files(&dir_path, dir_name, &mut referenced_files)?;
        }
    }

    // ── Build manifest with checksums ─────────────────────────────────────

    let oxo_version = env!("CARGO_PKG_VERSION").to_string();
    let mut manifest_files = Vec::new();
    // Staged into a uniquely named temp directory instead of
    // `temp_dir()/oxo-publish-<pid>` (see the extraction note in
    // `bundle.rs` for why the predictable pattern is unsafe); the TempDir
    // guard cleans it up on every exit path, replacing the manual
    // `remove_dir_all` that error paths previously skipped.
    let temp_dir = tempfile::Builder::new()
        .prefix("oxo-publish-")
        .tempdir()
        .context("failed to create temporary staging directory")?;
    let temp_dir = temp_dir.path().to_path_buf();

    // Copy the main workflow file
    let wf_filename = workflow_path
        .file_name()
        .and_then(|n| n.to_str())
        .unwrap_or("workflow.oxoflow");
    let wf_dest = temp_dir.join(wf_filename);
    std::fs::copy(&workflow_path, &wf_dest)?;
    let wf_checksum = compute_sha256(&workflow_path)?;
    let wf_size = std::fs::metadata(&workflow_path)?.len();
    manifest_files.push(serde_json::json!({
        "path": wf_filename,
        "sha256": wf_checksum,
        "size": wf_size,
    }));

    // Copy all env/script files to temp dir and compute checksums
    for (rel_path, abs_path) in &referenced_files {
        ensure_safe_bundle_path(rel_path)?;
        let checksum = compute_sha256(abs_path)?;
        let dest = temp_dir.join(rel_path);
        if let Some(parent) = dest.parent() {
            std::fs::create_dir_all(parent)?;
        }
        std::fs::copy(abs_path, &dest)?;

        let size = std::fs::metadata(abs_path)?.len();
        manifest_files.push(serde_json::json!({
            "path": rel_path,
            "sha256": checksum,
            "size": size,
        }));
    }

    // ── Build per-rule resource summary for the manifest ──────────────────
    // Helps bundle consumers assess hardware requirements before running.
    let config = oxo_flow_core::config::WorkflowConfig::from_file(&workflow_path)
        .with_context(|| format!("failed to parse {}", workflow_path.display()))?;
    let mut resource_summary = Vec::new();
    for rule in &config.rules {
        let mut entry = serde_json::json!({
            "rule": rule.name,
            "threads": rule.effective_threads(),
        });
        if let Some(mem) = rule.effective_memory() {
            entry["memory"] = serde_json::Value::String(mem.to_string());
        }
        if rule.resources.gpu.is_some() || rule.resources.gpu_spec.is_some() {
            let gpu = rule.resources.gpu.unwrap_or(0)
                + rule
                    .resources
                    .gpu_spec
                    .as_ref()
                    .map(|s| s.count)
                    .unwrap_or(0);
            if gpu > 0 {
                entry["gpu"] = serde_json::Value::Number(serde_json::Number::from(gpu));
            }
        }
        if let Some(ref disk) = rule.resources.disk {
            entry["disk"] = serde_json::Value::String(disk.clone());
        }
        if let Some(ref time_limit) = rule.resources.time_limit {
            entry["time_limit"] = serde_json::Value::String(time_limit.clone());
        }
        resource_summary.push(entry);
    }

    // Compute aggregate minimum requirements
    let max_threads = config
        .rules
        .iter()
        .map(|r| r.effective_threads())
        .max()
        .unwrap_or(1);
    let max_memory = config
        .rules
        .iter()
        .filter_map(|r: &oxo_flow_core::rule::Rule| {
            r.effective_memory()
                .and_then(oxo_flow_core::scheduler::parse_memory_mb)
        })
        .max();
    let total_gpu = config
        .rules
        .iter()
        .filter_map(|r| {
            let simple = r.resources.gpu.unwrap_or(0);
            let spec = r.resources.gpu_spec.as_ref().map(|s| s.count).unwrap_or(0);
            let total = simple + spec;
            if total > 0 { Some(total) } else { None }
        })
        .max()
        .unwrap_or(0);

    let mut recommendations = serde_json::json!({
        "min_threads": max_threads,
    });
    if let Some(mem_mb) = max_memory {
        recommendations["min_memory_mb"] =
            serde_json::Value::Number(serde_json::Number::from(mem_mb));
    }
    if total_gpu > 0 {
        recommendations["min_gpu"] = serde_json::Value::Number(serde_json::Number::from(total_gpu));
    }

    // Build manifest
    let checksum_count = manifest_files.len();
    let timestamp = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .map(|d| d.as_secs())
        .unwrap_or(0);

    let manifest = serde_json::json!({
        "format": "oxoflow-bundle-v1",
        "workflow": workflow_path.file_name().and_then(|s| s.to_str()),
        "oxo_flow_version": oxo_version,
        "created_at_epoch": timestamp,
        "entrypoint": workflow_path.file_name().and_then(|s| s.to_str()),
        "files": &manifest_files,
        "containers": &container_refs,
        "resources": {
            "rules": &resource_summary,
            "recommendations": &recommendations,
        },
        // Reserved for bundle signing. Always empty today — present so that adding
        // signatures later is an additive change rather than a manifest format bump.
        // Consumers read the manifest field-by-field, so an empty array is ignored
        // by older versions of oxo-flow.
        "signatures": serde_json::Value::Array(Vec::new()),
    });

    let manifest_json = serde_json::to_string_pretty(&manifest)?;
    let manifest_path = temp_dir.join("manifest.json");
    std::fs::write(&manifest_path, &manifest_json)?;

    // ── Build archive (two concrete formats, finalized via match below) ───

    let archive_file = std::fs::File::create(&output_archive)
        .with_context(|| format!("failed to create archive: {}", output_archive.display()))?;

    // Helper macro to add files to whichever builder we end up with.
    macro_rules! add_files {
        ($builder:expr) => {{
            // Add manifest first, then workflow, then all referenced files
            $builder.append_path_with_name(&manifest_path, "manifest.json")?;
            // Add workflow file
            $builder.append_path_with_name(&wf_dest, wf_filename)?;
            // Add all env/script files
            for (rel_path, _abs_path) in &referenced_files {
                // Backstop for the archive entry names themselves: a `../`
                // entry would be a zip-slip on the consumer side even if the
                // temp-dir copy above had somehow tolerated it.
                ensure_safe_bundle_path(rel_path)?;
                let dest = temp_dir.join(rel_path);
                $builder.append_path_with_name(&dest, rel_path)?;
            }
        }};
    }

    match bundle_format {
        crate::commands::bundle::BundleFormat::TarZst => {
            let enc = zstd::stream::write::Encoder::new(archive_file, 3)
                .context("failed to create zstd encoder")?;
            let mut builder = tar::Builder::new(enc);
            add_files!(builder);
            let enc = builder.into_inner().context("failed to finalize tar")?;
            enc.finish()
                .context("failed to finalize zstd compression")?;
        }
        crate::commands::bundle::BundleFormat::TarGz => {
            let enc = flate2::write::GzEncoder::new(archive_file, flate2::Compression::default());
            let mut builder = tar::Builder::new(enc);
            add_files!(builder);
            let enc = builder.into_inner().context("failed to finalize tar")?;
            enc.finish()
                .context("failed to finalize gzip compression")?;
        }
    }

    // TempDir guard cleans up the staging dir; the lockfiles it may contain
    // have already been read by the archive builder above.

    // ── Summary ───────────────────────────────────────────────────────────

    let archive_size = std::fs::metadata(&output_archive)
        .map(|m| m.len())
        .unwrap_or(0);
    let size_str = if archive_size > 1_048_576 {
        format!("{:.1} MB", archive_size as f64 / 1_048_576.0)
    } else if archive_size > 1_024 {
        format!("{:.1} KB", archive_size as f64 / 1_024.0)
    } else {
        format!("{} B", archive_size)
    };

    eprintln!(
        "{} Published to {}",
        "✓".green().bold(),
        output_archive.display()
    );
    eprintln!("  size:      {}", size_str);
    eprintln!(
        "  files:     {} (workflow + env + scripts/bin)",
        referenced_files.len() + 1
    );
    eprintln!("  checksums: {} files verified (SHA-256)", checksum_count);

    Ok(())
}

/// Reject bundle member paths that could escape the bundle root: absolute
/// paths and `..` components. `referenced_files` paths come straight from
/// workflow TOML (`metadata_file` / `pairs_file` / include references), so
/// a hostile workflow must not steer `temp_dir.join(rel)` outside the
/// staging dir on the publisher's machine, or plant `../` entry names the
/// consumer-side extractor would follow (issue #297 item 7; POSIX-only —
/// oxo-flow targets Linux/macOS).
pub(crate) fn ensure_safe_bundle_path(rel_path: &str) -> Result<()> {
    if rel_path.is_empty() {
        anyhow::bail!("bundle path is empty");
    }
    let path = Path::new(rel_path);
    if path.is_absolute() {
        anyhow::bail!(
            "bundle path '{rel_path}' is absolute — bundle members must be \
             relative to the workflow dir"
        );
    }
    if path
        .components()
        .any(|c| c == std::path::Component::ParentDir)
    {
        anyhow::bail!("bundle path '{rel_path}' contains '..' — it would escape the bundle root");
    }
    Ok(())
}

/// Recursively scan a workflow file (and its `[[include]]` children) for
/// environment file references.
fn scan_workflow_env_files(
    wf_path: &Path,
    workflow_dir: &Path,
    referenced_files: &mut Vec<(String, PathBuf)>,
    container_refs: &mut Vec<serde_json::Value>,
    scanned: &mut std::collections::HashSet<PathBuf>,
) -> Result<()> {
    let canonical = std::path::absolute(wf_path)?;
    if !scanned.insert(canonical) {
        return Ok(()); // Already scanned — avoid cycles
    }

    let content = std::fs::read_to_string(wf_path)
        .with_context(|| format!("failed to read workflow file: {}", wf_path.display()))?;
    let toml_value: toml::Table =
        toml::from_str(&content).context("failed to parse workflow as TOML")?;

    // ── Scan [[rules]] → [rules.environment] ──────────────────────────

    if let Some(rules) = toml_value.get("rules").and_then(|v| v.as_array()) {
        for rule in rules {
            let Some(env) = rule.get("environment") else {
                continue;
            };
            // All local-file environment fields (conda, mamba, pixi, venv, venv_requirements).
            for field in ["conda", "mamba", "pixi", "venv", "venv_requirements"] {
                add_env_file(env, field, workflow_dir, referenced_files);
            }
            // Container image references — record for reproducibility
            for field in ["docker", "singularity"] {
                if let Some(image) = env.get(field).and_then(|v| v.as_str())
                    && !container_refs.iter().any(|c| c["image"] == image)
                {
                    container_refs.push(serde_json::json!({
                        "type": field,
                        "image": image,
                    }));
                }
            }
        }
    }

    // ── Also scan [env_groups] for named environment specs ─────────────

    if let Some(env_groups) = toml_value.get("env_groups").and_then(|v| v.as_table()) {
        for (_group_name, env_spec) in env_groups {
            for field in ["conda", "mamba", "pixi", "venv", "venv_requirements"] {
                add_env_file(env_spec, field, workflow_dir, referenced_files);
            }
            for field in ["docker", "singularity"] {
                if let Some(image) = env_spec.get(field).and_then(|v| v.as_str())
                    && !container_refs.iter().any(|c| c["image"] == image)
                {
                    container_refs.push(serde_json::json!({
                        "type": field,
                        "image": image,
                    }));
                }
            }
        }
    }

    // ── Scan [workflow] for pairs_file / sample_groups_file / metadata_file ──

    if let Some(wf) = toml_value.get("workflow") {
        // metadata_file is data the workflow CANNOT parse without: a bundle
        // that omits it fails at plan time on the consumer side (issue #276).
        // pairs_file/sample_groups_file are collected the same way. All three
        // keep their RELATIVE path from the workflow dir (a file referenced as
        // `data/meta.tsv` must land at `data/meta.tsv` inside the bundle, not
        // flattened to the archive root).
        for key in &["pairs_file", "sample_groups_file", "metadata_file"] {
            if let Some(file_path) = wf.get(key).and_then(|v| v.as_str()) {
                // These three come straight from workflow TOML — a hostile
                // workflow must not steer bundle members outside the bundle
                // (issue #297 item 7). Fail the publish loudly rather than
                // silently trimming the reference.
                ensure_safe_bundle_path(file_path)?;
                let abs_path = workflow_dir.join(file_path);
                if abs_path.exists() {
                    if !referenced_files.iter().any(|(name, _)| name == file_path) {
                        referenced_files.push((file_path.to_string(), abs_path));
                    }
                } else {
                    eprintln!(
                        "  {} {} referenced but not found: {} (bundles that omit it \
                         cannot run on the consumer side)",
                        "⚠".yellow(),
                        key,
                        file_path
                    );
                }
            }
        }
    }

    // ── Follow [[include]] references ─────────────────────────────────
    //
    // Core resolves nested includes relative to the INCLUDING file's
    // parent (`next_base = inc_path.parent()` in `config/parse.rs`), so the
    // scan must track each file's own base dir — joining against the
    // top-level `workflow_dir` would silently skip a sub-workflow's
    // includes that live in a subdirectory.
    if let Some(includes) = toml_value.get("include").and_then(|v| v.as_array()) {
        let file_base = wf_path.parent().unwrap_or(workflow_dir).to_path_buf();
        for inc in includes {
            if let Some(inc_path) = inc.get("path").and_then(|v| v.as_str()) {
                let included_wf = file_base.join(inc_path);
                if included_wf.exists() {
                    scan_workflow_env_files(
                        &included_wf,
                        workflow_dir,
                        referenced_files,
                        container_refs,
                        scanned,
                    )?;
                }
            }
        }
    }

    Ok(())
}

/// Add a single environment file reference if it exists on disk.
fn add_env_file(
    env: &toml::Value,
    field: &str,
    workflow_dir: &Path,
    referenced_files: &mut Vec<(String, PathBuf)>,
) {
    if let Some(env_file) = env.get(field).and_then(|v| v.as_str()) {
        let abs_path = workflow_dir.join(env_file);
        if abs_path.exists() {
            // Keep the DECLARED relative path as the bundle member: the
            // bundled workflow still says `conda = "envs/qc.yaml"`, so
            // flattening the file to the archive root produced a bundle
            // that could not run (audit finding).
            let member = env_file.trim_start_matches("./").replace('\\', "/");
            if member.starts_with('/')
                || member
                    .split('/')
                    .any(|component| component.is_empty() || component == "..")
            {
                eprintln!(
                    "  {} env file path '{}' is not a clean relative path — skipping (field: {})",
                    "⚠".yellow(),
                    env_file,
                    field
                );
                return;
            }
            if !referenced_files.iter().any(|(name, _)| name == &member) {
                referenced_files.push((member, abs_path));
            }
        } else {
            eprintln!(
                "  {} env file referenced but not found: {} (field: {})",
                "⚠".yellow(),
                env_file,
                field
            );
        }
    }
}

/// Collect all files from a directory, preserving relative paths.
fn collect_directory_files(
    dir: &Path,
    prefix: &str,
    referenced_files: &mut Vec<(String, PathBuf)>,
) -> Result<()> {
    for entry in walkdir::WalkDir::new(dir)
        .follow_links(false)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|e| e.file_type().is_file())
    {
        let abs = entry.path().to_path_buf();
        let rel = Path::new(prefix).join(entry.path().strip_prefix(dir).unwrap());
        let rel_str = rel.to_string_lossy().to_string();
        if !referenced_files.iter().any(|(name, _)| name == &rel_str) {
            referenced_files.push((rel_str, abs));
        }
    }
    Ok(())
}

/// Generate conda lockfiles for collected environment YAML files.
///
/// Tries `conda-lock` first, then falls back to `conda env export`.
/// Lockfiles are added to `referenced_files` for inclusion in the bundle.
///
/// Returns the TempDir holding the generated lockfiles — the caller must
/// keep it alive until the bundle archive has been written, since the
/// `referenced_files` entries point inside it. Dropping the guard deletes
/// the directory.
fn generate_lockfiles(
    _workflow_dir: &Path,
    referenced_files: &mut Vec<(String, PathBuf)>,
) -> Option<tempfile::TempDir> {
    // Find conda-lock or compatible tool
    let lock_tool = if std::process::Command::new("conda-lock")
        .arg("--version")
        .output()
        .is_ok_and(|o| o.status.success())
    {
        Some("conda-lock")
    } else {
        None
    };

    if let Some(tool) = lock_tool {
        eprintln!("  {} Generating lockfiles with {}", "→".cyan(), tool);
    } else {
        eprintln!(
            "  {} conda-lock not found — install with: pip install conda-lock",
            "⚠".yellow()
        );
        eprintln!(
            "  {} lockfiles not generated; environments may resolve differently over time",
            "⚠".yellow()
        );
        return None;
    }

    let lock_dir = match tempfile::Builder::new().prefix("oxo-lock-").tempdir() {
        Ok(d) => d,
        Err(e) => {
            eprintln!(
                "  {} failed to create lockfile staging dir: {} — lockfiles not generated",
                "⚠".yellow(),
                e
            );
            return None;
        }
    };
    let temp_dir = lock_dir.path().to_path_buf();

    // Collect conda/mamba env files to lock
    let env_files: Vec<(String, PathBuf)> = referenced_files
        .iter()
        .filter(|(name, _)| name.ends_with(".yaml") || name.ends_with(".yml"))
        .map(|(name, path)| (name.clone(), path.clone()))
        .collect();

    for (name, abs_path) in &env_files {
        let lock_name = format!(
            "{}.lock.yml",
            Path::new(name).file_stem().unwrap().to_string_lossy()
        );
        let lock_path = temp_dir.join(&lock_name);

        eprintln!("    Locking {}...", name);
        let result = std::process::Command::new("conda-lock")
            .args([
                "lock",
                "--file",
                &abs_path.display().to_string(),
                "--platform",
                "linux-64",
                "--platform",
                "osx-64",
                "--lockfile",
                &lock_path.display().to_string(),
                "--quiet",
            ])
            .output();

        match result {
            Ok(output) if output.status.success() => {
                if lock_path.exists() && !referenced_files.iter().any(|(n, _)| n == &lock_name) {
                    eprintln!("      {} {} generated", "✓".green(), lock_name);
                    referenced_files.push((lock_name, lock_path));
                }
            }
            Ok(output) => {
                let stderr = String::from_utf8_lossy(&output.stderr);
                eprintln!(
                    "      {} conda-lock failed for {}: {}",
                    "⚠".yellow(),
                    name,
                    stderr.lines().next().unwrap_or("unknown error")
                );
            }
            Err(e) => {
                eprintln!(
                    "      {} failed to run conda-lock for {}: {}",
                    "⚠".yellow(),
                    name,
                    e
                );
            }
        }
    }

    // Not cleaned here — the caller holds the TempDir guard until the
    // archive builder has read the lockfiles out of it.
    Some(lock_dir)
}