usage/sdk/

mod.rs

use std::fmt;

use std::path::PathBuf;

use crate::case::AsPascalCase;
use indexmap::IndexMap;

use crate::spec::cmd::SpecCommand;
use crate::spec::output::Selector;
use crate::{Framing, Spec};

pub mod python;
pub mod typescript;

#[derive(Debug, Clone)]
pub enum SdkLanguage {
    TypeScript,
    Python,
}

#[derive(Debug, Clone)]
pub struct SdkOptions {
    pub language: SdkLanguage,
    pub package_name: Option<String>,
    pub source_file: Option<String>,
}

#[derive(Debug)]
pub struct SdkOutput {
    pub files: Vec<SdkFile>,
}

#[derive(Debug)]
pub struct SdkFile {
    pub path: PathBuf,
    pub content: String,
}

pub fn generate(spec: &Spec, opts: &SdkOptions) -> SdkOutput {
    match opts.language {
        SdkLanguage::TypeScript => typescript::generate(spec, opts),
        SdkLanguage::Python => python::generate(spec, opts),
    }
}

/// The choices a generated type can spell out as a union of literals.
///
/// None for `choices run=`: its values come from a command where the program runs, so a
/// type fixed at generation time would refuse ones the program accepts. Such a value is
/// typed as a plain string.
pub(crate) fn typed_choices(arg: &crate::SpecArg) -> Option<&crate::SpecChoices> {
    arg.choices
        .as_ref()
        .filter(|choices| choices.run().is_none())
}

/// Escape JSDoc-terminating sequences in comment text.
pub(crate) fn escape_jsdoc(s: &str) -> String {
    s.replace("*/", r"*\/")
}

/// Escape triple-quote sequences and backslashes in Python docstrings.
pub(crate) fn escape_py_docstring(s: &str) -> String {
    s.replace('\\', r"\\").replace(r#"""""#, r#"\"\"\""#)
}

/// Escape a string for a double-quoted literal.
///
/// Backslashes, quotes, and control characters — the last of these because neither language
/// can carry one literally inside a quoted string, so a value with a newline in it wrote a
/// module that fails to import rather than one that says something wrong. Help text and
/// config defaults both really do contain them.
///
/// Python and TypeScript spell all of these the same way, so one function serves both.
fn escape_string_literal(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    for c in s.chars() {
        match c {
            '\\' => out.push_str(r"\\"),
            '"' => out.push_str(r#"\""#),
            '\n' => out.push_str(r"\n"),
            '\r' => out.push_str(r"\r"),
            '\t' => out.push_str(r"\t"),
            c if c.is_control() => out.push_str(&format!("\\x{:02x}", c as u32)),
            c => out.push(c),
        }
    }
    out
}

/// Escape a string for a Python literal.
pub(crate) fn escape_py_string(s: &str) -> String {
    escape_string_literal(s)
}

/// Escape a string for a TypeScript literal.
pub(crate) fn escape_ts_string(s: &str) -> String {
    escape_string_literal(s)
}

/// A simple code writer with indentation management.
pub(crate) struct CodeWriter {
    buf: String,
    indent: usize,
    indent_str: &'static str,
}

impl CodeWriter {
    pub fn new() -> Self {
        Self {
            buf: String::new(),
            indent: 0,
            indent_str: "  ",
        }
    }

    /// Create a CodeWriter with custom indent string (e.g. "    " for Python).
    pub fn with_indent(indent_str: &'static str) -> Self {
        Self {
            buf: String::new(),
            indent: 0,
            indent_str,
        }
    }

    pub fn line(&mut self, s: &str) {
        if !s.is_empty() {
            for _ in 0..self.indent {
                self.buf.push_str(self.indent_str);
            }
        }
        self.buf.push_str(s);
        self.buf.push('\n');
    }

    pub fn indent(&mut self) {
        self.indent += 1;
    }

    pub fn dedent(&mut self) {
        self.indent = self.indent.saturating_sub(1);
    }

    pub fn finish(self) -> String {
        self.buf
    }
}

impl fmt::Display for CodeWriter {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.buf)
    }
}

pub(crate) fn generated_header(comment_prefix: &str, source: &Option<String>) -> String {
    match source {
        Some(s) => {
            format!("{comment_prefix} @generated by usage-cli from {s}. Do not edit manually.")
        }
        None => format!("{comment_prefix} @generated by usage-cli. Do not edit manually."),
    }
}

