packages/database/src/extension.ts

import { PrismaClient } from '@prisma/client';
import { getOrganizationId } from './context';

export const createTenantExtension = (explicitOrganizationId?: string) => {
  return {
    query: {
      $allModels: {
        async $allOperations({ model, operation, args, query }: any) {
          const organizationId = explicitOrganizationId || getOrganizationId();

          // 🚨 EXCEPTION: Certain models should never be filtered by organizationId
          const EXCLUDED_MODELS = ['Organization', 'TrainingData', 'NormalizationRule'];
          if (EXCLUDED_MODELS.includes(model) || !organizationId) {
            return query(args);
          }

          // List of operations where we want to enforce organizationId
          const filteredOperations = [
            'findMany',
            'findFirst',
            'count',
            'updateMany',
            'deleteMany',
            'aggregate',
            'groupBy'
          ];

          if (filteredOperations.includes(operation)) {
            args.where = { ...args.where, organizationId };
          }

          // Special handling for create: automatically inject organizationId
          if (operation === 'create') {
            args.data = { ...args.data, organizationId };
          }

          // Special handling for upsert: scope where, and inject in create/update
          if (operation === 'upsert') {
            args.where = { ...args.where, organizationId };
            args.create = { ...args.create, organizationId };
            args.update = { ...args.update, organizationId };
          }

          // For update/delete/findUnique
          if (['update', 'delete', 'findUnique'].includes(operation)) {
             args.where = { ...args.where, organizationId };
          }

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

/**
 * TenantClient exposes all standard PrismaClient delegates plus $forOrganization.
 * The double $extends() in the implementation produces a type that TypeScript cannot
 * easily infer, so we cast explicitly to this alias so callers don't need `as any`.
 */
export type TenantClient = PrismaClient & {
  $forOrganization(organizationId: string): PrismaClient;
};

/**
 * Extends a PrismaClient with a helper to get a tenant-bound client
 */
export function withTenantIsolation(prisma: PrismaClient): TenantClient {
  return prisma.$extends({
    client: {
      $forOrganization(organizationId: string) {
        return prisma.$extends(createTenantExtension(organizationId));
      }
    }
  }).$extends(createTenantExtension()) as unknown as TenantClient;
}