import type { AsyncLocalStorage } from "node:async_hooks";
import { getAsyncLocalStorage } from "@better-auth/core/async_hooks";
import type { DBAdapter, DBTransactionAdapter } from "../db/adapter";
import { __getBetterAuthGlobal } from "./global";

type HookContext = {
	adapter: DBTransactionAdapter;
	pendingHooks: Array<() => Promise<void>>;
};

const ensureAsyncStorage = async () => {
	const betterAuthGlobal = __getBetterAuthGlobal();
	if (!betterAuthGlobal.context.adapterAsyncStorage) {
		const AsyncLocalStorage = await getAsyncLocalStorage();
		betterAuthGlobal.context.adapterAsyncStorage = new AsyncLocalStorage();
	}
	return betterAuthGlobal.context
		.adapterAsyncStorage as AsyncLocalStorage<HookContext>;
};

/**
 * This is for internal use only. Most users should use `getCurrentAdapter` instead.
 *
 * It is exposed for advanced use cases where you need direct access to the AsyncLocalStorage instance.
 */
export const getCurrentDBAdapterAsyncLocalStorage = async () => {
	return ensureAsyncStorage();
};

export const getCurrentAdapter = async (
	fallback: DBTransactionAdapter,
): Promise<DBTransactionAdapter> => {
	return ensureAsyncStorage()
		.then((als) => {
			const store = als.getStore();
			return store?.adapter || fallback;
		})
		.catch(() => {
			return fallback;
		});
};

export const runWithAdapter = async <R>(
	adapter: DBAdapter,
	fn: () => R,
): Promise<R> => {
	let called = false;
	return ensureAsyncStorage()
		.then(async (als) => {
			called = true;
			const pendingHooks: Array<() => Promise<void>> = [];
			let result: Awaited<R>;
			let error: unknown;
			let hasError = false;
			try {
				result = await als.run({ adapter, pendingHooks }, fn);
			} catch (err) {
				error = err;
				hasError = true;
			}
			// Execute pending hooks after the function completes (even if it threw)
			for (const hook of pendingHooks) {
				await hook();
			}
			if (hasError) {
				throw error;
			}
			return result!;
		})
		.catch((err) => {
			if (!called) {
				return fn();
			}
			throw err;
		});
};

export const runWithTransaction = async <R>(
	adapter: DBAdapter,
	fn: () => R,
): Promise<R> => {
	let called = true;
	return ensureAsyncStorage()
		.then(async (als) => {
			called = true;
			const pendingHooks: Array<() => Promise<void>> = [];
			let result: Awaited<R>;
			let error: unknown;
			let hasError = false;
			try {
				result = await adapter.transaction(async (trx) => {
					return als.run({ adapter: trx, pendingHooks }, fn);
				});
			} catch (e) {
				hasError = true;
				error = e;
			}
			for (const hook of pendingHooks) {
				await hook();
			}
			if (hasError) {
				throw error;
			}
			return result!;
		})
		.catch((err) => {
			if (!called) {
				return fn();
			}
			throw err;
		});
};

/**
 * Queue a hook to be executed after the current transaction commits.
 * If not in a transaction, the hook will execute immediately.
 */
export const queueAfterTransactionHook = async (
	hook: () => Promise<void>,
): Promise<void> => {
	return ensureAsyncStorage()
		.then((als) => {
			const store = als.getStore();
			if (store) {
				// We're in a transaction context, queue the hook
				store.pendingHooks.push(hook);
			} else {
				// Not in a transaction, execute immediately
				return hook();
			}
		})
		.catch(() => {
			// No async storage available, execute immediately
			return hook();
		});
};