/// Returns the PascalCase type name for a command: the command name, or the package name for root.
pub(crate) fn command_type_name(cmd: &SpecCommand, package_name: &str) -> String {
    if cmd.name.is_empty() {
        AsPascalCase(package_name).to_string()
    } else {
        AsPascalCase(&cmd.name).to_string()
    }
}

/// A command type name qualified by its full path, for declarations emitted at module scope.
pub(crate) fn command_path_type_name(cmd: &SpecCommand, package_name: &str) -> String {
    if cmd.full_cmd.is_empty() {
        command_type_name(cmd, package_name)
    } else {
        cmd.full_cmd
            .iter()
            .map(|part| AsPascalCase(part).to_string())
            .collect()
    }
}

// ---------------------------------------------------------------------------
// Declared outputs, as generated client methods
// ---------------------------------------------------------------------------

/// One extra method on a generated command class, for one declared output.
///
/// `exec()` is never changed: a caller who wants raw text keeps getting it. What framing
/// buys is a method whose *shape* matches the wire format, because `json` is read to the
/// end and parsed once while `jsonl` arrives a line at a time and may never finish. Those
/// are different signatures, not different parse calls.
pub(crate) struct OutputMethod {
    /// Appended to `exec`, so `exec_jsonl` / `execJsonl`.
    pub suffix: String,
    pub framing: Framing,
    /// The words that pick this output, ready to append to an argv.
    pub select: Vec<String>,
    /// The selecting flag's property name on the flags bag, so a caller-supplied value of
    /// it can be left out rather than duplicated on the command line.
    pub omit: Option<String>,
    /// The generated type alias a parsed value flows through.
    pub type_alias: String,
    /// The generated constant holding the schema, and the schema itself, where one was
    /// declared.
    pub schema_const: Option<String>,
    pub schema: Option<String>,
    pub help: Option<String>,
}

/// The methods a command's declared outputs earn it.
///
/// Named after the **framing**, not the output's own token. That is the whole point of the
/// split: hk spells its line-delimited output `jsonl` and aube spells the identical format
/// `ndjson`, so `exec_ndjson()` for one and `exec_jsonl()` for the other would put the
/// per-CLI spelling back into every caller. The token goes into the argv this builds; the
/// caller never types it.
pub(crate) fn output_methods(
    cmd: &SpecCommand,
    spec: &Spec,
    package_name: &str,
) -> Vec<OutputMethod> {
    let chain = command_chain(spec, cmd);
    let outputs = crate::spec::output::effective_outputs_ref(spec, chain.iter().copied());
    let select = crate::spec::output::effective_select_ref(spec, &chain);
    let machine: Vec<_> = outputs
        .iter()
        .filter(|o| o.framing != Framing::Text)
        .collect();
    let type_prefix = command_type_name(cmd, package_name);
    machine
        .iter()
        .filter_map(|output| {
            let selector = output.select_argv_with(select.as_deref())?;
            let framing = output.framing.as_str();
            // Several outputs can share a framing — a CLI with both `json` and
            // `json-compact`. The default one keeps the plain name and the rest are
            // suffixed, so the common call stays short and none of them collide.
            let shares = machine
                .iter()
                .filter(|o| o.framing == output.framing)
                .count()
                > 1;
            let suffix = if shares && !output.default {
                format!("{framing}_{}", output.name.replace(['-', '.', ' '], "_"))
            } else {
                framing.to_string()
            };
            Some(OutputMethod {
                type_alias: format!("{type_prefix}{}Output", AsPascalCase(&output.name)),
                schema_const: output
                    .schema
                    .as_ref()
                    .map(|_| format!("{}_{}_SCHEMA", shouty(&type_prefix), shouty(&output.name))),
                schema: output.schema.clone(),
                omit: match &selector {
                    Selector::Value { flag, .. } => Some(flag.trim_start_matches('-').to_string()),
                    Selector::Present { .. } => None,
                },
                select: selector.argv(),
                suffix,
                framing: output.framing,
                help: output.help.clone(),
            })
        })
        .collect()
}

