Bitburner autocomplete

bitburnerUtils.ts

/**
 * `createFlagSchema` will use the schema passed in for type definitio
 * as well as creating an autocomplete functi
 * and a minor validation setup for using flags
 * @param flagSchema The schema object for the flags, the output type of the `flags` function return is based on this
 * @param args The default arguments outside the flags
 * @returns
 * @example
 * ```ts
 * const flagSchema = createFlagSchema(
 *   {
 *     required: { shortcut: 'r', default: '' as 'option1' | 'option2', options: ['option1', 'option2'], required: true },
 *   },
 *   {
 *     default: [] as string[],
 *     options: (data) => data.servers,
 *     required: true,
 *   },
 * );
 * 
 * export const autocomplete = flagSchema.autocomplete;
 * 
 * export async function main(ns: NS) {
 *   // Fully typed flags object
 *   const flags = flagSchema.flags(ns);
 *   // type = string[]
 *   const servers = flags._;
 *   // type = 'option1'|'option2'
 *   const required = flags.required;
 * }
 * ```
 */
export function createFlagSchema<FlagSchema extends Record<string, Omit<Flag, 'arg'>>, Args extends MainArgs>(
  flagSchema: FlagSchema,
  args?: Args,
): {
  /**
   * The schema converted into an autocomplete function
   */
  autocomplete: (data: AutocompleteData, args: ScriptArg[]) => string[];
  /**
   * A function to create an object with your flags from the arguments
   *
   * It calls the `ns.flags` function for you
   * @param ns
   * @returns
   */
  flags: (ns: NS) => SchemaToTypedFlags<FlagSchema, Args>;
} {
  const flags = Object.entries(flagSchema).map(([arg, value]) => ({ arg, ...value }));
  flags.push({
    arg: '_',
    ...args,
    default: args?.default ?? [],
  } as unknown as Flag);
  // Add one of the default args from the game
  const autocompleteFlags = flags.concat([{ arg: 'tail', default: false }]);
  return {
    autocomplete: createAutocomplete(autocompleteFlags),
    flags: (ns: NS) => flagData<SchemaToTypedFlags<FlagSchema, Args>>(ns, flags),
  };
}

type SchemaToTypedFlags<FlagSchema extends Record<string, Omit<Flag, 'arg'>>, Args extends MainArgs> = {
  [key in keyof (FlagSchema & { _: Args })]: DefaultToPrimitives<(FlagSchema & { _: Args })[key]['default']>;
};

//#region Utilities
/**
Matches any [primitive value](https://developer.mozilla.org/en-US/docs/Glossary/Primitive).

@link https://github.com/sindresorhus/type-fest

@category Type
*/
export type Primitive = null | undefined | string | number | boolean | symbol | bigint;

/**
Returns a boolean for whether the given type is `never`.

@link https://github.com/microsoft/TypeScript/issues/31751#issuecomment-498526919
@link https://stackoverflow.com/a/53984913/10292952
@link https://www.zhenghao.io/posts/ts-never
@link https://github.com/sindresorhus/type-fest

Useful in type utilities, such as checking if something does not occur.

@category Type Guard
@category Utilities
*/
export type IsNever<T> = [T] extends [never] ? true : false;

/**
Returns a boolean for whether the given type `T` is the specified `LiteralType`.

@link https://stackoverflow.com/a/52806744/10292952
@link https://github.com/sindresorhus/type-fest
*/
type LiteralCheck<T, LiteralType extends Primitive> = IsNever<T> extends false // Must be wider than `never`
  ? [T] extends [LiteralType] // Must be narrower than `LiteralType`
    ? [LiteralType] extends [T] // Cannot be wider than `LiteralType`
      ? false
      : true
    : false
  : false;

/**
Returns a boolean for whether the given type is a `true` or `false` [literal type](https://www.typescriptlang.org/docs/handbook/2/everyday-types.html#literal-types).

Useful for:
- providing strongly-typed functions when given literal arguments
- type utilities, such as when constructing parsers and ASTs

@link https://github.com/sindresorhus/type-fest

@category Type Guard
@category Utilities
*/
export type IsBooleanLiteral<T> = LiteralCheck<T, boolean>;

type DefaultToPrimitives<T, NeverType = string> = T extends (infer AT)[]
  ? [AT] extends [never]
    ? NeverType[]
    : AT[]
  : IsBooleanLiteral<T> extends true
  ? boolean
  : T;

