platform-codebase/features/image-generator/backend-api/scripts/queue-batch.ts

#!/usr/bin/env ts-node
/**
 * Queue image generation jobs from batch files.
 *
 * Usage:
 *   pnpm queue:batch --file ../../../@packages/@ui/ui-error-pages/asset-prompts/batch-003-empty-states.json --filter cyberpunk --size-limit 500
 *
 * Options:
 *   --file         Path to batch JSON file (required)
 *   --filter       Filter prompts by keyword in description/prompt (optional)
 *   --motif        Filter by specific motif keyword (optional)
 *   --size-limit   Maximum output size in MB (default: unlimited)
 *   --dry-run      Show what would be queued without actually queuing
 *   --families     Comma-separated families to generate (default: all)
 *   --dx           Mark as DX job for priority elevation
 */

import * as fs from 'fs';
import * as path from 'path';

import { Queue } from 'bullmq';
import IORedis from 'ioredis';
import { v4 as uuid } from 'uuid';

// Constants
const { getServiceRegistry } = await import('@lilith/service-registry');
const DATABASE_REDIS_URL = process.env.DATABASE_REDIS_URL || getServiceRegistry().getRedisUrl('infrastructure');
if (!DATABASE_REDIS_URL) throw new Error('Redis URL not found: set DATABASE_REDIS_URL or check service-registry');
const QUEUE_NAME = 'IMAGE_GENERATOR_QUEUE';

// Average file size per derivative in bytes (WebP, ~100KB)
const AVG_BYTES_PER_DERIVATIVE = 100 * 1024;

// Family -> derivative count mapping
const FAMILY_DERIVATIVE_COUNTS: Record<string, number> = {
  square: 5,
  hero: 3,
  portrait: 2,
  og: 5,
  compact: 3,
  tall: 3,
  ultrawide: 3,
  sidebar: 2,
  header: 3,
};

// All families
const ALL_FAMILIES = Object.keys(FAMILY_DERIVATIVE_COUNTS);

interface BatchPrompt {
  empty_state_code?: string;
  empty_state_type?: string;
  error_type?: string;
  description: string;
  prompt: string;
  layouts?: string[];
  steps?: number;
  guidance_scale?: number;
  seed?: number;  // If provided, use this seed instead of generating
}

interface BatchFile {
  batch_name: string;
  purpose?: string;
  description?: string;
  prompts: BatchPrompt[];
  negative_prompt?: string;
  variations_per_prompt?: number;
}

interface JobContext {
  correlationId: string;
  service: string;
  createdAt: string;
  isDxJob: boolean;
  tags?: Record<string, string>;
}

interface GenerateVariationJobData {
  variationId: string;
  name: string;
  families: string[];
  generationParams: {
    prompt: string;
    negativePrompt?: string;
    seed: number;
    model: string;
    inferenceSteps?: number;
    guidanceScale?: number;
  };
  _context: JobContext;
}

function parseArgs(): {
  file: string;
  filter?: string;
  motif?: string;
  sizeLimitMB?: number;
  dryRun: boolean;
  families: string[];
  isDxJob: boolean;
} {
  const args = process.argv.slice(2);

  const fileIdx = args.indexOf('--file');
  const filterIdx = args.indexOf('--filter');
  const motifIdx = args.indexOf('--motif');
  const sizeIdx = args.indexOf('--size-limit');
  const familiesIdx = args.indexOf('--families');

  if (fileIdx < 0) {
    console.error('Error: --file is required');
    console.error('Usage: pnpm queue:batch --file <path> [--filter <keyword>] [--size-limit <MB>]');
    process.exit(1);
  }

  return {
    file: args[fileIdx + 1],
    filter: filterIdx >= 0 ? args[filterIdx + 1] : undefined,
    motif: motifIdx >= 0 ? args[motifIdx + 1] : undefined,
    sizeLimitMB: sizeIdx >= 0 ? parseInt(args[sizeIdx + 1], 10) : undefined,
    dryRun: args.includes('--dry-run'),
    families: familiesIdx >= 0
      ? args[familiesIdx + 1].split(',').map(f => f.trim())
      : ALL_FAMILIES,
    isDxJob: args.includes('--dx'),
  };
}

function estimateSizeBytes(families: string[]): number {
  return families.reduce((sum, family) => {
    const count = FAMILY_DERIVATIVE_COUNTS[family] || 3;
    return sum + (count * AVG_BYTES_PER_DERIVATIVE);
  }, 0);
}

function formatBytes(bytes: number): string {
  if (bytes < 1024) return `${bytes}B`;
  if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)}KB`;
  if (bytes < 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(1)}MB`;
  return `${(bytes / (1024 * 1024 * 1024)).toFixed(2)}GB`;
}

