import { AsyncLocalStorage } from "node:async_hooks";
import type { PrismaClient } from "@prisma/client";

/**
 * Transaction-safe Prisma client type that excludes methods not available on transaction clients.
 * This prevents services from accidentally calling $connect, $disconnect, $transaction, etc.
 * on a transaction client, which would cause runtime errors.
 */
export type TransactionClient = Omit<
  PrismaClient,
  "$connect" | "$disconnect" | "$transaction" | "$on" | "$use"
>;

/**
 * AsyncLocalStorage for propagating RLS-scoped Prisma client through the call chain.
 * This allows the RlsContextInterceptor to set a transaction-scoped client that
 * services can access via getRlsClient() without explicit dependency injection.
 *
 * The RLS client is a Prisma transaction client that has SET LOCAL app.current_user_id
 * and app.current_workspace_id executed, enabling Row-Level Security policies.
 *
 * @see docs/design/credential-security.md for RLS architecture
 */
const rlsContext = new AsyncLocalStorage<TransactionClient>();

/**
 * Gets the current RLS-scoped Prisma client from AsyncLocalStorage.
 * Returns undefined if no RLS context is set (e.g., unauthenticated routes).
 *
 * Services should use this pattern:
 * ```typescript
 * const client = getRlsClient() ?? this.prisma;
 * ```
 *
 * This ensures they use the RLS-scoped client when available (for authenticated
 * requests) and fall back to the standard client otherwise.
 *
 * @returns The RLS-scoped Prisma transaction client, or undefined
 *
 * @example
 * ```typescript
 * @Injectable()
 * export class TasksService {
 *   constructor(private readonly prisma: PrismaService) {}
 *
 *   async findAll() {
 *     const client = getRlsClient() ?? this.prisma;
 *     return client.task.findMany(); // Automatically filtered by RLS
 *   }
 * }
 * ```
 */
export function getRlsClient(): TransactionClient | undefined {
  return rlsContext.getStore();
}

/**
 * Executes a function with an RLS-scoped Prisma client available via getRlsClient().
 * The client is propagated through the call chain using AsyncLocalStorage and is
 * automatically cleared after the function completes.
 *
 * This is used by RlsContextInterceptor to wrap request handlers.
 *
 * @param client - The RLS-scoped Prisma transaction client
 * @param fn - The function to execute with RLS context
 * @returns The result of the function
 *
 * @example
 * ```typescript
 * await prisma.$transaction(async (tx) => {
 *   await tx.$executeRaw`SET LOCAL app.current_user_id = ${userId}`;
 *
 *   return runWithRlsClient(tx, async () => {
 *     // getRlsClient() now returns tx
 *     return handler();
 *   });
 * });
 * ```
 */
export function runWithRlsClient<T>(client: TransactionClient, fn: () => T): T {
  return rlsContext.run(client, fn);
}