type FlagKey = Parameters<NS["flags"]>[number][number][0];
type FlagTypeMappings
  // type                 final     options   DefaultNSType
  = [StringConstructor,   string,   string[], string]
  | [StringConstructor[], string[], string[], string[]]
  | [NumberConstructor,   number,   number[], number]
  | [NumberConstructor[], number[], number[], (number | string)[]] // DefualtNSType = (number | string)[] is not a typo
  | [BooleanConstructor,  boolean,  never,    boolean];
type AllFlagTypes = FlagTypeMappings[0];
type GetMapping   <T extends AllFlagTypes> = Extract<FlagTypeMappings, [T, any, any, any]>;
type FinalType    <T extends AllFlagTypes> = GetMapping<T>[1];
type OptionsType  <T extends AllFlagTypes> = GetMapping<T>[2];
type DefaultNSType<T extends AllFlagTypes> = GetMapping<T>[3];

type SchemaDefinition = Record<FlagKey, FlagDefinition<AllFlagTypes>>;

export interface FlagDefinition<T extends AllFlagTypes> {
  type: T;
  default?: FinalType<T>,
  options?: OptionsType<T> | ((data: AutocompleteData, args: string[]) => OptionsType<T>);
  validate?: (ns: NS, flagValue: FinalType<T>) => undefined | string | string[];
  help?: string;
  position?: number;
}

export function createSchema<T extends {
  [K in keyof T]: K extends FlagKey ?
    K extends "_" ? "'_' is used to pass non-flag arguments" :
    K extends "t" ? "'t' is used by the run command" :
    K extends "tail" ? "'tail' is used by the run command" :
    K extends "ram-override" ? "'ram-override' is used by the run command" :
    "type" extends keyof T[K] ?
      T[K]["type"] extends AllFlagTypes ?
        FlagDefinition<T[K]["type"]>
      : never : never : never
}>(schema: T, withHelpFlag: false): {
  [K in keyof T]: FlagDefinition<T[K]["type"]>;
};
export function createSchema<T extends {
  [K in keyof T]: K extends FlagKey ?
    K extends "_" ? "'_' is used to pass non-flag arguments" :
    K extends "t" ? "'t' is used by the run command" :
    K extends "tail" ? "'tail' is used by the run command" :
    K extends "ram-override" ? "'ram-override' is used by the run command" :
    "type" extends keyof T[K] ?
      T[K]["type"] extends AllFlagTypes ?
        FlagDefinition<T[K]["type"]>
      : never : never : never
}>(schema: T, withHelpFlag?: true): "help" extends keyof T ? {
  [K in keyof T]: FlagDefinition<T[K]["type"]>;
} : ExpandProps<{
  [K in Exclude<keyof T, "help">]: FlagDefinition<T[K]["type"]>;
} & {
  help: FlagDefinition<BooleanConstructor>;
}>;
export function createSchema<T extends {
  [K in keyof T]: K extends FlagKey ?
    K extends "_" ? "'_' is used to pass non-flag arguments" :
    K extends "t" ? "'t' is used by the run command" :
    K extends "tail" ? "'tail' is used by the run command" :
    K extends "ram-override" ? "'ram-override' is used by the run command" :
    "type" extends keyof T[K] ?
      T[K]["type"] extends AllFlagTypes ?
        FlagDefinition<T[K]["type"]>
      : never : never : never
}>(schema: T, withHelpFlag = true) {
  const positionals = Object.entries(schema as SchemaDefinition)
    .filter(positionIsNumber)
    .sort((a, b) => a[1].position - b[1].position);
  for (let i = 0; i < positionals.length; i++)
    if (positionals[i][1].position !== i)
      throw new Error([
        "Positional arguments must be in order from 0 to #positionals. Got positions:",
        ...positionals.map(([key, value], i) => `  ${key}: position ${value.position}, expected ${i}`)
      ].join("\n"));
  if (!withHelpFlag)
    return schema;
  return {
    help: helpFlag(),
    ...schema,
  };
}

function positionIsNumber(
  kv: [string, FlagDefinition<AllFlagTypes>]
): kv is [string, FlagDefinition<AllFlagTypes> & { position: number }] {
  return kv[1].position !== undefined;
}