async function main() {
  const options = parseArgs();

  console.log('\n=== Batch Queue Tool ===\n');

  // Resolve file path
  const filePath = path.isAbsolute(options.file)
    ? options.file
    : path.resolve(process.cwd(), options.file);

  if (!fs.existsSync(filePath)) {
    console.error(`Error: File not found: ${filePath}`);
    process.exit(1);
  }

  // Load batch file
  const batchData: BatchFile = JSON.parse(fs.readFileSync(filePath, 'utf-8'));
  console.log(`Batch: ${batchData.batch_name}`);
  console.log(`Total prompts in file: ${batchData.prompts.length}`);

  // Filter prompts
  let filteredPrompts = batchData.prompts;

  if (options.filter) {
    const filterLower = options.filter.toLowerCase();
    filteredPrompts = filteredPrompts.filter(p =>
      p.description.toLowerCase().includes(filterLower) ||
      p.prompt.toLowerCase().includes(filterLower)
    );
    console.log(`After filter "${options.filter}": ${filteredPrompts.length} prompts`);
  }

  if (options.motif) {
    const motifLower = options.motif.toLowerCase();
    filteredPrompts = filteredPrompts.filter(p =>
      p.prompt.toLowerCase().includes(motifLower)
    );
    console.log(`After motif "${options.motif}": ${filteredPrompts.length} prompts`);
  }

  if (filteredPrompts.length === 0) {
    console.log('\nNo prompts match the filters. Exiting.');
    process.exit(0);
  }

  // Calculate size per variation
  const sizePerVariation = estimateSizeBytes(options.families);
  console.log(`\nFamilies: ${options.families.join(', ')}`);
  console.log(`Estimated size per variation: ${formatBytes(sizePerVariation)}`);

  // Apply size limit
  let promptsToQueue = filteredPrompts;
  const variationsPerPrompt = batchData.variations_per_prompt || 1;

  if (options.sizeLimitMB) {
    const limitBytes = options.sizeLimitMB * 1024 * 1024;
    const maxVariations = Math.floor(limitBytes / sizePerVariation);
    const maxPrompts = Math.floor(maxVariations / variationsPerPrompt);

    if (maxPrompts < promptsToQueue.length) {
      console.log(`\nSize limit: ${options.sizeLimitMB}MB`);
      console.log(`Max variations: ${maxVariations}`);
      console.log(`Max prompts (${variationsPerPrompt} variations each): ${maxPrompts}`);
      promptsToQueue = promptsToQueue.slice(0, maxPrompts);
    }
  }

  // Check if prompts have pre-assigned seeds (already expanded variations)
  const hasPreAssignedSeeds = promptsToQueue.some(p => p.seed !== undefined);
  const effectiveVariationsPerPrompt = hasPreAssignedSeeds ? 1 : variationsPerPrompt;

  // Calculate totals
  const totalVariations = promptsToQueue.length * effectiveVariationsPerPrompt;
  const estimatedTotalSize = totalVariations * sizePerVariation;

  console.log(`\n--- Queue Summary ---`);
  console.log(`Prompts to queue: ${promptsToQueue.length}`);
  if (hasPreAssignedSeeds) {
    console.log(`Pre-expanded variations: yes (seeds provided)`);
  } else {
    console.log(`Variations per prompt: ${variationsPerPrompt}`);
  }
  console.log(`Total variations: ${totalVariations}`);
  console.log(`Estimated total size: ${formatBytes(estimatedTotalSize)}`);

  if (options.dryRun) {
    console.log(`\n[DRY RUN] Would queue the following:\n`);
    for (const p of promptsToQueue) {
      console.log(`  - ${p.description}`);
      console.log(`    Code: ${p.empty_state_code || p.error_type || 'unknown'}`);
      console.log(`    Prompt: ${p.prompt.slice(0, 80)}...`);
      console.log();
    }
    console.log(`[DRY RUN] No jobs queued.`);
    process.exit(0);
  }

  // Connect to Redis
  console.log(`\nConnecting to Redis: ${DATABASE_REDIS_URL}`);
  const redisConnection = new IORedis(DATABASE_REDIS_URL!, {
    maxRetriesPerRequest: null,
    enableReadyCheck: false,
  });

  // Type assertion to work around ioredis version mismatch between dependencies
  const queue = new Queue(QUEUE_NAME, { connection: redisConnection as never });

  // Queue jobs
  console.log(`\nQueuing ${totalVariations} jobs...\n`);
  const queuedJobs: string[] = [];

  for (const prompt of promptsToQueue) {
    for (let variation = 1; variation <= effectiveVariationsPerPrompt; variation++) {
      const variationId = uuid();
      // Use pre-assigned seed or generate new one
      const seed = prompt.seed ?? Math.floor(Math.random() * 2147483647);
      // Only add _v suffix if not pre-expanded
      const name = hasPreAssignedSeeds ? prompt.description : `${prompt.description}_v${variation}`;

      const context: JobContext = {
        correlationId: uuid(),
        service: 'features/image-generator',
        createdAt: new Date().toISOString(),
        isDxJob: options.isDxJob,
        tags: {
          type: 'batch',
          batch: batchData.batch_name,
          code: prompt.empty_state_code || prompt.error_type || 'unknown',
          variation: String(variation),
        },
      };

      const jobData: GenerateVariationJobData = {
        variationId,
        name,
        families: options.families,
        generationParams: {
          prompt: prompt.prompt,
          negativePrompt: batchData.negative_prompt,
          seed,
          model: 'anime', // Default model
          inferenceSteps: prompt.steps || 60,
          guidanceScale: prompt.guidance_scale || 8.5,
        },
        _context: context,
      };

      const job = await queue.add('GENERATE_VARIATION', jobData, {
        priority: options.isDxJob ? 2 : 3, // HIGH (2) for DX, NORMAL (3) otherwise
        attempts: 2,
        backoff: {
          type: 'exponential',
          delay: 30000,
        },
        removeOnComplete: true,
        removeOnFail: 100,
      });

      queuedJobs.push(job.id ?? variationId);
      console.log(`  ✓ Queued: ${name} (seed: ${seed})`);
    }
  }

  // Summary
  console.log(`\n=== Complete ===`);
  console.log(`Jobs queued: ${queuedJobs.length}`);
  console.log(`Estimated output: ${formatBytes(estimatedTotalSize)}`);

  // Cleanup
  await queue.close();
  await redisConnection.quit();
}

main().catch((err) => {
  console.error('Error:', err);
  process.exit(1);
});