import { Prisma, PrismaClient } from '@prisma/client';
import { AsyncLocalStorage } from 'async_hooks';

// AsyncLocalStorage to support automatic ambient tenant isolation
export const tenantContext = new AsyncLocalStorage<string>();

const globalForPrisma = globalThis as unknown as {
  prisma: ReturnType<typeof createPrismaClient> | undefined;
};

// Models with tenantId field
const TENANT_MODELS = new Set(['User', 'ClassSession', 'Booking', 'MediaAsset', 'EWaiver']);

export function createTenantExtension(explicitTenantId?: string) {
  return Prisma.defineExtension({
    name: 'multiTenantExtension',
    query: {
      $allModels: {
        async $allOperations({ model, operation, args, query }) {
          const currentTenantId = explicitTenantId || tenantContext.getStore();

          if (!currentTenantId || !TENANT_MODELS.has(model)) {
            return query(args);
          }

          const queryArgs = ((args ?? {}) as Record<string, any>);

          switch (operation) {
            case 'findUnique':
            case 'findUniqueOrThrow':
            case 'findFirst':
            case 'findFirstOrThrow':
            case 'findMany':
            case 'count':
            case 'aggregate':
            case 'groupBy': {
              queryArgs.where = {
                ...queryArgs.where,
                tenantId: currentTenantId,
              };
              return query(queryArgs);
            }

            case 'create': {
              queryArgs.data = {
                ...queryArgs.data,
                tenantId: currentTenantId,
              };
              return query(queryArgs);
            }

            case 'createMany':
            case 'createManyAndReturn': {
              if (Array.isArray(queryArgs.data)) {
                queryArgs.data = queryArgs.data.map((item: any) => ({
                  ...item,
                  tenantId: currentTenantId,
                }));
              } else if (queryArgs.data) {
                queryArgs.data = {
                  ...queryArgs.data,
                  tenantId: currentTenantId,
                };
              }
              return query(queryArgs);
            }

            case 'update':
            case 'updateMany':
            case 'delete':
            case 'deleteMany':
            case 'upsert': {
              queryArgs.where = {
                ...queryArgs.where,
                tenantId: currentTenantId,
              };
              if (operation === 'upsert' && queryArgs.create) {
                queryArgs.create = {
                  ...queryArgs.create,
                  tenantId: currentTenantId,
                };
              }
              return query(queryArgs);
            }

            default:
              return query(args);
          }
        },
      },
    },
  });
}

export function createPrismaClient() {
  const baseClient = new PrismaClient({
    log: process.env.NODE_ENV === 'development' ? ['error', 'warn'] : ['error'],
  });

  return baseClient.$extends(createTenantExtension());
}

export const prisma = globalForPrisma.prisma ?? createPrismaClient();

if (process.env.NODE_ENV !== 'production') {
  globalForPrisma.prisma = prisma;
}

export type ExtendedPrismaClient = typeof prisma;

/**
 * Returns a tenant-scoped Prisma client using Prisma Client Extension
 */
export function getTenantPrisma(tenantId: string) {
  const baseClient = new PrismaClient();
  return baseClient.$extends(createTenantExtension(tenantId));
}

export default prisma;