import { Buffer } from "node:buffer";
import { timingSafeEqual } from "node:crypto";
import type { GenericEndpointContext } from "@better-auth/core";
import { APIError, BASE_ERROR_CODES } from "@better-auth/core/error";
import { SocialProviderListEnum } from "@better-auth/core/social-providers";
import { safeJSONParse } from "@better-auth/core/utils/json";
import { base64Url } from "@better-auth/utils/base64";
import { createHash } from "@better-auth/utils/hash";
import { betterFetch } from "@better-fetch/fetch";
import { createAuthEndpoint, sessionMiddleware } from "better-auth/api";
import { setSessionCookie } from "better-auth/cookies";
import type { User } from "better-auth/db";
import { parseUserOutput } from "better-auth/db";
import * as z from "zod";
import { ELECTRON_ERROR_CODES } from "./error-codes";
import type { ElectronOptions } from "./types";

const electronTokenBodySchema = z.object({
	token: z.string().nonempty(),
	state: z.string().nonempty(),
	code_verifier: z.string().nonempty(),
});

export const electronToken = (_opts: ElectronOptions) =>
	createAuthEndpoint(
		"/electron/token",
		{
			method: "POST",
			body: electronTokenBodySchema,
			metadata: {
				scope: "http",
				openapi: {
					description: "Exchange the electron token for a session",
					operationId: "electronToken",
					responses: {
						200: {
							description: "Returns the session token and user",
							content: {
								"application/json": {
									schema: {
										type: "object",
										properties: {
											token: {
												type: "string",
											},
											user: {
												type: "object",
												$ref: "#/components/schemas/User",
											},
										},
										required: ["token", "user"],
									},
								},
							},
						},
					},
				},
			},
		},
		async (ctx) => {
			const token = await ctx.context.internalAdapter.findVerificationValue(
				`electron:${ctx.body.token}`,
			);
			if (!token || token.expiresAt < new Date()) {
				throw APIError.from("NOT_FOUND", ELECTRON_ERROR_CODES.INVALID_TOKEN);
			}

			const tokenRecord = safeJSONParse<Record<string, any>>(token.value);
			if (!tokenRecord) {
				throw APIError.from(
					"INTERNAL_SERVER_ERROR",
					ELECTRON_ERROR_CODES.INVALID_TOKEN,
				);
			}

			if (tokenRecord.state !== ctx.body.state) {
				throw APIError.from("BAD_REQUEST", ELECTRON_ERROR_CODES.STATE_MISMATCH);
			}

			if (!tokenRecord.codeChallenge) {
				throw APIError.from(
					"BAD_REQUEST",
					ELECTRON_ERROR_CODES.MISSING_CODE_CHALLENGE,
				);
			}
			if (tokenRecord.codeChallengeMethod === "s256") {
				const codeChallenge = Buffer.from(
					base64Url.decode(tokenRecord.codeChallenge),
				);
				const codeVerifier = Buffer.from(
					await createHash("SHA-256").digest(ctx.body.code_verifier),
				);

				if (
					codeChallenge.length !== codeVerifier.length ||
					!timingSafeEqual(codeChallenge, codeVerifier)
				) {
					throw APIError.from(
						"BAD_REQUEST",
						ELECTRON_ERROR_CODES.INVALID_CODE_VERIFIER,
					);
				}
			} else {
				if (tokenRecord.codeChallenge !== ctx.body.code_verifier) {
					throw APIError.from(
						"BAD_REQUEST",
						ELECTRON_ERROR_CODES.INVALID_CODE_VERIFIER,
					);
				}
			}
			await ctx.context.internalAdapter.deleteVerificationByIdentifier(
				`electron:${ctx.body.token}`,
			);

			const user = await ctx.context.internalAdapter.findUserById(
				tokenRecord.userId,
			);
			if (!user) {
				throw APIError.from(
					"INTERNAL_SERVER_ERROR",
					BASE_ERROR_CODES.USER_NOT_FOUND,
				);
			}

			const session = await ctx.context.internalAdapter.createSession(user.id);
			if (!session) {
				throw APIError.from(
					"INTERNAL_SERVER_ERROR",
					BASE_ERROR_CODES.FAILED_TO_CREATE_SESSION,
				);
			}

			await setSessionCookie(ctx, {
				session,
				user,
			});

			return ctx.json({
				token: session.token,
				user: parseUserOutput(ctx.context.options, user) as User &
					Record<string, any>,
			});
		},
	);

const electronInitOAuthProxyQuerySchema = z.object({
	provider: z.string().nonempty(),
	state: z.string(),
	code_challenge: z.string(),
	code_challenge_method: z.string().optional(),
});