function dashifyFlag(flag: string) {
  return flag.length > 1 ? `--${flag}` : `-${flag}`;
}

export function autocompleteFlags(
  data: AutocompleteData, args: string[],
  schema: SchemaDefinition
): string[] {
  const lastArg = args.at(-1);

  if (!lastArg || !lastArg.startsWith("-")) {
    const setFlags = new Set<string>(args
      .filter(arg => arg.startsWith("-"))
      .map(arg => arg.replace(/^-{1,2}/, "")));
    return Object.entries(schema)
      .filter(([key, value]) => Array.isArray(value.default) || !setFlags.has(key))
      .map(([key, _]) => dashifyFlag(key));
  }

  const flag = lastArg.replace(/^-{1,2}/, "");
  if (!isKeyOf(schema, flag))
    return [];
    
  const options = schema[flag].options;
  if (typeof options === "function")
    return options(data, args).map(val => val.toString());
  return options?.map(val => val.toString()) ?? [];
}

function isFlagOfType<T extends AllFlagTypes>(
  flagDef: FlagDefinition<AllFlagTypes>,
  type: T
): flagDef is FlagDefinition<T> {
  const isArray = Array.isArray(type);
  if (Array.isArray(flagDef.type) !== isArray)
    return false;
  if (isArray)
    return (flagDef.type as Extract<AllFlagTypes, any[]>)[0] === type[0];
  return flagDef.type === type;
}

function isFlagOfTypeEx<T extends AllFlagTypes>(
  flag: {
    def: FlagDefinition<AllFlagTypes>,
    value: DefaultNSType<AllFlagTypes>
  },
  type: T
): flag is {
    def: FlagDefinition<T>,
    value: DefaultNSType<T>
} {
  return isFlagOfType(flag.def, type);
}

interface FlagError {
  flag: FlagKey,
  value?: DefaultNSType<AllFlagTypes>,
  msg: string
}

type OkResults<T extends SchemaDefinition> = ExpandProps<{
  [K in keyof T]: FinalType<T[K]["type"]>;
} & { _: ScriptArg[] }>;

type ErrorResults<T extends SchemaDefinition> = ExpandProps<{
  [K in keyof T]: DefaultNSType<T[K]["type"]> | undefined;
} & { _: ScriptArg[] }>;