/**
 * Recursive deepEqual implementation
 *
 * Probably not that robust...
 * @param a
 * @param b
 * @returns
 */
function deepEqual(a: unknown, b: unknown): a is typeof b {
  if (Object.is(a, b)) return true;
  if (typeof a !== 'object' || typeof b !== 'object' || a === null || b === null) {
    return false;
  }

  if (Array.isArray(a) || Array.isArray(b)) {
    if (!Array.isArray(a) || !Array.isArray(b) || a.length !== b.length) {
      return false;
    }
    for (let i = 0; i < a.length; ++i) {
      if (!deepEqual(a[i], b[i])) return false;
    }
  } else {
    // This is an object!
    if (Object.getPrototypeOf(a) !== Object.getPrototypeOf(b)) {
      return false;
    }
    const aEntries = Object.entries(a);
    const bEntries = Object.entries(b);
    if (aEntries.length !== bEntries.length) {
      return false;
    }
    for (let i = 0; i < aEntries.length; ++i) {
      const [key, value] = aEntries[i];
      if (!(key in b) || !deepEqual(value, (b as Record<string, unknown>)[key])) {
        return false;
      }
    }
  }
  return true;
}

//#endregion Utilities

//#region flags

type Alphabet =
  | 'a'
  | 'b'
  | 'c'
  | 'd'
  | 'e'
  | 'f'
  | 'g'
  | 'h'
  | 'i'
  | 'j'
  | 'k'
  | 'l'
  | 'm'
  | 'n'
  | 'o'
  | 'p'
  | 'q'
  | 'r'
  | 's'
  | 't'
  | 'u'
  | 'v'
  | 'w'
  | 'x'
  | 'y'
  | 'z';
type AllShortcuts = `${Exclude<Alphabet, 't'> | Uppercase<Alphabet>}`;

type OptionsValue<T extends ScriptArg> = ((data: AutocompleteData, args: ScriptArg[]) => T[]) | T[];

export interface Flag {
  /**
   * The name of the argument (--`arg`)
   */
  arg: string;
  /**
   * What is the single letter shortcut for this arg?
   * e.g. `r`
   *
   * This will set the `arg` property of the flags to the shortcut value
   * when the shortcut has a non-default value and the main value is the default value
   */
  shortcut?: AllShortcuts;
  /**
   * The default value for the flag
   *
   * This sets up the type definitions for the flag
   *
   * If you use a type cast, you can tell it that it's a union or something.
   */
  default: ScriptArg | string[];
  /**
   * The allowed options for the item, this is used for validating the inputs
   * to ensure the options provided are used
   *
   * Due to not having `data` outside the autocomplete function,
   * if you use a function, the resulting input won't be validated.
   *
   * This will throw when parsing the flags when validation fails (i.e. a value provided doesn't match an option provided)
   */
  options?: OptionsValue<string>;
  /**
   * Is this argument required? This will throw errors if the argument isn't provided
   *
   * For arrays, this will require at least one value in the array,
   * for primitives, this will require not being the default value.
   */
  required?: boolean;
}

export type MainArgs = Omit<Flag, 'arg' | 'default' | 'shortcut' | 'options'> & {
  default?: ScriptArg[];
  options?: OptionsValue<ScriptArg>;
};

export interface PartialCommand {
  flag: Flag;
  start: number;
  args: ScriptArg[];
}

