lab/bastion/scripts/generate-completions.ts

#!/usr/bin/env tsx
/**
 * generate-completions.ts -- auto-generates shell completions from the commander.js command tree.
 *
 * Usage:
 *   tsx scripts/generate-completions.ts           # print generated files to stdout
 *   tsx scripts/generate-completions.ts --write   # write completions/ files
 *   tsx scripts/generate-completions.ts --check   # exit 0 if files match, 1 if stale
 *
 * Requires `pnpm build` to have run first (workspace packages must be compiled).
 */

import { Command, type Option, type Argument } from 'commander';
import { readFileSync, writeFileSync, mkdirSync } from 'node:fs';
import { join, dirname } from 'node:path';
import { fileURLToPath } from 'node:url';

const __dirname = dirname(fileURLToPath(import.meta.url));
const ROOT = join(__dirname, '..');

// ============================================================
// Command tree extraction
// ============================================================

interface CmdInfo {
  name: string;
  description: string;
  hidden: boolean;
  options: OptInfo[];
  args: ArgInfo[];
  subcommands: CmdInfo[];
}

interface OptInfo {
  short?: string;
  long: string;
  description: string;
  takesValue: boolean;
  choices?: string[];
  negate: boolean;
}

interface ArgInfo {
  name: string;
  description: string;
  required: boolean;
  variadic: boolean;
  choices?: string[];
}

function extractOption(opt: Option): OptInfo {
  return {
    short: (opt as unknown as Record<string, string>).short || undefined,
    long: (opt as unknown as Record<string, string>).long,
    description: opt.description,
    takesValue: (opt as unknown as Record<string, boolean>).required || (opt as unknown as Record<string, boolean>).optional || false,
    choices: (opt as unknown as Record<string, string[] | undefined>).argChoices || undefined,
    negate: (opt as unknown as Record<string, boolean>).negate || false,
  };
}

function extractArgument(arg: Argument): ArgInfo {
  return {
    name: (arg as unknown as Record<string, string>)._name ?? arg.name(),
    description: arg.description,
    required: (arg as unknown as Record<string, boolean>).required,
    variadic: (arg as unknown as Record<string, boolean>).variadic,
    choices: (arg as unknown as Record<string, string[] | undefined>)._choices || undefined,
  };
}

function extractCommand(cmd: Command): CmdInfo {
  const options = (cmd.options as Option[])
    .filter((o) => {
      const long = (o as unknown as Record<string, string>).long;
      return long !== '--help' && long !== '--version';
    })
    .map(extractOption);

  const args = ((cmd as unknown as Record<string, Argument[]>).registeredArguments ?? [])
    .map(extractArgument);

  const subcommands = (cmd.commands as Command[])
    .filter((sub) => sub.name() !== 'help')
    .map(extractCommand);

  if ((cmd.commands as Command[]).some((sub) => sub.name() === 'help')) {
    subcommands.push({
      name: 'help',
      description: 'display help for command',
      hidden: false,
      options: [],
      args: [],
      subcommands: [],
    });
  }

  return {
    name: cmd.name(),
    description: cmd.description(),
    hidden: (cmd as unknown as Record<string, boolean>)._hidden ?? false,
    options,
    args,
    subcommands,
  };
}

async function extractTree(): Promise<CmdInfo> {
  const { createProgram } = await import('../src/cli/src/index.js') as { createProgram: () => Command };
  const program = createProgram();
  return extractCommand(program);
}

// ============================================================
// Utilities
// ============================================================