export function getFlags<T extends SchemaDefinition>(
  ns: NS,
  schema: T
) {
  const array = Object.entries(schema)
    .map(([key, value]) => {
      if (isFlagOfType(value, String))
        return [key, value.default ?? ""] as
          [FlagKey, DefaultNSType<StringConstructor>];
      if (isFlagOfType(value, [String]))
        return [key, value.default ?? []] as
          [FlagKey, DefaultNSType<StringConstructor[]>];
      if (isFlagOfType(value, Number))
        return [key, value.default ?? Number.NaN] as
          [FlagKey, DefaultNSType<NumberConstructor>];
      if (isFlagOfType(value, [Number]))
        return [key, value.default ?? [] as string[]] as
          [FlagKey, DefaultNSType<NumberConstructor[]>];
      if (isFlagOfType(value, Boolean))
        return [key, false] as
          [FlagKey, DefaultNSType<BooleanConstructor>];
      throw new Error(`Unknown flag '${key}' of type '${value.type}'`);
    });
  const rawFlags = ns.flags(array as any) as Record<FlagKey | "_", DefaultNSType<AllFlagTypes>>;

  const missingKeys = new Set(Object.entries(schema)
    .filter(([_, value]) => value.default === undefined && !isFlagOfType(value, Boolean))
    .map(([key, _]) => key));
  const errors: FlagError[] = [];
  for (const [key, value] of Object.entries(rawFlags)) {
    if (!isKeyOf(schema, key)) {
      if (key !== "_")
        delete rawFlags[key];
      continue;
    }

    const flag = {
      def: schema[key],
      value,
    };
    if (Array.isArray(flag.value)) {
      if (flag.value.length > 0) {
        missingKeys.delete(key);

        if (isFlagOfTypeEx(flag, [Number])) {
          const values = flag.value;
          for (let i = 0; i < values.length; i++) {
            const value = values[i];
            if (typeof value === "number")
              continue;
            const newVal = Number(value);
            if (Number.isNaN(newVal))
              errors.push({
                flag: key,
                value: value,
                msg: `${dashifyFlag(key)}: '${value}' is not a number`,
              });
            else
              values[i] = newVal as any;
          }
        }
      }
    } else if (isFlagOfTypeEx(flag, String)) {
      if (flag.value !== "")
        missingKeys.delete(key);
    } else if (isFlagOfTypeEx(flag, Number)) {
      if (!Number.isNaN(flag.value))
        missingKeys.delete(key);
    } else if (isFlagOfTypeEx(flag, Boolean)) {
      if (flag.def.default)
        rawFlags[key] = !flag.value;
    } else
      throw new Error(`Unknown flag '${key}' of type '${flag.def.type}' with value '${flag.value}'`);
  }

  if (missingKeys.size > 0) {
    for (const key of missingKeys)
      errors.push({
        flag: key,
        msg: `${dashifyFlag(key)}: Flag is required, but none was provided`,
      });
  }
  if (errors.length > 0)
    return { errors, flags: rawFlags as unknown as ErrorResults<T> };

  const flags = rawFlags as Record<FlagKey | "_", FinalType<AllFlagTypes>>;
  for (const [key, value] of Object.entries(flags)) {
    if (key === "_")
      continue;
    const flagDef = schema[key];
    if (flagDef.validate)
      errors.push(...
        [flagDef.validate(ns, value)]
          .flatMap(msg => msg)
          .filter(isNotUndefined)
          .map(msg => ({
            flag: key,
            msg: `${dashifyFlag(key)}: ${msg}`,
          })));
  }
  if (errors.length > 0)
    return { errors, flags: flags as unknown as OkResults<T> };
  return { flags: flags as unknown as OkResults<T> };
}

export function hackableServerFlag(help?: string, defaultServer?: string): FlagDefinition<StringConstructor>;
export function hackableServerFlag(array: false, help?: string, defaultServer?: string): FlagDefinition<StringConstructor>;
export function hackableServerFlag(array: true, help?: string, defaultServer?: string[]): FlagDefinition<StringConstructor[]>;
export function hackableServerFlag(arg1?: boolean | string, arg2?: string, arg3?: string | string[]) {
  let array: boolean;
  let help: string | undefined;
  let defaultServer: string | string[] | undefined;
  if (arg1 === undefined || typeof arg1 === "string") {
    array = false;
    help = arg1;
    defaultServer = arg2;
  } else {
    array = arg1;
    help = arg2;
    defaultServer = arg3;
  }
  return {
    type: array ? [String] : String,
    default: defaultServer,
    options: (data) => data.servers,
    validate: (ns, flagValues) => [flagValues]
      .flatMap(x => x)
      .map(flagValue => {
        if (!ns.serverExists(flagValue))
          return `Server '${flagValue}' does not exist`;
        if (!ns.hasRootAccess(flagValue))
          return `You do not have root access on '${flagValue}'`;
        return;
      }),
    help,
  } as FlagDefinition<StringConstructor | StringConstructor[]>;
}

export function helpFlag(): FlagDefinition<BooleanConstructor> {
  return {
    type: Boolean,
    default: false,
    help: "Print usage help",
  };
}

function typeHint(flag: FlagDefinition<AllFlagTypes>) { 
  const defaults = flag.options !== undefined && typeof flag.options !== "function" ?
    "|" + flag.options
      .map(opt => `${opt}`)
      .join("|") :
    "";
  if (Array.isArray(flag.type))
    return `<${flag.type[0].name.toLowerCase()}${defaults}>...`;
  return `<${flag.type.name.toLowerCase()}${defaults}>`;
}

function defaultHint(flag: FlagDefinition<AllFlagTypes>) {
  if (flag.default === undefined)
    return "";
  if (Array.isArray(flag.default)) {
    if (flag.default.length === 0)
      return "<none>";
    return flag.default.join(", ");
  }
  return flag.default;
}