function getCurrentFlag(schema: Flag[], args: ScriptArg[]): PartialCommand[] {
  // boolean flags have no possible options!
  const availableFlags = new Map(
    schema.flatMap((flag) =>
      flag.shortcut
        ? [
            [flag.arg, flag],
            [flag.shortcut, flag],
          ]
        : [[flag.arg, flag]],
    ),
  );
  const basicArgsFlag = availableFlags.get('_');

  // convert all args to strings to make life easier
  const strings = args.map((arg) => `${arg}`);
  if (strings.at(-1)?.includes(' ')) {
    // The last string is a full quoted text
    // we know we can move to the next autocomplete!
    // if the last string is a single word, then we can't short-circuit like this :(
    const commands: PartialCommand[] = basicArgsFlag
      ? [
          {
            flag: basicArgsFlag,
            start: args.length,
            args: [],
          },
        ]
      : [];
    return commands;
  }
  const lastQuote = strings.findLastIndex((arg) => arg.startsWith('"'));
  const usefulArgs = lastQuote === -1 ? args.slice(-2) : args.slice(lastQuote - 1);
  const offset = args.length - usefulArgs.length;
  const lastFlagIndex = usefulArgs.findLastIndex(
    (arg) => typeof arg === 'string' && arg.startsWith('-') && !arg.includes(' '),
  );
  const flagText = lastFlagIndex !== -1 ? String(usefulArgs[lastFlagIndex]).replace(/^-{1,2}/, '') : '_';
  const commandStart = lastQuote !== -1 ? lastQuote : offset + 1 + (lastFlagIndex !== -1 ? lastFlagIndex : 0);
  const command = args.slice(commandStart);

  // a boolean flag has no args after it! So, any arguments would *have* to be basic args
  const flag =
    typeof availableFlags.get(flagText)?.default === 'boolean' ? basicArgsFlag : availableFlags.get(flagText);
  const commands: PartialCommand[] = [];
  if (flag) {
    commands.push({
      flag,
      start: commandStart,
      args: command,
    });
  }
  if (lastQuote === -1 && basicArgsFlag) {
    if (command.length > 0) {
      // you *could* auto fill here from either flag or main args
      // since we don't see spaces in the autocomplete function
      // In any case, outside of the lastQuote existing,
      // you can *always* autofill from the basic commands
      // for the next word (though not the current one)...
      commands.push({
        flag: basicArgsFlag,
        start: args.length,
        args: [],
      });
    }
  }
  // Only bother to return things that actually have options to return
  return commands.filter((command) => command.flag.options);
}

function transformSchema(schema: Flag[], data: AutocompleteData, args: ScriptArg[]): string[] {
  const commands = getCurrentFlag(schema, args);
  const options = new Set<string>();
  for (const command of commands) {
    // Get the available options for each command
    // add surrounding quotes if the option is multi-word
    const baseCommandOptions = command.flag.options
      ? Array.isArray(command.flag.options)
        ? command.flag.options
        : command.flag.options(data, command.args)
      : [];
    const commandOpts = baseCommandOptions.map((option) =>
      option.includes(' ') && !option.startsWith('"') ? `"${option}"` : option,
    );
    const isFlag = command.flag.arg !== '_';

    if (!commandOpts.length) {
      if (isFlag && command.args.length === 0) {
        // Prevent autocompletion - this flag has no data!
        break;
      }
      continue;
    }
    const commandText = command.args.join(' ');
    const availableOptions = command.args.length
      ? commandOpts
          .filter((option) => commandText && option.startsWith(commandText))
          .map((option) =>
            option
              .split(' ')
              .slice(command.args.length - 1)
              .join(' '),
          )
      : commandOpts;
    availableOptions.forEach((opt) => options.add(opt));
    if (options.has(commandText)) {
      options.delete(commandText);
      // Don't suggest the full text again!
    }
    if (command.flag.arg !== '_') {
      if (!commandText || commandText.startsWith('"') || !baseCommandOptions.includes(commandText)) {
        // prevent basic args from autofilling
        // since this flag isn't finished yet!
        break;
      }
    }
  }
  return [...options];
}
function transformSchemaToNSFlags(schema: Flag[], args: ScriptArg[] = []): Parameters<NS['flags']>[0] {
  const flagArgs = args.reduce((acc, arg, index) => {
    if (index !== args.length - 1 && typeof arg === 'string' && arg.startsWith('-')) {
      // Don't include the last index since that might be the current thing they are autocompleting
      if (arg[1] === '-') {
        acc.add(arg.slice(2));
      } else {
        for (const char of arg.slice(1).split('')) {
          acc.add(char);
        }
      }
    }
    return acc;
  }, new Set<string>());

  return schema.reduce<Parameters<NS['flags']>[0]>((acc, flag) => {
    if (
      flag.arg !== '_' &&
      (Array.isArray(flag.default) || (!flagArgs.has(flag.arg) && !flagArgs.has(flag.shortcut!)))
    ) {
      // if the flag is already used and it's not an array (multiple values),
      // then it shouldn't be up for autocomplete anymore
      // also, the main arg ('_') shouldn't be autocompleted in general.
      acc.push([flag.arg, flag.default]);
      if (flag.shortcut) {
        acc.push([flag.shortcut, flag.default]);
      }
    }
    return acc;
  }, []);
}

/**
 * Only validates non-function options
 * @param schema
 * @param flags
 */