/// Whether a flag answers to a selector spelling, dashes or not.
pub(crate) fn flag_names(flag: &crate::SpecFlag, selector: &str) -> bool {
    let bare = selector.trim_start_matches('-');
    flag.long.iter().any(|l| l == bare)
        || flag.short.iter().any(|s| s.to_string() == bare)
        || flag.name == bare
}

/// SHOUTY_SNAKE, for the generated constant names.
pub(crate) fn shouty(value: &str) -> String {
    let separated = value
        .chars()
        .flat_map(|c| {
            if c.is_uppercase() {
                vec!['_', c]
            } else if c == '-' || c == '.' || c == ' ' {
                vec!['_']
            } else {
                vec![c.to_ascii_uppercase()]
            }
        })
        .collect::<String>()
        .trim_start_matches('_')
        .to_string();
    separated
        .split('_')
        .filter(|part| !part.is_empty())
        .collect::<Vec<_>>()
        .join("_")
}

/// The exit codes a generated client should document, folded from the spec's.
pub(crate) fn exit_codes_for(cmd: &SpecCommand, spec: &Spec) -> Vec<crate::SpecExitCode> {
    let chain = command_chain(spec, cmd);
    crate::spec::exit_code::effective_exit_codes_ref(spec, chain.iter().copied())
}

/// Commands from the first subcommand through `cmd`, recovered from its stamped path.
///
/// SDK renderers recurse with only the current command, while outputs and exit codes inherit
/// through every ancestor. Looking up the chain here keeps those renderers from silently
/// skipping declarations on an intermediate command.
fn command_chain<'a>(spec: &'a Spec, cmd: &'a SpecCommand) -> Vec<&'a SpecCommand> {
    let mut current = &spec.cmd;
    let mut chain = Vec::with_capacity(cmd.full_cmd.len());
    for name in &cmd.full_cmd {
        let Some(next) = current.subcommands.get(name) else {
            // A programmatically constructed command may not belong to `spec`. Preserve the
            // old root-plus-command behavior in that case; parsed specs always take the path.
            return vec![cmd];
        };
        chain.push(next);
        current = next;
    }
    chain
}

// ---------------------------------------------------------------------------
// Choice type collection with collision detection
// ---------------------------------------------------------------------------

/// Maps choice type definitions and provides collision-aware type name lookup.
///
/// When two commands have the same arg/flag name with different choices,
/// the type name is prefixed with the command's PascalCase name to avoid collision.
pub(crate) struct ChoiceTypeMap {
    /// Resolved type name -> choice values
    pub types: IndexMap<String, Vec<String>>,
    /// (cmd_name, item_name) -> resolved type name
    name_map: IndexMap<(String, String), String>,
}

impl ChoiceTypeMap {
    /// Look up the resolved type name for a choice arg/flag.
    /// `cmd_name` is the command's name (empty string for root).
    /// `item_name` is the arg or flag name.
    pub fn lookup(&self, cmd_name: &str, item_name: &str) -> Option<&str> {
        self.name_map
            .get(&(cmd_name.to_string(), item_name.to_string()))
            .map(|s| s.as_str())
    }

    /// Iterate over type definitions (name, choices).
    pub fn iter(&self) -> indexmap::map::Iter<'_, String, Vec<String>> {
        self.types.iter()
    }

    /// Check if there are no choice types.
    pub fn is_empty(&self) -> bool {
        self.types.is_empty()
    }
}

struct ChoiceEntry {
    base_name: String,
    item_name: String,
    cmd_name: String,
    cmd_prefix: String,
    choices: Vec<String>,
}

/// Collects all unique choice types across a command tree.
/// When two commands have the same arg/flag name with different choices,
/// the type name is prefixed with the command's PascalCase name to avoid collision.
pub(crate) fn collect_choice_types(cmd: &SpecCommand) -> ChoiceTypeMap {
    let mut all_entries: Vec<ChoiceEntry> = Vec::new();
    collect_choice_entries(cmd, &mut all_entries);

    // Group by base type name, check for choice differences
    let mut base_groups: IndexMap<String, Vec<&ChoiceEntry>> = IndexMap::new();
    for entry in &all_entries {
        base_groups
            .entry(entry.base_name.clone())
            .or_default()
            .push(entry);
    }

    let mut types = IndexMap::new();
    let mut name_map = IndexMap::new();
    for (base_name, entries) in &base_groups {
        let all_same = entries.windows(2).all(|w| w[0].choices == w[1].choices);
        if all_same {
            types.insert(base_name.clone(), entries[0].choices.clone());
            for entry in entries {
                name_map.insert(
                    (entry.cmd_name.clone(), entry.item_name.clone()),
                    base_name.clone(),
                );
            }
        } else {
            for entry in entries {
                let prefixed = format!("{}{}", entry.cmd_prefix, base_name);
                types.insert(prefixed.clone(), entry.choices.clone());
                name_map.insert((entry.cmd_name.clone(), entry.item_name.clone()), prefixed);
            }
        }
    }

    ChoiceTypeMap { types, name_map }
}

