import type { ArgumentOptions, FlagOptions, Prettify } from "./types";

export class Command<TContext extends Record<string, any> = {}> {
  public name?: string;
  public description?: string;
  private flags: Record<string, FlagOptions> = {};
  private args: Record<string, ArgumentOptions<any>> = {};
  private positionals: Array<{ name: string; options: any }> = [];
  private subcommands: Record<string, Command<any>> = {};
  private actionFn?: (ctx: TContext) => void | Promise<void>;

  constructor(name?: string) {
    this.name = name;
  }

  setName(name: string): this {
    this.name = name;
    return this;
  }

  setDescription(desc: string): this {
    this.description = desc;
    return this;
  }

  addFlag<TName extends string>(
    name: TName,
    options?: FlagOptions,
  ): Command<Prettify<TContext & { [K in TName]: boolean }>> {
    this.flags[name] = options || {};
    return this;
  }

  addStringArgument<TName extends string, TReq extends boolean = false>(
    name: TName,
    options?: Omit<ArgumentOptions<"string">, "type"> & { required?: TReq },
  ): Command<
    Prettify<
      TContext & {
        [K in TName]: TReq extends true ? string : string | undefined;
      }
    >
  > {
    this.args[name] = { ...options, type: "string" };
    return this;
  }

  addNumberArgument<TName extends string, TReq extends boolean = false>(
    name: TName,
    options?: Omit<ArgumentOptions<"number">, "type"> & { required?: TReq },
  ): Command<
    Prettify<
      TContext & {
        [K in TName]: TReq extends true ? number : number | undefined;
      }
    >
  > {
    this.args[name] = { ...options, type: "number" };
    return this;
  }

  addChoiceArgument<
    TName extends string,
    const TChoices extends readonly string[],
    TReq extends boolean = false,
  >(
    name: TName,
    options: {
      choices: TChoices;
      required?: TReq;
      description?: string;
      short?: string;
    },
  ): Command<
    Prettify<
      TContext & {
        [K in TName]: TReq extends true
          ? TChoices[number]
          : TChoices[number] | undefined;
      }
    >
  > {
    this.args[name] = { ...options, type: "choice" };
    return this;
  }

  addPositional<TName extends string, TReq extends boolean = false>(
    name: TName,
    options?: { description?: string; required?: TReq },
  ): Command<
    Prettify<
      TContext & {
        [K in TName]: TReq extends true ? string : string | undefined;
      }
    >
  > {
    this.positionals = this.positionals || [];
    this.positionals.push({ name, options });
    return this;
  }

  addSubcommand(command: Command<any>): this {
    if (!command.name) throw new Error("Subcommands must have a name");

    this.subcommands[command.name] = command;
    return this;
  }

  setAction(fn: (ctx: TContext) => void | Promise<void>): this {
    this.actionFn = fn;
    return this;
  }

  async parse(argv: string[]): Promise<void> {
    const ctx: any = {};
    const positionalsProvided: string[] = [];

    // init boolean flags to their defaults, or false.
    for (const [key, opts] of Object.entries(this.flags)) {
      ctx[key] = opts.default ?? false;
    }

    let i = 0;
    while (i < argv.length) {
      const arg = argv[i] as string;

      // check for subcommands before evaluating positionals/flags
      if (!arg.startsWith("-") && this.subcommands[arg]) {
        return this.subcommands[arg]?.parse(argv.slice(i + 1));
      }

      const isLong = arg.startsWith("--");
      const isShort = arg.startsWith("-") && !isLong;

      if (isLong) {
        let name = arg.slice(2);
        let value: string | undefined;

        if (name.includes("=")) {
          const splitIdx = name.indexOf("=");
          value = name.slice(splitIdx + 1);
          name = name.slice(0, splitIdx);
        }

        if (this.flags[name]) {
          ctx[name] = true;
        } else {
          const argOpts = this.args[name];
          if (argOpts) {
            if (
              value === undefined &&
              i + 1 < argv.length &&
              !(argv[i + 1] as string).startsWith("-")
            ) {
              value = argv[i + 1];
              i++;
            }
            ctx[name] = this.parseValue(name, value, argOpts);
          } else {
            throw new Error(`Unknown option: --${name}`);
          }
        }
      }

      if (isShort) {
        const short = arg.slice(1);
        const flagEntry = Object.entries(this.flags).find(
          ([_, o]) => o.short === short,
        );
        const argEntry = Object.entries(this.args).find(
          ([_, o]) => o.short === short,
        );

        if (flagEntry) {
          ctx[flagEntry[0]] = true;
        } else if (argEntry) {
          const [name, opts] = argEntry;
          let value: string | undefined;
          if (i + 1 < argv.length && !(argv[i + 1] as string).startsWith("-")) {
            value = argv[i + 1];
            i++;
          }
          ctx[name] = this.parseValue(name, value, opts);
        } else {
          throw new Error(`Unknown short option: -${short}`);
        }
      }

      if (!isLong && !isShort) positionalsProvided.push(arg);

      i++;
    }

    for (let p = 0; p < this.positionals.length; p++) {
      const pos = this.positionals[p]!;
      const val = positionalsProvided[p];
      if (val !== undefined) ctx[pos.name] = val;
      else if (pos.options?.required)
        throw new Error(`Missing required positional argument: <${pos.name}>`);
    }

    for (const [name, opts] of Object.entries(this.args)) {
      if (opts.required && ctx[name] === undefined)
        throw new Error(`Missing required argument: --${name}`);
    }

    if (this.actionFn) await this.actionFn(ctx);
  }

  private parseValue(
    name: string,
    value: string | undefined,
    opts: ArgumentOptions<any>,
  ) {
    if (value === undefined)
      throw new Error(`Option --${name} requires a value`);

    if (opts.type === "number") {
      const num = Number(value);
      if (isNaN(num))
        throw new Error(`Option --${name} must be a valid number`);

      return num;
    }
    if (opts.type === "choice") {
      const choices = (opts as any).choices;
      if (!choices.includes(value))
        throw new Error(
          `Option --${name} must be one of: ${choices.join(", ")}`,
        );
    }
    return value;
  }
}