iforgor 0.3.3

use {
    anyhow::{bail, Context},
    serde::{Deserialize, Serialize},
    std::{
        collections::{BTreeMap, BTreeSet},
        io::Write,
        path::PathBuf,
        process,
    },
};

use crate::{
    arg_history::ArgHistory,
    command::{AfterRun, ArgDef, ArgType, CommandId, CommandsSource, UserCommand},
    config::Config,
    ctrlc_handler, discover, execute,
    on_disk::OnDisk,
};

pub struct RunResult {
    pub status: process::ExitStatus,
    pub after_run: AfterRun,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct LegacySources {
    pub sources: BTreeSet<PathBuf>,
}

pub fn run_script_by_id(
    commands: &BTreeMap<CommandId, UserCommand>,
    id: &CommandId,
    app_dir: &std::path::Path,
) -> anyhow::Result<RunResult> {
    let Some(entry) = commands.get(id) else {
        bail!("Unknown command ID {id}")
    };

    let after_run = entry.after_run.clone().unwrap_or_default();

    // Load arg history for remembered values.
    let mut arg_hist = ArgHistory::load(app_dir);

    if entry
        .args
        .iter()
        .any(|a| matches!(a.arg_type, ArgType::Text))
    {
        println!("Use Ctrl+C to abort.\n");
    }

    let mut args_values = Vec::new();
    let mut skipped_args = Vec::new();

    for arg in &entry.args {
        let remembered = if arg.remember && !arg.secret {
            arg_hist.get(id, &arg.name).map(String::from)
        } else {
            None
        };

        let (raw, transformed) = collect_arg(arg, remembered.as_deref())?;

        if transformed.is_empty() && !matches!(arg.arg_type, ArgType::Text) {
            skipped_args.push(arg.name.clone());
        }

        if arg.remember && !arg.secret {
            arg_hist.set(id, &arg.name, &raw);
        }

        args_values.push(transformed);
    }

    if !skipped_args.is_empty() {
        println!(
            "Skipped (no choices available): {}. Check choices/source config.",
            skipped_args.join(", ")
        );
    }

    if let Err(e) = arg_hist.save() {
        eprintln!("Warning: failed to save arg history: {e}");
    }

    if entry.risky {
        print!("⚠️ Script is tagged as risky, are you sure you want to run it? [y/N]: ");
        let mut buf = String::new();
        std::io::stdout().flush()?;
        std::io::stdin().read_line(&mut buf)?;

        if !["y", "yes"].contains(&buf.to_lowercase().trim()) {
            bail!("Aborted execution of risky script")
        }
    }

    let shell_str = entry
        .shell
        .as_ref()
        .map_or("default".into(), |s| format!("{s}"));
    println!("💭 Running \"{}\" with shell \"{shell_str}\"\n", entry.name);

    ctrlc_handler::set_mode(ctrlc_handler::Mode::Ignore);
    let status = execute::execute_command(entry, &args_values)?;
    ctrlc_handler::set_mode(ctrlc_handler::Mode::Kill);

    Ok(RunResult { status, after_run })
}

/// Collect a single argument value based on its ArgDef type.
/// Returns (raw_value, transformed_value). Raw is stored in history,
/// transformed is used for script execution.
fn collect_arg(arg: &ArgDef, remembered: Option<&str>) -> anyhow::Result<(String, String)> {
    let prompt = arg.prompt.as_deref().unwrap_or(&arg.name);

    let raw_value = match &arg.arg_type {
        ArgType::Text => {
            let default = remembered.or(arg.default.as_deref());
            collect_text_arg(prompt, default, arg.secret)?
        }
        ArgType::Select => collect_select_arg(prompt, arg, false, remembered)?,
        ArgType::MultiSelect => collect_select_arg(prompt, arg, true, remembered)?,
    };

    let transformed = match &arg.post_transform {
        Some(cmd) => apply_transform(cmd, &raw_value)?,
        None => raw_value.clone(),
    };

    Ok((raw_value, transformed))
}

/// Collect a plain text argument with optional default.
fn collect_text_arg(prompt: &str, default: Option<&str>, secret: bool) -> anyhow::Result<String> {
    let mut buf = String::new();
    if let Some(def) = default.filter(|d| !d.trim().is_empty()) {
        if secret {
            println!("- {prompt} [Default: *****]:");
        } else {
            println!("- {prompt} [Default: {def}]:");
        }
        std::io::stdin().read_line(&mut buf)?;
        if buf.trim_end_matches(['\n', '\r']).is_empty() {
            return Ok(def.to_string());
        }
    } else {
        print!("- {prompt}: ");
        std::io::stdout().flush()?;
        std::io::stdin().read_line(&mut buf)?;
    }

    Ok(buf.trim().to_string())
}

/// Collect a select or multi-select argument using ichoose TUI.
/// `remembered` is the raw value from last run (newline-separated for multi).
fn collect_select_arg(
    prompt: &str,
    arg: &ArgDef,
    multi: bool,
    remembered: Option<&str>,
) -> anyhow::Result<String> {
    let mut items = resolve_choices(arg)?;

    if items.is_empty() {
        return Ok(String::new());
    }

    // For single select: move remembered value to top so it's highlighted first.
    if !multi {
        if let Some(rem) = remembered {
            if let Some(pos) = items.iter().position(|s| s == rem) {
                let item = items.remove(pos);
                items.insert(0, item);
            }
        }
    }

    // Track which items were previously selected.
    let prev_selected: std::collections::BTreeSet<&str> = remembered
        .map(|rem| rem.lines().collect())
        .unwrap_or_default();

    let entries: Vec<_> = items
        .iter()
        .enumerate()
        .map(|(i, c)| {
            let description = if prev_selected.contains(c.as_str()) {
                Some("(last used)".to_string())
            } else {
                None
            };
            ichoose::ListEntry {
                key: i,
                name: c.clone(),
                description,
            }
        })
        .collect();

    // For multi-select: pre-select previously chosen items.
    let initial_selection = if multi {
        entries
            .iter()
            .filter(|e| prev_selected.contains(e.name.as_str()))
            .map(|e| e.key)
            .collect()
    } else {
        std::collections::BTreeSet::new()
    };

    let selected = ichoose::ListSearch {
        items: &entries,
        filter_callback: None,
        preview_callback: None,
        extra: ichoose::ListSearchExtra {
            title: format!(" {prompt} "),
            multi_select: multi,
            initial_selection,
            preserve_order: true,
            ..Default::default()
        },
    }
    .run()?
    .selected;

    if selected.is_empty() {
        bail!("No selection made for argument \"{prompt}\"");
    }

    let selected_values: Vec<_> = selected
        .iter()
        .filter_map(|&idx| items.get(idx))
        .cloned()
        .collect();

    Ok(selected_values.join("\n"))
}

/// Resolve the list of choices from static options or a dynamic source command.
fn resolve_choices(arg: &ArgDef) -> anyhow::Result<Vec<String>> {
    if !arg.choices.is_empty() {
        return Ok(arg.choices.clone());
    }

    if let Some(ref source_cmd) = arg.source {
        let output = process::Command::new("sh")
            .args(["-c", source_cmd])
            .output()
            .context("failed to execute source command")?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            bail!("Source command failed: {stderr}");
        }

        let stdout = String::from_utf8_lossy(&output.stdout);
        let lines: Vec<String> = stdout
            .lines()
            .map(|l| l.trim().to_string())
            .filter(|l| !l.is_empty())
            .collect();

        return Ok(lines);
    }