fn collect_choice_entries(cmd: &SpecCommand, entries: &mut Vec<ChoiceEntry>) {
    if cmd.hide {
        return;
    }

    let cmd_prefix = if cmd.name.is_empty() {
        String::new()
    } else {
        AsPascalCase(&cmd.name).to_string()
    };
    let cmd_name = cmd.name.clone();

    for arg in &cmd.args {
        if arg.hide {
            continue;
        }
        if let Some(choices) = typed_choices(arg) {
            let base_name = format!("{}Choice", AsPascalCase(&arg.name));
            entries.push(ChoiceEntry {
                base_name,
                item_name: arg.name.clone(),
                cmd_name: cmd_name.clone(),
                cmd_prefix: cmd_prefix.clone(),
                choices: choices.choices.clone(),
            });
        }
    }

    for flag in &cmd.flags {
        if flag.hide {
            continue;
        }
        if let Some(arg) = &flag.arg
            && let Some(choices) = typed_choices(arg)
        {
            let base_name = format!("{}Choice", AsPascalCase(&flag.name));
            entries.push(ChoiceEntry {
                base_name,
                item_name: flag.name.clone(),
                cmd_name: cmd_name.clone(),
                cmd_prefix: cmd_prefix.clone(),
                choices: choices.choices.clone(),
            });
        }
    }

    for subcmd in cmd.subcommands.values() {
        collect_choice_entries(subcmd, entries);
    }
}

/// Collects type names that need to be imported from the types module.
pub(crate) fn collect_type_imports(
    cmd: &SpecCommand,
    package_name: &str,
    choice_types: &ChoiceTypeMap,
    spec: &Spec,
) -> Vec<String> {
    let mut imports = Vec::new();
    collect_type_imports_recursive(cmd, package_name, choice_types, spec, &mut imports);
    imports.sort();
    imports.dedup();
    imports
}

