import { PrismaClient } from '@prisma/client';
import { PrismaPg } from '@prisma/adapter-pg';
import { Pool } from 'pg';
import { UserSession } from './auth';

let pool: Pool;
let globalPrisma: PrismaClient;

const getPrismaClient = (): PrismaClient => {
  if (!globalPrisma) {
    const dbUrl = new URL(process.env.DATABASE_URL!);
    pool = new Pool({
      host: dbUrl.hostname,
      port: dbUrl.port ? parseInt(dbUrl.port) : 5432,
      user: decodeURIComponent(dbUrl.username),
      password: decodeURIComponent(dbUrl.password),
      database: dbUrl.pathname.substring(1).split('?')[0],
      ssl: false
    });
    const adapter = new PrismaPg(pool);
    globalPrisma = new PrismaClient({
      adapter,
      log: ['query', 'info', 'warn', 'error']
    });
  }
  return globalPrisma;
};

export const getPrisma = (session?: UserSession | null) => {
  const prisma = getPrismaClient();

  if (!session) {
    return prisma;
  }

  return prisma.$extends({
    client: {
      async $transaction(args: any, options?: any) {
        if (typeof args === 'function') {
          const originalFn = args;
          return prisma.$transaction(async (tx) => {
            await tx.$executeRawUnsafe(`SET LOCAL app.current_user_id = '${session.userId}';`);
            await tx.$executeRawUnsafe(`SET LOCAL app.current_user_role = '${session.role}';`);
            await tx.$executeRawUnsafe(`SET LOCAL app.current_hotel_id = '${session.hotelId}';`);
            await tx.$executeRawUnsafe(`SET LOCAL app.current_region_id = '${session.regionId}';`);
            return originalFn(tx);
          }, options);
        } else if (Array.isArray(args)) {
          const rlsQueries = [
            prisma.$executeRawUnsafe(`SET LOCAL app.current_user_id = '${session.userId}';`),
            prisma.$executeRawUnsafe(`SET LOCAL app.current_user_role = '${session.role}';`),
            prisma.$executeRawUnsafe(`SET LOCAL app.current_hotel_id = '${session.hotelId}';`),
            prisma.$executeRawUnsafe(`SET LOCAL app.current_region_id = '${session.regionId}';`),
          ];
          const results = await prisma.$transaction([...rlsQueries, ...args], options);
          return results.slice(rlsQueries.length);
        }
        return prisma.$transaction(args, options);
      }
    },
    query: {
      $allModels: {
        async $allOperations({ args, query, __internalParams }: any) {
          // Check if we are inside a transaction (either batch or interactive transaction)
          const isInsideTransaction = __internalParams?.transaction !== undefined;

          if (isInsideTransaction) {
            return query(args);
          }

          // Execute RLS parameter setting followed by the original query in a batch transaction
          const results = await prisma.$transaction([
            prisma.$executeRawUnsafe(`SET LOCAL app.current_user_id = '${session.userId}';`),
            prisma.$executeRawUnsafe(`SET LOCAL app.current_user_role = '${session.role}';`),
            prisma.$executeRawUnsafe(`SET LOCAL app.current_hotel_id = '${session.hotelId}';`),
            prisma.$executeRawUnsafe(`SET LOCAL app.current_region_id = '${session.regionId}';`),
            query(args)
          ]);
          // Return the results of the query block
          return results[4];
        }
      }
    }
  });
};