    bail!(
        "Argument \"{}\" is select/multi-select but has no choices or source",
        arg.name
    );
}

/// Pipe a value through a post_transform shell command.
fn apply_transform(transform_cmd: &str, input: &str) -> anyhow::Result<String> {
    let mut child = process::Command::new("sh")
        .args(["-c", transform_cmd])
        .stdin(process::Stdio::piped())
        .stdout(process::Stdio::piped())
        .spawn()
        .context("failed to start post_transform command")?;

    if let Some(mut stdin) = child.stdin.take() {
        stdin.write_all(input.as_bytes())?;
    }

    let output = child.wait_with_output()?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        bail!("post_transform command failed: {stderr}");
    }

    Ok(String::from_utf8_lossy(&output.stdout).trim().to_string())
}

/// Compute the domain name for a source file relative to its `.iforgor/` dir.
/// E.g., `.iforgor/ci/docker.toml` → `ci/docker`.
fn compute_domain(
    file_path: &std::path::Path,
    iforgor_dir: Option<&std::path::Path>,
) -> Option<String> {
    if let Some(base) = iforgor_dir {
        if let Ok(relative) = file_path.strip_prefix(base) {
            let domain = relative.with_extension("");
            return domain.to_str().map(|s| s.to_string());
        }
    }
    // Fallback: just the file stem.
    file_path
        .file_stem()
        .and_then(|s| s.to_str())
        .map(|s| s.to_string())
}