function validateFlags<T extends object>(schema: Flag[], flags: Record<string, ScriptArg | string[]>): ToFlagData<T> {
  const needsValidation = schema.filter(
    (flag): flag is typeof flag & { options: string[] } =>
      Array.isArray(flag.options) && (typeof flag.default === 'string' || Array.isArray(flag.default)),
  );
  if (!needsValidation.length) {
    return flags as ToFlagData<T>;
  }
  for (const flag of needsValidation) {
    const options = new Set(flag.options);

    const data = flags[flag.arg];
    const isArray = Array.isArray(data);
    const dataAsArray = isArray ? data : data == null ? [] : [data];
    if (!dataAsArray.length || data === flag.default) {
      // The data is empty or unset from the default!
      if (flag.required) {
        throw new TypeError(
          `${flag.arg === '_' ? 'Script args' : `Flag --${flag.arg}`} is required, but is empty or the default!`,
        );
      }
      flags[flag.arg] = flag.default;
      // Ignore these values in the rest of the validation!
      continue;
    }
    const invalidValues = dataAsArray.filter((value) => !options.has(`${value}`));
    if (invalidValues.length) {
      const hasDefault = invalidValues.some((value) => value === flag.default);
      throw new TypeError(
        `${flag.arg === '_' ? 'Script args' : `Flag --${flag.arg}`} has invalid values!\n${
          hasDefault ? `Value remained default, but the default value isn't allowed!` : ''
        }\n${isArray ? `Args(s): ${JSON.stringify(data)}\n` : ''}Invalid Value(s): ${
          isArray ? JSON.stringify(invalidValues) : `"${data}"`
        }`,
      );
    }
  }
  return flags as ToFlagData<T>;
}

function createAutocomplete(flagSchema: Flag[]): (data: AutocompleteData, args: ScriptArg[]) => string[] {
  return (data, args) => {
    data.flags(transformSchemaToNSFlags(flagSchema, args));
    return transformSchema(flagSchema, data, args);
  };
}

type ToFlagData<T> = T extends { _: ScriptArg[] } ? T : T & { _: ScriptArg[] };

function flagData<T extends object>(ns: NS, flagSchema: Flag[]): ToFlagData<T> {
  const flags = transformDataWithShortcuts(ns.flags(transformSchemaToNSFlags(flagSchema)), flagSchema);
  return validateFlags<T>(flagSchema, flags);
}

function transformDataWithShortcuts(
  flags: {
    [key: string]: ScriptArg | string[];
  },
  schema: Flag[],
) {
  for (const flag of schema) {
    if (!flag.shortcut) {
      continue;
    }
    const shortcutValue = flags[flag.shortcut];
    const mainValue = flags[flag.arg];
    // Remove the shortcut from the options!
    delete flags[flag.shortcut];
    const shortcutDefaulted = deepEqual(shortcutValue, flag.default);
    const mainDefaulted = deepEqual(mainValue, flag.default);
    if (shortcutDefaulted || !mainDefaulted) {
      continue;
    }
    // Replace the value if main was default
    // and shortcut wasn't default.
    flags[flag.arg] = shortcutValue;
  }
  return flags;
}

//#endregion

script.ts

import { createFlagSchema } from './bitburnerUtils';

const contractTypes = [
  'Find Largest Prime Factor',
  'Subarray with Maximum Sum',
  'Total Ways to Sum',
  'Total Ways to Sum II',
  'Spiralize Matrix',
  'Array Jumping Game',
  'Array Jumping Game II',
];
const flagSchema = createFlagSchema(
  {
    contractTypes: { default: '', options: contractTypes, shortcut: 'e' },
    bool: { default: false, shortcut: 'b' },
    optionsString: {
      default: '',
      options: ['option1', 'option2'],
    },
    optionsArr: {
      default: [],
      options: ['option1', 'option2'],
    },
    multiStringArr: {
      default: ['no options here!', 'Just defaults'],
    },
    literalArr: {
      default: [''] as ('' | 'one' | 'two')[],
      options: ['one', 'two'],
    },
    requiredArg: {
      default: '',
      required: true,
      shortcut: 'r',
    },
    number: {
      default: 3,
    },
  },
  { default: ['One', 'Two'] as string[], options: ['One', 'Two', 'Three', 'Four', 'Five', 'Fourteenth'] },
);

export const autocomplete = flagSchema.autocomplete;

export async function main(ns: NS) {
  // Fully typed flags object
  const flags = flagSchema.flags(ns);
}