import { AsyncLocalStorage } from "node:async_hooks";
import { db } from "./db";

export type Transaction = Parameters<Parameters<typeof db.transaction>[0]>[0];

const storage = new AsyncLocalStorage<Transaction>();

/** The transaction opened by authTransaction() around the running code, if any. */
export const currentTransaction = (): Transaction | undefined => storage.getStore();

// What the one auth instance is given as its database. Every access goes to
// the transaction opened by authTransaction() around the running code, or to
// the pool when there is none, so no second auth instance is needed to make
// Better Auth write on a transaction's connection.
export const authDatabase = new Proxy(db, {
    get(_target, property) {
        const target = storage.getStore() ?? db;
        const value = Reflect.get(target, property, target);

        return typeof value === "function" ? value.bind(target) : value;
    },
});

/**
 * Runs the callback in a transaction that everything the auth instance does
 * inside it joins: its endpoints, its hooks and its internal adapter all
 * commit or roll back together with the callback's own writes.
 */
export const authTransaction = <T>(callback: (tx: Transaction) => Promise<T>): Promise<T> =>
    db.transaction(tx => storage.run(tx, () => callback(tx)));