pub fn load_scripts_for_source(
    commands: &mut BTreeMap<CommandId, UserCommand>,
    path: PathBuf,
    iforgor_dir: Option<&std::path::Path>,
) -> anyhow::Result<()> {
    let source_path = path.clone();
    let scripts = OnDisk::<CommandsSource>::open(path)?.into_inner();

    for mut script in scripts.entries {
        // Ignore scripts incompatible with current platform.
        if let Some(ref platform) = script.only_on {
            if !platform.eq_ignore_ascii_case(std::env::consts::OS) {
                continue;
            }
        }

        script.source_path = Some(source_path.clone());
        script.project_dir = iforgor_dir
            .and_then(|d| d.parent())
            .or_else(|| source_path.parent())
            .map(|p| p.to_path_buf());
        script.domain = compute_domain(&source_path, iforgor_dir);
        let id = script.resolve_id();

        // Warn on name collision.
        if let Some(existing) = commands.values().find(|c| c.name == script.name) {
            eprintln!(
                "Warning: duplicate command name \"{}\" (overwriting previous)",
                existing.name
            );
        }

        commands.insert(id, script);
    }

    Ok(())
}

/// Load commands from all auto-discovered `.iforgor/` dirs + legacy registered sources.
pub fn load_all_commands(
    start_dir: &std::path::Path,
    legacy_sources: &BTreeSet<PathBuf>,
) -> BTreeMap<CommandId, UserCommand> {
    let mut commands = BTreeMap::new();

    // Load from discovered .iforgor/ folders.
    let iforgor_dirs = discover::discover_iforgor_dirs(start_dir);
    for dir in &iforgor_dirs {
        let domain_files = discover::find_domain_files(dir);
        for file in domain_files {
            if let Err(e) = load_scripts_for_source(&mut commands, file, Some(dir)) {
                eprintln!("Warning: failed to load source: {e}");
            }
        }
    }

    // Load from legacy registered sources (no .iforgor/ dir context).
    for path in legacy_sources {
        if let Err(e) = load_scripts_for_source(&mut commands, path.clone(), None) {
            eprintln!("Warning: failed to load legacy source: {e}");
        }
    }

    // Apply config defaults to commands that don't specify shell/after_run.
    let config = Config::load_merged(start_dir);
    for cmd in commands.values_mut() {
        if cmd.shell.is_none() {
            cmd.shell = config.default_shell.clone();
        }
        if cmd.after_run.is_none() {
            cmd.after_run = config.default_after_run.clone();
        }
    }

    commands
}

#[cfg(test)]
mod tests {
    use {super::*, std::path::Path};

    #[test]
    fn domain_simple_file() {
        let domain = compute_domain(
            Path::new("/project/.iforgor/docker.toml"),
            Some(Path::new("/project/.iforgor")),
        );
        assert_eq!(domain.as_deref(), Some("docker"));
    }

    #[test]
    fn domain_nested() {
        let domain = compute_domain(
            Path::new("/project/.iforgor/ci/docker.toml"),
            Some(Path::new("/project/.iforgor")),
        );
        assert_eq!(domain.as_deref(), Some("ci/docker"));
    }

    #[test]
    fn domain_deeply_nested() {
        let domain = compute_domain(
            Path::new("/project/.iforgor/ci/deploy/prod.toml"),
            Some(Path::new("/project/.iforgor")),
        );
        assert_eq!(domain.as_deref(), Some("ci/deploy/prod"));
    }

    #[test]
    fn domain_legacy_no_iforgor_dir() {
        let domain = compute_domain(Path::new("/some/path/scripts.toml"), None);
        assert_eq!(domain.as_deref(), Some("scripts"));
    }

    #[test]
    fn domain_file_not_under_iforgor_dir_falls_back() {
        let domain = compute_domain(
            Path::new("/other/path/stuff.toml"),
            Some(Path::new("/project/.iforgor")),
        );
        // Can't strip prefix, falls back to stem
        assert_eq!(domain.as_deref(), Some("stuff"));
    }
}