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

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

  private shortToFlag: Record<string, string> = {};
  private shortToArg: Record<string, string> = {};

  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 || {};
    if (options?.short) this.shortToFlag[options.short] = name;
    return this as unknown as Command<
      Prettify<TContext & { [K in TName]: boolean }>
    >;
  }

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

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

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

  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 as unknown as Command<
      Prettify<
        TContext & {
          [K in TName]: TReq extends true ? string : string | undefined;
        }
      >
    >;
  }

  addSubcommand<T extends Record<string, unknown>>(command: Command<T>): this {
    if (!command.name) throw new Error("Subcommands must have a name");

    this.subcommands[command.name] = command as unknown as Command<
      Record<string, unknown>
    >;
    return this;
  }

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

  async parse(argv: string[]): Promise<void> {
    const ctx: Record<string, unknown> = {};
    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;
    }

    // init args to their defaults
    for (const [key, opts] of Object.entries(this.args)) {
      if (opts.default !== undefined) ctx[key] = opts.default;
    }

    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}`);
          }
        }
      } else if (isShort) {
        const shorts = arg.slice(1).split("");
        for (let j = 0; j < shorts.length; j++) {
          const short = shorts[j] as string;
          const flagName = this.shortToFlag[short];
          const argName = this.shortToArg[short];

          if (flagName) {
            ctx[flagName] = true;
          } else if (argName) {
            const opts = this.args[argName];
            let value: string | undefined;
            if (j === shorts.length - 1) {
              if (
                i + 1 < argv.length &&
                !(argv[i + 1] as string).startsWith("-")
              ) {
                value = argv[i + 1];
                i++;
              }
            }
            ctx[argName] = this.parseValue(argName, value, opts!);
          } else {
            throw new Error(`Unknown short option: -${short}`);
          }
        }
      } else {
        positionalsProvided.push(arg);
      }

      i++;
    }

    if (positionalsProvided.length > this.positionals.length)
      throw new Error(
        `Unknown positional argument: ${positionalsProvided[this.positionals.length]}`,
      );

    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 as unknown as TContext);
  }

  private parseValue(
    name: string,
    value: string | undefined,
    opts: ArgumentOptions<unknown, unknown>,
  ) {
    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 unknown as { choices: string[] }).choices;
      if (!choices.includes(value))
        throw new Error(
          `Option --${name} must be one of: ${choices.join(", ")}`,
        );
    }
    return value;
  }
}

export * from "./types";