fn collect_type_imports_recursive(
    cmd: &SpecCommand,
    package_name: &str,
    choice_types: &ChoiceTypeMap,
    spec: &Spec,
    imports: &mut Vec<String>,
) {
    if cmd.hide {
        return;
    }

    let name = command_type_name(cmd, package_name);
    let has_args = cmd.args.iter().any(|a| !a.hide);
    let has_flags = cmd.flags.iter().any(|f| !f.hide);

    if has_args {
        imports.push(format!("{name}Args"));
    }
    if has_flags {
        imports.push(format!("{name}Flags"));
    }

    for arg in &cmd.args {
        if !arg.hide
            && typed_choices(arg).is_some()
            && let Some(type_name) = choice_types.lookup(&cmd.name, &arg.name)
        {
            imports.push(type_name.to_string());
        }
    }
    for flag in &cmd.flags {
        if !flag.hide
            && let Some(arg) = &flag.arg
            && typed_choices(arg).is_some()
            && let Some(type_name) = choice_types.lookup(&cmd.name, &flag.name)
        {
            imports.push(type_name.to_string());
        }
    }

    // The alias a parsed output flows through, so a generated signature never names
    // `unknown` directly and the follow-up that fills it in touches no call site.
    for output in output_methods(cmd, spec, package_name) {
        imports.push(output.type_alias);
    }

    for subcmd in cmd.subcommands.values() {
        collect_type_imports_recursive(subcmd, package_name, choice_types, spec, imports);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_code_writer_display() {
        let mut w = CodeWriter::with_indent("    ");
        w.line("hello");
        w.line("world");
        let displayed = format!("{w}");
        assert!(displayed.contains("hello"));
        assert!(displayed.contains("world"));
    }

    #[test]
    fn test_command_type_name_empty() {
        let cmd = SpecCommand::default();
        assert!(cmd.name.is_empty());
        let result = command_type_name(&cmd, "mypackage");
        assert_eq!(result, "Mypackage");
    }

    #[test]
    fn shouty_collapses_every_separator_run() {
        assert_eq!(shouty("a---b...c"), "A_B_C");
    }

    #[test]
    fn test_generated_header_with_source() {
        let result = generated_header("//", &Some("test.kdl".to_string()));
        assert!(result.contains("test.kdl"));
    }

    #[test]
    fn test_generated_header_without_source() {
        let result = generated_header("//", &None);
        assert!(!result.contains("test.kdl"));
        assert!(result.contains("@generated"));
    }

    /// Hidden command with choices — covers collect_choice_entries skip paths
    /// and collect_type_imports_recursive cmd.hide path.
    #[test]
    fn test_hidden_command_with_choices() {
        let spec: crate::Spec = r##"
            bin "app"
            cmd "visible" help="Visible" {
                arg "env" help="Environment" {
                    choices "dev" "prod"
                }
            }
            cmd "hidden" hide=#true help="Hidden" {
                arg "mode" help="Mode" {
                    choices "fast" "slow"
                }
                flag "--level <n>" help="Level" {
                    choices "1" "2" "3"
                }
            }
        "##
        .parse()
        .unwrap();
        let choice_types = collect_choice_types(&spec.cmd);
        // hidden command's choices should not be collected
        assert!(choice_types.lookup("hidden", "mode").is_none());
        assert!(choice_types.lookup("hidden", "level").is_none());
        // visible command's choices should be collected
        assert!(choice_types.lookup("visible", "env").is_some());
    }

    /// Hidden arg/flag with choices — covers skip paths in collect_choice_entries.
    #[test]
    fn test_hidden_arg_flag_with_choices() {
        let spec: crate::Spec = r##"
            bin "app"
            arg "visible_choice" help="Visible" {
                choices "a" "b"
            }
            arg "hidden_choice" hide=#true help="Hidden" {
                choices "x" "y"
            }
            flag "--visible-flag <val>" help="Visible" {
                choices "m" "n"
            }
            flag "--hidden-flag <val>" hide=#true help="Hidden" {
                choices "p" "q"
            }
        "##
        .parse()
        .unwrap();
        let choice_types = collect_choice_types(&spec.cmd);
        assert!(choice_types.lookup("app", "visible_choice").is_some());
        assert!(choice_types.lookup("app", "hidden_choice").is_none());
        assert!(choice_types.lookup("app", "visible-flag").is_some());
        assert!(choice_types.lookup("app", "hidden-flag").is_none());
    }

    /// Flag with arg choices — covers flag.arg choices import path.
    #[test]
    fn test_flag_arg_choices_import() {
        let spec: crate::Spec = r##"
            bin "app"
            flag "--shell <shell>" help="Shell type" {
                choices "bash" "zsh" "fish"
            }
        "##
        .parse()
        .unwrap();
        let choice_types = collect_choice_types(&spec.cmd);
        let mut imports = Vec::new();
        collect_type_imports_recursive(&spec.cmd, "app", &choice_types, &spec, &mut imports);
        assert!(imports.iter().any(|i| i.contains("Choice")));
    }

    #[test]
    fn nested_commands_inherit_outputs_and_exit_codes_through_the_full_path() {
        let spec: crate::Spec = r#"
            bin "app"
            exit_code 130 "interrupted"
            cmd "report" {
                flag "--format <FORMAT>" global=#true
                output "json" framing="json"
                select "--format"
                exit_code 2 "report failed"
                cmd "watch" {
                    output "jsonl" framing="jsonl"
                    exit_code 3 "stream failed"
                }
            }
        "#
        .parse()
        .unwrap();
        let watch = &spec.cmd.subcommands["report"].subcommands["watch"];

        let methods = output_methods(watch, &spec, "app");
        assert_eq!(
            methods
                .iter()
                .map(|method| method.framing)
                .collect::<Vec<_>>(),
            [Framing::Json, Framing::Jsonl]
        );
        assert_eq!(
            exit_codes_for(watch, &spec)
                .iter()
                .map(|code| code.code)
                .collect::<Vec<_>>(),
            [130, 2, 3]
        );
    }
}