1import path from "node:path";
2import process from "node:process";
3import { readdir, readFile, writeFile } from "node:fs/promises";
4 
5type CliArgs = {
6  outlinePath: string | null;
7  promptsDir: string | null;
8  outputPath: string | null;
9  imagesDir: string | null;
10  refsDir: string;
11  provider: string;
12  model: string | null;
13  aspectRatio: string;
14  quality: string;
15  jobs: number | null;
16  help: boolean;
17};
18 
19type OutlineEntry = {
20  index: number;
21  filename: string;
22};
23 
24type PromptReference = {
25  filename: string;
26  usage: "direct" | "style" | "palette";
27};
28 
29function printUsage(): void {
30  console.log(`Usage:
31  bun <baseDir>/scripts/build-batch.ts --outline outline.md --prompts prompts --output batch.json --images-dir attachments
32  npx -y tsx <baseDir>/scripts/build-batch.ts --outline outline.md --prompts prompts --output batch.json --images-dir attachments
33 
34Options:
35  --outline <path>     Path to outline.md
36  --prompts <path>     Path to prompts directory
37  --output <path>      Path to output batch.json
38  --images-dir <path>  Directory for generated images
39  --refs-dir <path>    Directory holding reference images, relative to batch file (default: references)
40  --provider <name>    Provider for baoyu-image-gen batch tasks (default: replicate)
41  --model <id>         Explicit model for baoyu-image-gen batch tasks (default: resolved by baoyu-image-gen config/env)
42  --ar <ratio>         Aspect ratio for all tasks (default: 16:9)
43  --quality <level>    Quality for all tasks (default: 2k)
44  --jobs <count>       Recommended worker count metadata (optional)
45  -h, --help           Show help`);
46}
47 
48function parseArgs(argv: string[]): CliArgs {
49  const args: CliArgs = {
50    outlinePath: null,
51    promptsDir: null,
52    outputPath: null,
53    imagesDir: null,
54    refsDir: "references",
55    provider: "replicate",
56    model: null,
57    aspectRatio: "16:9",
58    quality: "2k",
59    jobs: null,
60    help: false,
61  };
62 
63  for (let i = 0; i < argv.length; i++) {
64    const current = argv[i]!;
65    if (current === "--outline") args.outlinePath = argv[++i] ?? null;
66    else if (current === "--prompts") args.promptsDir = argv[++i] ?? null;
67    else if (current === "--output") args.outputPath = argv[++i] ?? null;
68    else if (current === "--images-dir") args.imagesDir = argv[++i] ?? null;
69    else if (current === "--refs-dir") args.refsDir = argv[++i] ?? args.refsDir;
70    else if (current === "--provider") args.provider = argv[++i] ?? args.provider;
71    else if (current === "--model") args.model = argv[++i] ?? args.model;
72    else if (current === "--ar") args.aspectRatio = argv[++i] ?? args.aspectRatio;
73    else if (current === "--quality") args.quality = argv[++i] ?? args.quality;
74    else if (current === "--jobs") {
75      const value = argv[++i];
76      args.jobs = value ? parseInt(value, 10) : null;
77    } else if (current === "--help" || current === "-h") {
78      args.help = true;
79    }
80  }
81  return args;
82}
83 
84function parsePromptReferences(content: string): PromptReference[] {
85  const fmMatch = content.match(/^---\s*\n([\s\S]*?)\n---\s*(?:\n|$)/);
86  if (!fmMatch) return [];
87  const lines = fmMatch[1]!.split(/\r?\n/);
88 
89  const refs: PromptReference[] = [];
90  let current: Partial<PromptReference> | null = null;
91  let inReferences = false;
92  let listIndent = 0;
93 
94  const flush = () => {
95    if (current?.filename) {
96      refs.push({
97        filename: current.filename,
98        usage: (current.usage ?? "direct") as PromptReference["usage"],
99      });
100    }
101    current = null;
102  };
103 
104  const unquote = (raw: string): string => raw.trim().replace(/^["']|["']$/g, "");
105 
106  for (const line of lines) {
107    if (!line.trim() || line.trim().startsWith("#")) continue;
108 
109    const keyMatch = line.match(/^(\S[^:]*):\s*(.*)$/);
110    if (keyMatch) {
111      flush();
112      if (keyMatch[1] === "references") {
113        inReferences = true;
114        listIndent = 0;
115        continue;
116      }
117      inReferences = false;
118      continue;
119    }
120 
121    if (!inReferences) continue;
122 
123    const itemMatch = line.match(/^(\s*)-\s*(.*)$/);
124    if (itemMatch) {
125      flush();
126      listIndent = itemMatch[1]!.length;
127      current = {};
128      const rest = itemMatch[2]!.trim();
129      if (rest) {
130        const kv = rest.match(/^(\w+)\s*:\s*(.*)$/);
131        if (kv && (kv[1] === "filename" || kv[1] === "usage")) {
132          (current as Record<string, string>)[kv[1]] = unquote(kv[2]!);
133        }
134      }
135      continue;
136    }
137 
138    const kvMatch = line.match(/^(\s+)(\w+)\s*:\s*(.*)$/);
139    if (kvMatch && kvMatch[1]!.length > listIndent && current) {
140      if (kvMatch[2] === "filename" || kvMatch[2] === "usage") {
141        (current as Record<string, string>)[kvMatch[2]!] = unquote(kvMatch[3]!);
142      }
143    }
144  }
145  flush();
146 
147  return refs;
148}
149 
150function parseOutline(content: string): OutlineEntry[] {
151  const entries: OutlineEntry[] = [];
152  const blocks = content.split(/^## Illustration\s+/m).slice(1);
153 
154  for (const block of blocks) {
155    const indexMatch = block.match(/^(\d+)/);
156    const filenameMatch = block.match(/\*\*Filename\*\*:\s*(.+)/);
157    if (indexMatch && filenameMatch) {
158      entries.push({
159        index: parseInt(indexMatch[1]!, 10),
160        filename: filenameMatch[1]!.trim(),
161      });
162    }
163  }
164  return entries;
165}
166 
167async function findPromptFile(promptsDir: string, entry: OutlineEntry): Promise<string | null> {
168  const files = await readdir(promptsDir);
169  const prefix = String(entry.index).padStart(2, "0");
170  const match = files.find((f) => f.startsWith(prefix) && f.endsWith(".md"));
171  return match ? path.join(promptsDir, match) : null;
172}
173 
174async function main(): Promise<void> {
175  const args = parseArgs(process.argv.slice(2));
176  if (args.help) {
177    printUsage();
178    return;
179  }
180 
181  if (!args.outlinePath) {
182    console.error("Error: --outline is required");
183    process.exit(1);
184  }
185  if (!args.promptsDir) {
186    console.error("Error: --prompts is required");
187    process.exit(1);
188  }
189  if (!args.outputPath) {
190    console.error("Error: --output is required");
191    process.exit(1);
192  }
193 
194  const outlineContent = await readFile(args.outlinePath, "utf8");
195  const entries = parseOutline(outlineContent);
196 
197  if (entries.length === 0) {
198    console.error("No illustration entries found in outline.");
199    process.exit(1);
200  }
201 
202  const tasks = [];
203  for (const entry of entries) {
204    const promptFile = await findPromptFile(args.promptsDir, entry);
205    if (!promptFile) {
206      console.error(`Warning: No prompt file found for illustration ${entry.index}, skipping.`);
207      continue;
208    }
209 
210    const imageDir = args.imagesDir ?? path.dirname(args.outputPath);
211    const promptContent = await readFile(promptFile, "utf8");
212    const refs = parsePromptReferences(promptContent)
213      .filter((r) => r.usage === "direct")
214      .map((r) => path.posix.join(args.refsDir, r.filename));
215 
216    const task: Record<string, unknown> = {
217      id: `illustration-${String(entry.index).padStart(2, "0")}`,
218      promptFiles: [promptFile],
219      image: path.join(imageDir, entry.filename),
220      provider: args.provider,
221      ar: args.aspectRatio,
222      quality: args.quality,
223    };
224    if (args.model) task.model = args.model;
225    if (refs.length > 0) task.ref = refs;
226    tasks.push(task);
227  }
228 
229  const output: Record<string, unknown> = { tasks };
230  if (args.jobs) output.jobs = args.jobs;
231 
232  await writeFile(args.outputPath, JSON.stringify(output, null, 2) + "\n");
233  console.log(`Batch file written: ${args.outputPath} (${tasks.length} tasks)`);
234}
235 
236main().catch((error) => {
237  console.error(error instanceof Error ? error.message : String(error));
238  process.exit(1);
239});
240