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";