1import path from "node:path";
2import { readFile } from "node:fs/promises";
3 
4import type { CliArgs } from "../types";
5 
6export type SeedreamModelFamily =
7  | "seedream5"
8  | "seedream45"
9  | "seedream40"
10  | "seedream30"
11  | "unknown";
12 
13type SeedreamRequestImage = string | string[];
14 
15type SeedreamRequestBody = {
16  model: string;
17  prompt: string;
18  size: string;
19  response_format: "url";
20  watermark: boolean;
21  image?: SeedreamRequestImage;
22  output_format?: "png";
23};
24 
25type SeedreamImageResponse = {
26  model?: string;
27  created?: number;
28  data?: Array<{
29    url?: string;
30    b64_json?: string;
31    size?: string;
32    error?: {
33      code?: string;
34      message?: string;
35    };
36  }>;
37  usage?: {
38    generated_images: number;
39    output_tokens: number;
40    total_tokens: number;
41  };
42  error?: {
43    code?: string;
44    message?: string;
45  };
46};
47 
48export function getDefaultModel(): string {
49  return process.env.SEEDREAM_IMAGE_MODEL || "doubao-seedream-5-0-260128";
50}
51 
52function getApiKey(): string | null {
53  return process.env.ARK_API_KEY || null;
54}
55 
56function getBaseUrl(): string {
57  return process.env.SEEDREAM_BASE_URL || "https://ark.cn-beijing.volces.com/api/v3";
58}
59 
60function parsePixelSize(value: string): { width: number; height: number } | null {
61  const match = value.trim().match(/^(\d+)\s*[xX]\s*(\d+)$/);
62  if (!match) return null;
63 
64  const width = parseInt(match[1]!, 10);
65  const height = parseInt(match[2]!, 10);
66  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
67    return null;
68  }
69 
70  return { width, height };
71}
72 
73function normalizePixelSize(value: string): string | null {
74  const parsed = parsePixelSize(value);
75  if (!parsed) return null;
76  return `${parsed.width}x${parsed.height}`;
77}
78 
79function normalizeSizePreset(value: string): string | null {
80  const upper = value.trim().toUpperCase();
81  if (upper === "ADAPTIVE") return "adaptive";
82  if (upper === "1K" || upper === "2K" || upper === "3K" || upper === "4K") return upper;
83  return null;
84}
85 
86function normalizeSizeValue(value: string): string | null {
87  return normalizeSizePreset(value) ?? normalizePixelSize(value);
88}
89 
90function getMimeType(filename: string): string {
91  const ext = path.extname(filename).toLowerCase();
92  if (ext === ".jpg" || ext === ".jpeg") return "image/jpeg";
93  if (ext === ".webp") return "image/webp";
94  if (ext === ".gif") return "image/gif";
95  if (ext === ".bmp") return "image/bmp";
96  if (ext === ".tiff" || ext === ".tif") return "image/tiff";
97  return "image/png";
98}
99 
100async function readImageAsDataUrl(filePath: string): Promise<string> {
101  const bytes = await readFile(filePath);
102  return `data:${getMimeType(filePath)};base64,${bytes.toString("base64")}`;
103}
104 
105export function getModelFamily(model: string): SeedreamModelFamily {
106  const normalized = model.trim();
107  if (/^doubao-seedream-5-0(?:-lite)?-\d+$/.test(normalized)) return "seedream5";
108  if (/^doubao-seedream-4-5-\d+$/.test(normalized)) return "seedream45";
109  if (/^doubao-seedream-4-0-\d+$/.test(normalized)) return "seedream40";
110  if (/^doubao-seedream-3-0-t2i-\d+$/.test(normalized)) return "seedream30";
111  return "unknown";
112}
113 
114function isRemovedSeededitModel(model: string): boolean {
115  return /^doubao-seededit-3-0-i2i-\d+$/.test(model.trim());
116}
117 
118function assertSupportedModel(model: string): void {
119  if (isRemovedSeededitModel(model)) {
120    throw new Error(
121      `${model} is no longer supported. SeedEdit 3.0 support has been removed from this tool; use Seedream 5.0/4.5/4.0/3.0 instead.`
122    );
123  }
124}
125 
126export function supportsReferenceImages(model: string): boolean {
127  const family = getModelFamily(model);
128  return family === "seedream5" || family === "seedream45" || family === "seedream40";
129}
130 
131function supportsOutputFormat(model: string): boolean {
132  return getModelFamily(model) === "seedream5";
133}
134 
135export function getDefaultOutputExtension(model: string): ".png" | ".jpg" {
136  assertSupportedModel(model);
137  return supportsOutputFormat(model) ? ".png" : ".jpg";
138}
139 
140export function getDefaultSeedreamSize(model: string, args: CliArgs): string {
141  assertSupportedModel(model);
142  const family = getModelFamily(model);
143 
144  if (family === "seedream5") return "2K";
145  if (family === "seedream45") return "2K";
146  if (family === "seedream40") return args.quality === "normal" ? "1K" : "2K";
147  if (family === "seedream30") return args.quality === "2k" ? "2048x2048" : "1024x1024";
148  return "2K";
149}
150 
151export function resolveSeedreamSize(model: string, args: CliArgs): string {
152  assertSupportedModel(model);
153  const family = getModelFamily(model);
154  const requested = args.size || args.imageSize || null;
155  const normalized = requested ? normalizeSizeValue(requested) : null;
156 
157  if (!normalized) {
158    return getDefaultSeedreamSize(model, args);
159  }
160 
161  if (family === "seedream30") {
162    const pixelSize = normalizePixelSize(normalized);
163    if (!pixelSize) {
164      throw new Error("Seedream 3.0 only supports explicit WxH sizes such as 1024x1024.");
165    }
166    return pixelSize;
167  }
168 
169  if (family === "seedream5") {
170    if (normalized === "4K" || normalized === "1K" || normalized === "adaptive") {
171      throw new Error("Seedream 5.0 only supports 2K, 3K, or explicit WxH sizes.");
172    }
173    return normalized;
174  }
175 
176  if (family === "seedream45") {
177    if (normalized === "1K" || normalized === "3K" || normalized === "adaptive") {
178      throw new Error("Seedream 4.5 only supports 2K, 4K, or explicit WxH sizes.");
179    }
180    return normalized;
181  }
182 
183  if (family === "seedream40") {
184    if (normalized === "3K" || normalized === "adaptive") {
185      throw new Error("Seedream 4.0 only supports 1K, 2K, 4K, or explicit WxH sizes.");
186    }
187    return normalized;
188  }
189 
190  if (normalized === "adaptive") {
191    throw new Error("Adaptive size is not supported by Seedream image generation.");
192  }
193 
194  if (normalized === "1K" || normalized === "3K" || normalized === "4K") {
195    throw new Error(
196      "Unknown Seedream model ID. Use a documented model ID or pass an explicit WxH size instead of preset imageSize."
197    );
198  }
199 
200  return normalized;
201}
202 
203export function validateArgs(model: string, args: CliArgs): void {
204  assertSupportedModel(model);
205  const family = getModelFamily(model);
206  const refCount = args.referenceImages.length;
207 
208  if (refCount === 0) {
209    resolveSeedreamSize(model, args);
210    return;
211  }
212 
213  if (family === "unknown") {
214    throw new Error(
215      "Reference images with Seedream require a known model ID. Use Seedream 5.0/4.5/4.0 model IDs instead of an endpoint ID."
216    );
217  }
218 
219  if (!supportsReferenceImages(model)) {
220    throw new Error(`${model} does not support reference images.`);
221  }
222 
223  if ((family === "seedream5" || family === "seedream45" || family === "seedream40") && refCount > 14) {
224    throw new Error(`${model} supports at most 14 reference images.`);
225  }
226 
227  resolveSeedreamSize(model, args);
228}
229 
230export async function buildImageInput(
231  model: string,
232  referenceImages: string[],
233): Promise<SeedreamRequestImage | undefined> {
234  if (referenceImages.length === 0) return undefined;
235  assertSupportedModel(model);
236 
237  const encoded = await Promise.all(referenceImages.map((refPath) => readImageAsDataUrl(refPath)));
238 
239  return encoded.length === 1 ? encoded[0]! : encoded;
240}
241 
242export function buildRequestBody(
243  prompt: string,
244  model: string,
245  args: CliArgs,
246  imageInput?: SeedreamRequestImage,
247): SeedreamRequestBody {
248  validateArgs(model, args);
249 
250  const requestBody: SeedreamRequestBody = {
251    model,
252    prompt,
253    size: resolveSeedreamSize(model, args),
254    response_format: "url",
255    watermark: false,
256  };
257 
258  if (imageInput) {
259    requestBody.image = imageInput;
260  }
261 
262  if (supportsOutputFormat(model)) {
263    requestBody.output_format = "png";
264  }
265 
266  return requestBody;
267}
268 
269async function downloadImage(url: string): Promise<Uint8Array> {
270  const imgResponse = await fetch(url);
271  if (!imgResponse.ok) {
272    throw new Error(`Failed to download image from ${url}`);
273  }
274 
275  const buffer = await imgResponse.arrayBuffer();
276  return new Uint8Array(buffer);
277}
278 
279export async function extractImageFromResponse(result: SeedreamImageResponse): Promise<Uint8Array> {
280  const first = result.data?.find((item) => item.url || item.b64_json || item.error);
281 
282  if (!first) {
283    throw new Error("No image data in Seedream response");
284  }
285 
286  if (first.error) {
287    throw new Error(first.error.message || "Seedream returned an image generation error");
288  }
289 
290  if (first.b64_json) {
291    return Uint8Array.from(Buffer.from(first.b64_json, "base64"));
292  }
293 
294  if (first.url) {
295    console.error(`Downloading image from ${first.url}...`);
296    return downloadImage(first.url);
297  }
298 
299  throw new Error("No image URL or base64 data in Seedream response");
300}
301 
302export async function generateImage(
303  prompt: string,
304  model: string,
305  args: CliArgs,
306): Promise<Uint8Array> {
307  const apiKey = getApiKey();
308  if (!apiKey) {
309    throw new Error(
310      "ARK_API_KEY is required. " +
311        "Get your API key from https://console.volcengine.com/ark"
312    );
313  }
314 
315  validateArgs(model, args);
316  const imageInput = await buildImageInput(model, args.referenceImages);
317  const requestBody = buildRequestBody(prompt, model, args, imageInput);
318 
319  console.error(`Calling Seedream API (${model}) with size: ${requestBody.size}`);
320 
321  const response = await fetch(`${getBaseUrl()}/images/generations`, {
322    method: "POST",
323    headers: {
324      "Content-Type": "application/json",
325      Authorization: `Bearer ${apiKey}`,
326    },
327    body: JSON.stringify(requestBody),
328  });
329 
330  if (!response.ok) {
331    const err = await response.text();
332    throw new Error(`Seedream API error (${response.status}): ${err}`);
333  }
334 
335  const result = (await response.json()) as SeedreamImageResponse;
336  if (result.error) {
337    throw new Error(result.error.message || "Seedream API returned an error");
338  }
339 
340  return extractImageFromResponse(result);
341}
342