function esc(s: string): string {
  return s.replace(/'/g, "\\'");
}

/** Collect all commands recursively with their full path. */
function collectCommands(cmd: CmdInfo, prefix: string[] = []): { path: string[]; cmd: CmdInfo }[] {
  const result: { path: string[]; cmd: CmdInfo }[] = [];
  for (const sub of cmd.subcommands) {
    const fullPath = [...prefix, sub.name];
    result.push({ path: fullPath, cmd: sub });
    result.push(...collectCommands(sub, fullPath));
  }
  return result;
}

// ============================================================
// Fish completion generator
// ============================================================

function generateFish(root: CmdInfo): string {
  const lines: string[] = [];
  const emit = (s: string): void => { lines.push(s); };
  const BIN = root.name;

  emit(`# ${BIN} fish completions -- auto-generated by scripts/generate-completions.ts`);
  emit('# DO NOT EDIT MANUALLY -- run: pnpm completions:generate');
  emit('');
  emit(`complete -c ${BIN} -e`);
  emit(`complete -c ${BIN} -f`);
  emit('');

  // Global options
  emit('# Global options');
  emit(`complete -c ${BIN} -s v -l version -d 'Show version'`);
  emit(`complete -c ${BIN} -s h -l help -d 'Show help'`);
  emit('');

  const allCmds = collectCommands(root);

  // Helper: test if EXACTLY the given subcommand chain is present (for subcommand suggestions)
  emit('# Helper: test if exactly a subcommand chain is active (no extra positional args)');
  emit(`function __${BIN}_using_cmd`);
  emit('    set -l tokens (commandline -opc)');
  emit('    set -l expected $argv');
  emit('    set -l depth (count $expected)');
  emit('    set -l found 0');
  emit('    set -l i 1');
  emit('    for tok in $tokens[2..]');
  emit('        if string match -q -- "-*" $tok');
  emit('            continue');
  emit('        end');
  emit('        set i (math $i + 1)');
  emit('        set -l idx (math $i - 1)');
  emit('        if test $idx -le $depth');
  emit('            if test "$tok" != "$expected[$idx]"');
  emit('                return 1');
  emit('            end');
  emit('            set found (math $found + 1)');
  emit('        else');
  emit('            return 1');
  emit('        end');
  emit('    end');
  emit('    test $found -eq $depth');
  emit('end');
  emit('');

  // Helper: test if command chain STARTS WITH the given prefix (for options that apply after args)
  emit('# Helper: test if command starts with a subcommand chain (options still apply after args)');
  emit(`function __${BIN}_in_cmd`);
  emit('    set -l tokens (commandline -opc)');
  emit('    set -l expected $argv');
  emit('    set -l depth (count $expected)');
  emit('    set -l found 0');
  emit('    for tok in $tokens[2..]');
  emit('        if string match -q -- "-*" $tok');
  emit('            continue');
  emit('        end');
  emit('        set found (math $found + 1)');
  emit('        if test $found -le $depth');
  emit('            if test "$tok" != "$expected[$found]"');
  emit('                return 1');
  emit('            end');
  emit('        end');
  emit('    end');
  emit('    test $found -ge $depth');
  emit('end');
  emit('');

  // Dynamic completions: fetch machine data from bastion API
  emit('# Dynamic: fetch machine hostnames from bastion (installed + queued)');
  emit(`function __${BIN}_installed_hosts`);
  emit('    curl -s http://localhost:8080/api/machines 2>/dev/null | ');
  emit("        python3 -c 'import sys,json; d=json.load(sys.stdin); hosts=[v.get(\"hostname\",\"\") for v in {**d.get(\"install_queue\",{}), **d.get(\"installed\",{})}.values() if v.get(\"hostname\")]; [print(h) for h in set(hosts)]' 2>/dev/null");
  emit('end');
  emit('');

  emit('# Dynamic: fetch all known MAC addresses (discovered + queue + installed)');
  emit(`function __${BIN}_known_macs`);
  emit('    curl -s http://localhost:8080/api/machines 2>/dev/null | ');
  emit("        python3 -c 'import sys,json; d=json.load(sys.stdin); [print(k) for k in {**d.get(\"discovered\",{}), **d.get(\"install_queue\",{}), **d.get(\"installed\",{})}]' 2>/dev/null");
  emit('end');
  emit('');

  emit('# Dynamic: fetch hostnames and MACs from all states');
  emit(`function __${BIN}_hosts_and_macs`);
  emit('    curl -s http://localhost:8080/api/machines 2>/dev/null | ');
  emit("        python3 -c 'import sys,json; d=json.load(sys.stdin); a={**d.get(\"discovered\",{}), **d.get(\"install_queue\",{}), **d.get(\"installed\",{})}; macs=list(a.keys()); hosts=[v.get(\"hostname\",\"\") for v in {**d.get(\"install_queue\",{}), **d.get(\"installed\",{})}.values() if v.get(\"hostname\")]; [print(x) for x in set(macs+hosts)]' 2>/dev/null");
  emit('end');
  emit('');

  // Target completions for commands that accept hostname/IP/MAC
  emit('# Target argument completions');
  // app k3s — takes hostname/IP
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd app k3s install" -a "(__${BIN}_installed_hosts)" -d 'installed host'`);
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd app k3s health" -a "(__${BIN}_installed_hosts)" -d 'installed host'`);
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd app labcontroller deploy" -a "(__${BIN}_installed_hosts)" -d 'installed host'`);
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd app labcontroller status" -a "(__${BIN}_installed_hosts)" -d 'installed host'`);
  // provision install — takes MAC then hostname
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd provision install" -a "(__${BIN}_known_macs)" -d 'MAC address'`);
  // provision reprovision/forget/logs — takes MAC or hostname
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd provision reprovision" -a "(__${BIN}_hosts_and_macs)" -d 'host or MAC'`);
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd provision forget" -a "(__${BIN}_hosts_and_macs)" -d 'host or MAC'`);
  emit(`complete -c ${BIN} -n "__${BIN}_using_cmd provision logs" -a "(__${BIN}_hosts_and_macs)" -d 'host or MAC'`);
  emit('');

  // Top-level commands
  const topCmds = root.subcommands.filter((c) => !c.hidden);
  emit('# Top-level commands');
  for (const cmd of topCmds) {
    emit(`complete -c ${BIN} -n "not __fish_seen_subcommand_from ${topCmds.map((c) => c.name).join(' ')}" -a ${cmd.name} -d '${esc(cmd.description)}'`);
  }
  emit('');

  // Subcommands and options at each level
  for (const { path, cmd } of allCmds) {
    if (cmd.hidden) continue;

    // If this command has subcommands, offer them
    const visibleSubs = cmd.subcommands.filter((s) => !s.hidden);
    if (visibleSubs.length > 0) {
      const parentCondition = `__${BIN}_using_cmd ${path.join(' ')}`;
      emit(`# ${path.join(' ')} subcommands`);
      for (const sub of visibleSubs) {
        emit(`complete -c ${BIN} -n "${parentCondition}" -a ${sub.name} -d '${esc(sub.description)}'`);
      }
      emit('');
    }

    // Options for this command (use __in_cmd so options complete even after positional args)
    if (cmd.options.length > 0) {
      const condition = `__${BIN}_in_cmd ${path.join(' ')}`;
      emit(`# ${path.join(' ')} options`);
      for (const opt of cmd.options) {
        const parts = [`complete -c ${BIN} -n "${condition}"`];
        if (opt.short) parts.push(`-s ${opt.short.replace('-', '')}`);
        parts.push(`-l ${opt.long.replace(/^--/, '')}`);
        parts.push(`-d '${esc(opt.description)}'`);
        if (opt.takesValue) {
          if (opt.choices) {
            parts.push(`-xa '${opt.choices.join(' ')}'`);
          } else {
            parts.push('-x');
          }
        }
        emit(parts.join(' '));
      }
      emit('');
    }
  }

  return lines.join('\n') + '\n';
}

// ============================================================
// Bash completion generator
// ============================================================

function generateBash(root: CmdInfo): string {
  const lines: string[] = [];
  const emit = (s: string): void => { lines.push(s); };
  const BIN = root.name;

  emit(`# ${BIN} bash completions -- auto-generated by scripts/generate-completions.ts`);
  emit('# DO NOT EDIT MANUALLY -- run: pnpm completions:generate');
  emit('');

  const allCmds = collectCommands(root);
  const topCmds = root.subcommands.filter((c) => !c.hidden).map((c) => c.name);

  emit(`_${BIN}() {`);
  emit('  local cur prev words cword');
  emit('  _init_completion || return');
  emit('');
  emit(`  local top_commands="${topCmds.join(' ')}"`);
  emit('');

  // Build chain of subcommands from command line
  emit('  # Extract the subcommand chain (skip options and their values)');
  emit('  local -a subcmd_chain=()');
  emit('  local i skip_next=false');
  emit('  for ((i=1; i < cword; i++)); do');
  emit('    if $skip_next; then skip_next=false; continue; fi');
  emit('    case "${words[i]}" in');
  emit('      -*) ;;  # skip options');
  emit('      *) subcmd_chain+=("${words[i]}") ;;');
  emit('    esac');
  emit('  done');
  emit('');
  emit('  local chain_len=${#subcmd_chain[@]}');
  emit('  local chain_str="${subcmd_chain[*]}"');
  emit('');

  // Build case statement for each command path
  emit('  case "$chain_str" in');

  // Start with the deepest paths first to match longest
  const sortedCmds = [...allCmds].sort((a, b) => b.path.length - a.path.length);

  for (const { path, cmd } of sortedCmds) {
    if (cmd.hidden) continue;
    const pathStr = path.join(' ');
    const visibleSubs = cmd.subcommands.filter((s) => !s.hidden).map((s) => s.name);
    const optFlags: string[] = [];
    for (const opt of cmd.options) {
      if (opt.short) optFlags.push(opt.short);
      optFlags.push(opt.long);
    }
    optFlags.push('-h', '--help');

    const completions = [...visibleSubs, ...optFlags].join(' ');
    emit(`    "${pathStr}")`);
    emit(`      COMPREPLY=($(compgen -W "${completions}" -- "$cur"))`);
    emit('      return ;;');
  }

  // Top-level (no subcommand yet)
  emit('    "")');
  emit(`      COMPREPLY=($(compgen -W "$top_commands -h --help -v --version" -- "$cur"))`);
  emit('      return ;;');

  // Default
  emit('    *)');
  emit('      COMPREPLY=($(compgen -W "-h --help" -- "$cur"))');
  emit('      return ;;');

  emit('  esac');
  emit('}');
  emit('');
  emit(`complete -F _${BIN} ${BIN}`);

  return lines.join('\n') + '\n';
}

// ============================================================
// Main
// ============================================================

async function main(): Promise<void> {
  const mode = process.argv[2] ?? '';

  let tree: CmdInfo;
  try {
    tree = await extractTree();
  } catch (err) {
    console.error('Failed to extract command tree from createProgram().');
    console.error('Make sure workspace packages are built: pnpm build');
    console.error(err);
    process.exit(1);
  }

  const fishContent = generateFish(tree);
  const bashContent = generateBash(tree);

  const completionsDir = join(ROOT, 'completions');
  const fishPath = join(completionsDir, 'labctl.fish');
  const bashPath = join(completionsDir, 'labctl.bash');

  if (mode === '--check') {
    let stale = false;
    try {
      const currentFish = readFileSync(fishPath, 'utf-8');
      if (currentFish !== fishContent) {
        console.error('completions/labctl.fish is stale');
        stale = true;
      }
    } catch {
      console.error('completions/labctl.fish does not exist');
      stale = true;
    }
    try {
      const currentBash = readFileSync(bashPath, 'utf-8');
      if (currentBash !== bashContent) {
        console.error('completions/labctl.bash is stale');
        stale = true;
      }
    } catch {
      console.error('completions/labctl.bash does not exist');
      stale = true;
    }
    if (stale) {
      console.error('Run: pnpm completions:generate');
      process.exit(1);
    }
    console.log('Completions are up to date.');
    process.exit(0);
  }

  if (mode === '--write') {
    mkdirSync(completionsDir, { recursive: true });
    writeFileSync(fishPath, fishContent);
    writeFileSync(bashPath, bashContent);
    console.log(`Wrote ${fishPath}`);
    console.log(`Wrote ${bashPath}`);
    process.exit(0);
  }

  // Default: print to stdout
  console.log('=== completions/labctl.fish ===');
  console.log(fishContent);
  console.log('=== completions/labctl.bash ===');
  console.log(bashContent);
}

main().catch((err) => {
  console.error(err);
  process.exit(1);
});