export const electronInitOAuthProxy = (opts: ElectronOptions) =>
	createAuthEndpoint(
		"/electron/init-oauth-proxy",
		{
			method: "GET",
			query: electronInitOAuthProxyQuerySchema,
			metadata: {
				scope: "http",
				openapi: {
					description: "Initialize the OAuth proxy for the electron app",
					operationId: "electronInitOAuthProxy",
					responses: {
						200: {
							description:
								"Returns the URL to redirect to and if the redirect should be performed",
							content: {
								"application/json": {
									schema: {
										type: "object",
										properties: {
											url: {
												type: "string",
												nullable: true,
											},
											redirect: {
												type: "boolean",
											},
											user: {
												type: "object",
												$ref: "#/components/schemas/User",
											},
											token: {
												type: "string",
											},
										},
										required: ["url", "redirect"],
									},
								},
							},
						},
					},
				},
			},
		},
		async (ctx) => {
			const isSocialProvider = SocialProviderListEnum.safeParse(
				ctx.query.provider,
			);
			if (!isSocialProvider && !ctx.context.getPlugin("generic-oauth")) {
				throw APIError.from("BAD_REQUEST", BASE_ERROR_CODES.PROVIDER_NOT_FOUND);
			}

			const headers = new Headers(ctx.request?.headers);
			headers.set("origin", new URL(ctx.context.baseURL).origin);
			let setCookie: string | null = null;
			const searchParams = new URLSearchParams();
			searchParams.set("client_id", opts.clientID || "electron");
			searchParams.set("code_challenge", ctx.query.code_challenge);
			searchParams.set(
				"code_challenge_method",
				ctx.query.code_challenge_method || "plain",
			);
			searchParams.set("state", ctx.query.state);
			const res = await betterFetch<{
				url: string | undefined;
				redirect: boolean;
				user?: User & Record<string, any>;
				token?: string;
			}>(
				`${isSocialProvider ? "/sign-in/social" : "/sign-in/oauth2"}?${searchParams.toString()}`,
				{
					baseURL: ctx.context.baseURL,
					method: "POST",
					body: {
						provider: ctx.query.provider,
					},
					onResponse: (ctx) => {
						const headers = ctx.response.headers;
						setCookie = headers.get("set-cookie") ?? null;
					},
					headers,
				},
			);

			if (res.error) {
				throw new APIError("INTERNAL_SERVER_ERROR", {
					message: res.error.message || "An unknown error occurred.",
				});
			}

			if (setCookie) {
				ctx.setHeader("set-cookie", setCookie);
			}
			if (res.data.url && res.data.redirect) {
				ctx.setHeader("Location", res.data.url);
				ctx.setStatus(302);
				return;
			}
			return ctx.json(res.data);
		},
	);

const electronTransferUserQuerySchema = z.object({
	client_id: z.string(),
	state: z.string(),
	code_challenge: z.string(),
	code_challenge_method: z.string().optional(),
});
const electronTransferUserBodySchema = z.object({
	callbackURL: z.string().optional(),
});

export const electronTransferUser = (
	_opts: ElectronOptions,
	{
		handleTransfer,
	}: {
		handleTransfer: (
			ctx: GenericEndpointContext,
			payload: {
				client_id: string;
				state: string;
				code_challenge: string;
				code_challenge_method?: string | undefined;
			},
		) => Promise<string | null>;
	},
) =>
	createAuthEndpoint(
		"/electron/transfer-user",
		{
			method: "POST",
			query: electronTransferUserQuerySchema,
			body: electronTransferUserBodySchema,
			use: [sessionMiddleware],
			requireHeaders: true,
			metadata: {
				openapi: {
					description: "Transfer the user to the electron app",
					operationId: "electronTransferUser",
					responses: {
						200: {
							description:
								"Returns the URL to redirect to and if the redirect should be performed",
							content: {
								"application/json": {
									schema: {
										type: "object",
										properties: {
											url: {
												type: "string",
												nullable: true,
											},
											redirect: {
												type: "boolean",
											},
											electron_authorization_code: {
												type: "string",
											},
										},
										required: [
											"url",
											"redirect",
											"electron_authorization_code",
										],
									},
								},
							},
						},
					},
				},
			},
		},
		async (ctx) => {
			const identifier = await handleTransfer(ctx, ctx.query);
			if (identifier === null) {
				throw APIError.from(
					"BAD_REQUEST",
					ELECTRON_ERROR_CODES.INVALID_CLIENT_ID,
				);
			}

			return ctx.json({
				url: ctx.body.callbackURL ? ctx.body.callbackURL : null,
				redirect: ctx.body.callbackURL ? true : false,
				electron_authorization_code: identifier,
			});
		},
	);
