import { createAuthEndpoint } from "@better-auth/core/api";
import { APIError, BASE_ERROR_CODES } from "@better-auth/core/error";
import { safeJSONParse } from "@better-auth/core/utils/json";
import * as z from "zod";
import { sessionMiddleware } from "../../../api";
import type { SecretConfig } from "../../../crypto";
import { symmetricDecrypt, symmetricEncrypt } from "../../../crypto";
import { generateRandomString } from "../../../crypto/random";
import { parseUserOutput } from "../../../db/schema";
import { shouldRequirePassword } from "../../../utils/password";
import { PACKAGE_VERSION } from "../../../version";
import { TWO_FACTOR_ERROR_CODES } from "../error-code";
import type {
	TwoFactorProvider,
	TwoFactorTable,
	UserWithTwoFactor,
} from "../types";
import { verifyTwoFactor } from "../verify-two-factor";

export interface BackupCodeOptions {
	/**
	 * The amount of backup codes to generate
	 *
	 * @default 10
	 */
	amount?: number | undefined;
	/**
	 * The length of the backup codes
	 *
	 * @default 10
	 */
	length?: number | undefined;
	/**
	 * An optional custom function to generate backup codes
	 */
	customBackupCodesGenerate?: (() => string[]) | undefined;
	/**
	 * How to store the backup codes in the database, whether encrypted or plain.
	 */
	storeBackupCodes?:
		| (
				| "plain"
				| "encrypted"
				| {
						encrypt: (token: string) => Promise<string>;
						decrypt: (token: string) => Promise<string>;
				  }
		  )
		| undefined;
	/**
	 * Allow generating backup codes without a password when the user does not
	 * have a credential account.
	 * When enabled, password is still required if a credential account exists.
	 * @default false
	 */
	allowPasswordless?: boolean | undefined;
}

function generateBackupCodesFn(options?: BackupCodeOptions | undefined) {
	return Array.from({ length: options?.amount ?? 10 })
		.fill(null)
		.map(() => generateRandomString(options?.length ?? 10, "a-z", "0-9", "A-Z"))
		.map((code) => `${code.slice(0, 5)}-${code.slice(5)}`);
}

export async function encodeBackupCodes(
	codes: string[],
	secret: string | SecretConfig,
	options?: BackupCodeOptions | undefined,
): Promise<string> {
	const json = JSON.stringify(codes);
	if (options?.storeBackupCodes === "encrypted") {
		return symmetricEncrypt({ data: json, key: secret });
	}
	if (
		typeof options?.storeBackupCodes === "object" &&
		"encrypt" in options?.storeBackupCodes
	) {
		return options.storeBackupCodes.encrypt(json);
	}
	return json;
}

export async function generateBackupCodes(
	secret: string | SecretConfig,
	options?: BackupCodeOptions | undefined,
) {
	const backupCodes = options?.customBackupCodesGenerate
		? options.customBackupCodesGenerate()
		: generateBackupCodesFn(options);
	return {
		backupCodes,
		encryptedBackupCodes: await encodeBackupCodes(backupCodes, secret, options),
	};
}

export async function verifyBackupCode(
	data: {
		backupCodes: string;
		code: string;
	},
	key: string | SecretConfig,
	options?: BackupCodeOptions | undefined,
) {
	const codes = await getBackupCodes(data.backupCodes, key, options);
	if (!codes) {
		return {
			status: false,
			updated: null,
		};
	}
	return {
		status: codes.includes(data.code),
		updated: codes.filter((code) => code !== data.code),
	};
}

export async function getBackupCodes(
	backupCodes: string,
	key: string | SecretConfig,
	options?: BackupCodeOptions | undefined,
) {
	if (options?.storeBackupCodes === "encrypted") {
		const decrypted = await symmetricDecrypt({ key, data: backupCodes });
		return safeJSONParse<string[]>(decrypted);
	}
	if (
		typeof options?.storeBackupCodes === "object" &&
		"decrypt" in options?.storeBackupCodes
	) {
		const decrypted = await options?.storeBackupCodes.decrypt(backupCodes);
		return safeJSONParse<string[]>(decrypted);
	}

	return safeJSONParse<string[]>(backupCodes);
}

const verifyBackupCodeBodySchema = z.object({
	code: z.string().meta({
		description: `A backup code to verify. Eg: "123456"`,
	}),
	/**
	 * Disable setting the session cookie
	 */
	disableSession: z
		.boolean()
		.meta({
			description: "If true, the session cookie will not be set.",
		})
		.optional(),
	/**
	 * if true, the device will be trusted
	 * for 30 days. It'll be refreshed on
	 * every sign in request within this time.
	 */
	trustDevice: z
		.boolean()
		.meta({
			description:
				"If true, the device will be trusted for 30 days. It'll be refreshed on every sign in request within this time. Eg: true",
		})
		.optional(),
});