function splitOnWidth(text: string, width: number) {
  const segmenter = new Intl.Segmenter("en", { granularity: "word" });
  const lines: string[] = [];
  let current = "";
  let whitespace = "";
  for (const segment of segmenter.segment(text)) {
    if (/^\s*$/.test(segment.segment))
      whitespace = whitespace + segment.segment;
    else {
      if (current.length + whitespace.length + segment.segment.length > width) {
        lines.push(current);
        current = segment.segment;
      } else
        current = current + whitespace + segment.segment;
      whitespace = "";
    }
  }
  if (current !== "")
    lines.push(current);
  return lines;
}

function flagUsageLines(
  flags: [FlagKey, FlagDefinition<AllFlagTypes>][],
  maxKeyLength: number
) {
  return flags
    .flatMap(([key, value]) => {
      const lines = [
        `${dashifyFlag(key).padStart(maxKeyLength)}: ` + 
          `${typeHint(value)}`,
      ];
      if (value.default !== undefined)
        lines.push(`${" ".repeat(maxKeyLength)}  Default: ${defaultHint(value)}`);
      if (value.help)
        lines.push(...splitOnWidth(value.help, 60)
          .map(line => `${" ".repeat(maxKeyLength)}  ${line}`));
      return lines;
    });
}

export function flagUsage<T extends SchemaDefinition>(
  ns: NS,
  schema: T
) {
  const maxKeyLength = 4 + Math.max(...
    Object.keys(schema)
      .map(key => dashifyFlag(key).length));
  
  const requiredFlags = Object.entries(schema)
    .filter(([_, value]) => value.default === undefined);
  const optionalFlags = Object.entries(schema)
    .filter(([_, value]) => value.default !== undefined);
  
  const lines = [
    `Usage: ${ns.getScriptName()} ${
      requiredFlags
        .map(([key, value]) => `${dashifyFlag(key)} ${typeHint(value)}`)
        .join(" ")
    }`,
  ];
  if (requiredFlags.length > 0) {
    lines.push("  Required flags:",
      ...flagUsageLines(requiredFlags, maxKeyLength));
  }
  if (optionalFlags.length > 0) {
    lines.push("  Optional flags:",
      ...flagUsageLines(optionalFlags, maxKeyLength));
  }
  return lines;
}

export function printFlagUsage<T extends SchemaDefinition>(
  ns: NS,
  schema: T
) {
  ns.tprint("\n" + flagUsage(ns, schema).join("\n"));
}

export function printFlagErrorsAndUsage<T extends SchemaDefinition>(
  ns: NS, schema: T,
  errors: FlagError[]
) {
  ns.tprint("\nError:\n" + [
      ...errors
        .map(e => `  ${e.msg}`),
      "",
      ...flagUsage(ns, schema)
    ].join("\n"));
}

export function tryGetFlags<T extends SchemaDefinition>(
  ns: NS,
  schema: T
) {
  const { errors, flags } = getFlags(ns, schema);
  if ("help" in flags && typeof flags.help === "boolean" && flags.help) {
    printFlagUsage(ns, schema);
    return undefined;
  }
  if (errors !== undefined) {
    printFlagErrorsAndUsage(ns, schema, errors);
    return undefined;
  }
  return flags;
}

const flagSchema = createSchema({
  target: { ...hackableServerFlag("Target server"), position: 0 },
  test: { type: Number, position: 3 },
});

export async function main(ns: NS) {
  const flags = tryGetFlags(ns, flagSchema);
  if (flags === undefined)
    return;
  ns.tprint(flags);
    
  const hackMoneyPerThread = ns.hackAnalyze(flags.target);
  const hackCanFail = ns.hackAnalyzeChance(flags.target) === 1;
  const hackSecurityPerThread = ns.hackAnalyzeSecurity(1);
}

export function autocomplete(data: AutocompleteData, args: string[]): string[] {
  return autocompleteFlags(data, args, flagSchema);
}

function isKeyOf<T extends Record<PropertyKey, unknown>>(obj: T, key: PropertyKey): key is keyof T {
  return key in obj;
}

type ExpandProps<T> = T extends infer O
  ? { [K in keyof O]: O[K] }
  : never;

function isNotUndefined<T>(x: T): x is Exclude<T, undefined> {
  return x !== undefined;
}