const viewBackupCodesBodySchema = z.object({
	userId: z.coerce.string().meta({
		description: `The user ID to view all backup codes. Eg: "user-id"`,
	}),
});

export const backupCode2fa = (opts: BackupCodeOptions) => {
	const twoFactorTable = "twoFactor";
	const passwordSchema = z.string().meta({
		description: "The users password.",
	});
	const generateBackupCodesBodySchema = opts.allowPasswordless
		? z.object({
				password: passwordSchema.optional(),
			})
		: z.object({
				password: passwordSchema,
			});

	return {
		id: "backup_code",
		version: PACKAGE_VERSION,
		endpoints: {
			/**
			 * ### Endpoint
			 *
			 * POST `/two-factor/verify-backup-code`
			 *
			 * ### API Methods
			 *
			 * **server:**
			 * `auth.api.verifyBackupCode`
			 *
			 * **client:**
			 * `authClient.twoFactor.verifyBackupCode`
			 *
			 * @see [Read our docs to learn more.](https://better-auth.com/docs/plugins/2fa#api-method-two-factor-verify-backup-code)
			 */
			verifyBackupCode: createAuthEndpoint(
				"/two-factor/verify-backup-code",

				{
					method: "POST",
					body: verifyBackupCodeBodySchema,
					metadata: {
						openapi: {
							description: "Verify a backup code for two-factor authentication",
							responses: {
								"200": {
									description: "Backup code verified successfully",
									content: {
										"application/json": {
											schema: {
												type: "object",
												properties: {
													user: {
														type: "object",
														properties: {
															id: {
																type: "string",
																description: "Unique identifier of the user",
															},
															email: {
																type: "string",
																format: "email",
																nullable: true,
																description: "User's email address",
															},
															emailVerified: {
																type: "boolean",
																nullable: true,
																description: "Whether the email is verified",
															},
															name: {
																type: "string",
																nullable: true,
																description: "User's name",
															},
															image: {
																type: "string",
																format: "uri",
																nullable: true,
																description: "User's profile image URL",
															},
															twoFactorEnabled: {
																type: "boolean",
																description:
																	"Whether two-factor authentication is enabled for the user",
															},
															createdAt: {
																type: "string",
																format: "date-time",
																description:
																	"Timestamp when the user was created",
															},
															updatedAt: {
																type: "string",
																format: "date-time",
																description:
																	"Timestamp when the user was last updated",
															},
														},
														required: [
															"id",
															"twoFactorEnabled",
															"createdAt",
															"updatedAt",
														],
														description:
															"The authenticated user object with two-factor details",
													},
													session: {
														type: "object",
														properties: {
															token: {
																type: "string",
																description: "Session token",
															},
															userId: {
																type: "string",
																description:
																	"ID of the user associated with the session",
															},
															createdAt: {
																type: "string",
																format: "date-time",
																description:
																	"Timestamp when the session was created",
															},
															expiresAt: {
																type: "string",
																format: "date-time",
																description:
																	"Timestamp when the session expires",
															},
														},
														required: [
															"token",
															"userId",
															"createdAt",
															"expiresAt",
														],
														description:
															"The current session object, included unless disableSession is true",
													},
												},
												required: ["user", "session"],
											},
										},
									},
								},
							},
						},
					},
				},
				async (ctx) => {
					const { session, valid } = await verifyTwoFactor(ctx);
					const user = session.user as UserWithTwoFactor;
					const twoFactor = await ctx.context.adapter.findOne<TwoFactorTable>({
						model: twoFactorTable,
						where: [
							{
								field: "userId",
								value: user.id,
							},
						],
					});
					if (!twoFactor) {
						throw APIError.from(
							"BAD_REQUEST",
							TWO_FACTOR_ERROR_CODES.BACKUP_CODES_NOT_ENABLED,
						);
					}
					const validate = await verifyBackupCode(
						{
							backupCodes: twoFactor.backupCodes,
							code: ctx.body.code,
						},
						ctx.context.secretConfig,
						opts,
					);
					if (!validate.status || !validate.updated) {
						throw APIError.from(
							"UNAUTHORIZED",
							TWO_FACTOR_ERROR_CODES.INVALID_BACKUP_CODE,
						);
					}
					const updatedBackupCodes = await encodeBackupCodes(
						validate.updated,
						ctx.context.secretConfig,
						opts,
					);

					const updated = await ctx.context.adapter.update({
						model: twoFactorTable,
						update: {
							backupCodes: updatedBackupCodes,
						},
						where: [
							{
								field: "id",
								value: twoFactor.id,
							},
							{
								field: "backupCodes",
								value: twoFactor.backupCodes,
							},
						],
					});
					if (!updated) {
						throw APIError.fromStatus("CONFLICT", {
							message: "Failed to verify backup code. Please try again.",
						});
					}

					if (!ctx.body.disableSession) {
						return valid(ctx);
					}
					return ctx.json({
						token: session.session?.token,
						user: parseUserOutput(ctx.context.options, session.user),
					});
				},
			),
			/**
			 * ### Endpoint
			 *
			 * POST `/two-factor/generate-backup-codes`
			 *
			 * ### API Methods
			 *
			 * **server:**
			 * `auth.api.generateBackupCodes`
			 *
			 * **client:**
			 * `authClient.twoFactor.generateBackupCodes`
			 *
			 * @see [Read our docs to learn more.](https://better-auth.com/docs/plugins/2fa#api-method-two-factor-generate-backup-codes)
			 */
			generateBackupCodes: createAuthEndpoint(
				"/two-factor/generate-backup-codes",
				{
					method: "POST",
					body: generateBackupCodesBodySchema,
					use: [sessionMiddleware],
					metadata: {
						openapi: {
							description:
								"Generate new backup codes for two-factor authentication",
							responses: {
								"200": {
									description: "Backup codes generated successfully",
									content: {
										"application/json": {
											schema: {
												type: "object",
												properties: {
													status: {
														type: "boolean",
														description:
															"Indicates if the backup codes were generated successfully",
														enum: [true],
													},
													backupCodes: {
														type: "array",
														items: { type: "string" },
														description:
															"Array of generated backup codes in plain text",
													},
												},
												required: ["status", "backupCodes"],
											},
										},
									},
								},
							},
						},
					},
				},
				async (ctx) => {
					const user = ctx.context.session.user as UserWithTwoFactor;
					if (!user.twoFactorEnabled) {
						throw APIError.from(
							"BAD_REQUEST",
							TWO_FACTOR_ERROR_CODES.TWO_FACTOR_NOT_ENABLED,
						);
					}
					const requirePassword = await shouldRequirePassword(
						ctx,
						user.id,
						opts.allowPasswordless,
					);
					if (requirePassword) {
						if (!ctx.body.password) {
							throw APIError.from(
								"BAD_REQUEST",
								BASE_ERROR_CODES.INVALID_PASSWORD,
							);
						}
						await ctx.context.password.checkPassword(user.id, ctx);
					}

					// First, find the twoFactor record to get its id
					const twoFactor = await ctx.context.adapter.findOne<TwoFactorTable>({
						model: twoFactorTable,
						where: [
							{
								field: "userId",
								value: user.id,
							},
						],
					});

					if (!twoFactor) {
						throw APIError.from(
							"BAD_REQUEST",
							TWO_FACTOR_ERROR_CODES.TWO_FACTOR_NOT_ENABLED,
						);
					}

					const backupCodes = await generateBackupCodes(
						ctx.context.secretConfig,
						opts,
					);

					// Use the id to update the record
					await ctx.context.adapter.update({
						model: twoFactorTable,
						update: {
							backupCodes: backupCodes.encryptedBackupCodes,
						},
						where: [
							{
								field: "id",
								value: twoFactor.id,
							},
						],
					});
					return ctx.json({
						status: true,
						backupCodes: backupCodes.backupCodes,
					});
				},
			),
			/**
			 * ### Endpoint
			 *
			 * POST `/two-factor/view-backup-codes`
			 *
			 * ### API Methods
			 *
			 * **server:**
			 * `auth.api.viewBackupCodes`
			 *
			 * @see [Read our docs to learn more.](https://better-auth.com/docs/plugins/2fa#api-method-two-factor-view-backup-codes)
			 */
			viewBackupCodes: createAuthEndpoint(
				{
					method: "POST",
					body: viewBackupCodesBodySchema,
				},
				async (ctx) => {
					const twoFactor = await ctx.context.adapter.findOne<TwoFactorTable>({
						model: twoFactorTable,
						where: [
							{
								field: "userId",
								value: ctx.body.userId,
							},
						],
					});
					if (!twoFactor) {
						throw APIError.from(
							"BAD_REQUEST",
							TWO_FACTOR_ERROR_CODES.BACKUP_CODES_NOT_ENABLED,
						);
					}
					const decryptedBackupCodes = await getBackupCodes(
						twoFactor.backupCodes,
						ctx.context.secretConfig,
						opts,
					);

					if (!decryptedBackupCodes) {
						throw APIError.from(
							"BAD_REQUEST",
							TWO_FACTOR_ERROR_CODES.INVALID_BACKUP_CODE,
						);
					}
					return ctx.json({
						status: true,
						backupCodes: decryptedBackupCodes,
					});
				},
			),
		},
	} satisfies TwoFactorProvider;
};
