1import type { CliArgs, Quality } from "../types";
2 
3type DashScopeModelFamily = "qwen2" | "qwenFixed" | "legacy";
4 
5type DashScopeModelSpec = {
6  family: DashScopeModelFamily;
7  defaultSize: string;
8};
9 
10const DEFAULT_MODEL = "qwen-image-2.0-pro";
11const MIN_QWEN_2_TOTAL_PIXELS = 512 * 512;
12const MAX_QWEN_2_TOTAL_PIXELS = 2048 * 2048;
13const SIZE_STEP = 16;
14const QWEN_NEGATIVE_PROMPT =
15  "低分辨率，低画质，肢体畸形，手指畸形，画面过饱和，蜡像感，人脸无细节，过度光滑，画面具有AI感，构图混乱，文字模糊，扭曲";
16 
17const QWEN_2_TARGET_PIXELS: Record<Quality, number> = {
18  normal: 1024 * 1024,
19  "2k": 1536 * 1536,
20};
21 
22const QWEN_2_RECOMMENDED: Record<string, Record<Quality, string>> = {
23  "1:1": { normal: "1024*1024", "2k": "1536*1536" },
24  "2:3": { normal: "768*1152", "2k": "1024*1536" },
25  "3:2": { normal: "1152*768", "2k": "1536*1024" },
26  "3:4": { normal: "960*1280", "2k": "1080*1440" },
27  "4:3": { normal: "1280*960", "2k": "1440*1080" },
28  "9:16": { normal: "720*1280", "2k": "1080*1920" },
29  "16:9": { normal: "1280*720", "2k": "1920*1080" },
30  "21:9": { normal: "1344*576", "2k": "2048*872" },
31};
32 
33const QWEN_FIXED_SIZES_BY_RATIO: Record<string, string> = {
34  "16:9": "1664*928",
35  "4:3": "1472*1104",
36  "1:1": "1328*1328",
37  "3:4": "1104*1472",
38  "9:16": "928*1664",
39};
40 
41const QWEN_FIXED_SIZES = Object.values(QWEN_FIXED_SIZES_BY_RATIO);
42 
43const LEGACY_STANDARD_SIZES: [number, number][] = [
44  [1024, 1024],
45  [1280, 720],
46  [720, 1280],
47  [1024, 768],
48  [768, 1024],
49  [1536, 1024],
50  [1024, 1536],
51  [1536, 864],
52  [864, 1536],
53];
54 
55const LEGACY_STANDARD_SIZES_2K: [number, number][] = [
56  [1536, 1536],
57  [2048, 1152],
58  [1152, 2048],
59  [1536, 1024],
60  [1024, 1536],
61  [1536, 864],
62  [864, 1536],
63  [2048, 2048],
64];
65 
66const QWEN_2_SPEC: DashScopeModelSpec = {
67  family: "qwen2",
68  defaultSize: "1024*1024",
69};
70 
71const QWEN_FIXED_SPEC: DashScopeModelSpec = {
72  family: "qwenFixed",
73  defaultSize: QWEN_FIXED_SIZES_BY_RATIO["16:9"],
74};
75 
76const LEGACY_SPEC: DashScopeModelSpec = {
77  family: "legacy",
78  defaultSize: "1536*1536",
79};
80 
81const MODEL_SPEC_ALIASES: Record<string, DashScopeModelSpec> = {
82  "qwen-image-2.0-pro": QWEN_2_SPEC,
83  "qwen-image-2.0-pro-2026-03-03": QWEN_2_SPEC,
84  "qwen-image-2.0": QWEN_2_SPEC,
85  "qwen-image-2.0-2026-03-03": QWEN_2_SPEC,
86  "qwen-image-max": QWEN_FIXED_SPEC,
87  "qwen-image-max-2025-12-30": QWEN_FIXED_SPEC,
88  "qwen-image-plus": QWEN_FIXED_SPEC,
89  "qwen-image-plus-2026-01-09": QWEN_FIXED_SPEC,
90  "qwen-image": QWEN_FIXED_SPEC,
91};
92 
93export function getDefaultModel(): string {
94  return process.env.DASHSCOPE_IMAGE_MODEL || DEFAULT_MODEL;
95}
96 
97function getApiKey(): string | null {
98  return process.env.DASHSCOPE_API_KEY || null;
99}
100 
101function getBaseUrl(): string {
102  const base = process.env.DASHSCOPE_BASE_URL || "https://dashscope.aliyuncs.com";
103  return base.replace(/\/+$/g, "");
104}
105 
106function getModelSpec(model: string): DashScopeModelSpec {
107  return MODEL_SPEC_ALIASES[model.trim().toLowerCase()] || LEGACY_SPEC;
108}
109 
110export function getModelFamily(model: string): DashScopeModelFamily {
111  return getModelSpec(model).family;
112}
113 
114function normalizeQuality(quality: CliArgs["quality"]): Quality {
115  return quality === "normal" ? "normal" : "2k";
116}
117 
118export function parseAspectRatio(ar: string): { width: number; height: number } | null {
119  const match = ar.match(/^(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)$/);
120  if (!match) return null;
121  const w = parseFloat(match[1]!);
122  const h = parseFloat(match[2]!);
123  if (w <= 0 || h <= 0) return null;
124  return { width: w, height: h };
125}
126 
127export function normalizeSize(size: string): string {
128  return size.replace("x", "*");
129}
130 
131export function parseSize(size: string): { width: number; height: number } | null {
132  const match = normalizeSize(size).match(/^(\d+)\*(\d+)$/);
133  if (!match) return null;
134  const width = Number(match[1]);
135  const height = Number(match[2]);
136  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
137    return null;
138  }
139  return { width, height };
140}
141 
142function formatSize(width: number, height: number): string {
143  return `${width}*${height}`;
144}
145 
146function getRatioValue(ar: string): number | null {
147  const parsed = parseAspectRatio(ar);
148  if (!parsed) return null;
149  return parsed.width / parsed.height;
150}
151 
152function findKnownRatioKey(ar: string, candidates: string[], tolerance = 0.02): string | null {
153  const targetRatio = getRatioValue(ar);
154  if (targetRatio == null) return null;
155 
156  let bestKey: string | null = null;
157  let bestDiff = Infinity;
158 
159  for (const candidate of candidates) {
160    const candidateRatio = getRatioValue(candidate);
161    if (candidateRatio == null) continue;
162    const diff = Math.abs(candidateRatio - targetRatio);
163    if (diff < bestDiff) {
164      bestDiff = diff;
165      bestKey = candidate;
166    }
167  }
168 
169  return bestDiff <= tolerance ? bestKey : null;
170}
171 
172function roundToStep(value: number): number {
173  return Math.max(SIZE_STEP, Math.round(value / SIZE_STEP) * SIZE_STEP);
174}
175 
176function fitToPixelBudget(
177  width: number,
178  height: number,
179  minPixels: number,
180  maxPixels: number,
181): { width: number; height: number } {
182  let nextWidth = width;
183  let nextHeight = height;
184  let pixels = nextWidth * nextHeight;
185 
186  if (pixels > maxPixels) {
187    const scale = Math.sqrt(maxPixels / pixels);
188    nextWidth *= scale;
189    nextHeight *= scale;
190  } else if (pixels < minPixels) {
191    const scale = Math.sqrt(minPixels / pixels);
192    nextWidth *= scale;
193    nextHeight *= scale;
194  }
195 
196  let roundedWidth = roundToStep(nextWidth);
197  let roundedHeight = roundToStep(nextHeight);
198  pixels = roundedWidth * roundedHeight;
199 
200  while (pixels > maxPixels && (roundedWidth > SIZE_STEP || roundedHeight > SIZE_STEP)) {
201    if (roundedWidth >= roundedHeight && roundedWidth > SIZE_STEP) {
202      roundedWidth -= SIZE_STEP;
203    } else if (roundedHeight > SIZE_STEP) {
204      roundedHeight -= SIZE_STEP;
205    } else {
206      break;
207    }
208    pixels = roundedWidth * roundedHeight;
209  }
210 
211  while (pixels < minPixels) {
212    if (roundedWidth <= roundedHeight) {
213      roundedWidth += SIZE_STEP;
214    } else {
215      roundedHeight += SIZE_STEP;
216    }
217    pixels = roundedWidth * roundedHeight;
218  }
219 
220  return { width: roundedWidth, height: roundedHeight };
221}
222 
223export function getSizeFromAspectRatio(ar: string | null, quality: CliArgs["quality"]): string {
224  const normalizedQuality = normalizeQuality(quality);
225  const sizes = normalizedQuality === "2k" ? LEGACY_STANDARD_SIZES_2K : LEGACY_STANDARD_SIZES;
226  const defaultSize = normalizedQuality === "2k" ? "1536*1536" : "1024*1024";
227 
228  if (!ar) return defaultSize;
229 
230  const parsed = parseAspectRatio(ar);
231  if (!parsed) return defaultSize;
232 
233  const targetRatio = parsed.width / parsed.height;
234  let best = defaultSize;
235  let bestDiff = Infinity;
236 
237  for (const [width, height] of sizes) {
238    const diff = Math.abs(width / height - targetRatio);
239    if (diff < bestDiff) {
240      bestDiff = diff;
241      best = formatSize(width, height);
242    }
243  }
244 
245  return best;
246}
247 
248export function getQwen2SizeFromAspectRatio(ar: string | null, quality: CliArgs["quality"]): string {
249  const normalizedQuality = normalizeQuality(quality);
250 
251  if (!ar) {
252    return QWEN_2_RECOMMENDED["1:1"][normalizedQuality];
253  }
254 
255  const recommendedRatio = findKnownRatioKey(ar, Object.keys(QWEN_2_RECOMMENDED));
256  if (recommendedRatio) {
257    return QWEN_2_RECOMMENDED[recommendedRatio][normalizedQuality];
258  }
259 
260  const parsed = parseAspectRatio(ar);
261  if (!parsed) {
262    return QWEN_2_RECOMMENDED["1:1"][normalizedQuality];
263  }
264 
265  const targetRatio = parsed.width / parsed.height;
266  const targetPixels = QWEN_2_TARGET_PIXELS[normalizedQuality];
267  const rawWidth = Math.sqrt(targetPixels * targetRatio);
268  const rawHeight = Math.sqrt(targetPixels / targetRatio);
269  const fitted = fitToPixelBudget(
270    rawWidth,
271    rawHeight,
272    MIN_QWEN_2_TOTAL_PIXELS,
273    MAX_QWEN_2_TOTAL_PIXELS,
274  );
275 
276  return formatSize(fitted.width, fitted.height);
277}
278 
279function getQwenFixedSizeFromAspectRatio(ar: string | null, quality: CliArgs["quality"]): string {
280  if (quality === "normal") {
281    console.warn(
282      "DashScope qwen-image-max/plus/image models use fixed output sizes; --quality normal does not change the generated resolution."
283    );
284  }
285 
286  if (!ar) return QWEN_FIXED_SPEC.defaultSize;
287 
288  const ratioKey = findKnownRatioKey(ar, Object.keys(QWEN_FIXED_SIZES_BY_RATIO));
289  if (!ratioKey) {
290    throw new Error(
291      `DashScope model supports only fixed ratios ${Object.keys(QWEN_FIXED_SIZES_BY_RATIO).join(", ")}. ` +
292      `For custom ratios like "${ar}", use --model qwen-image-2.0-pro.`
293    );
294  }
295 
296  return QWEN_FIXED_SIZES_BY_RATIO[ratioKey]!;
297}
298 
299function validateSizeFormat(size: string): { width: number; height: number } {
300  const parsed = parseSize(size);
301  if (!parsed) {
302    throw new Error(`Invalid DashScope size "${size}". Expected <width>x<height> or <width>*<height>.`);
303  }
304  return parsed;
305}
306 
307function validateQwen2Size(size: string): string {
308  const normalized = normalizeSize(size);
309  const parsed = validateSizeFormat(normalized);
310  const totalPixels = parsed.width * parsed.height;
311  if (totalPixels < MIN_QWEN_2_TOTAL_PIXELS || totalPixels > MAX_QWEN_2_TOTAL_PIXELS) {
312    throw new Error(
313      `DashScope qwen-image-2.0* models require total pixels between ${MIN_QWEN_2_TOTAL_PIXELS} ` +
314      `and ${MAX_QWEN_2_TOTAL_PIXELS}. Received ${normalized} (${totalPixels} pixels).`
315    );
316  }
317  return normalized;
318}
319 
320function validateQwenFixedSize(size: string): string {
321  const normalized = normalizeSize(size);
322  validateSizeFormat(normalized);
323  if (!QWEN_FIXED_SIZES.includes(normalized)) {
324    throw new Error(
325      `DashScope qwen-image-max/plus/image models support only these sizes: ${QWEN_FIXED_SIZES.join(", ")}. ` +
326      `Received ${normalized}.`
327    );
328  }
329  return normalized;
330}
331 
332export function resolveSizeForModel(
333  model: string,
334  args: Pick<CliArgs, "size" | "aspectRatio" | "quality">,
335): string {
336  const spec = getModelSpec(model);
337 
338  if (args.size) {
339    if (spec.family === "qwen2") return validateQwen2Size(args.size);
340    if (spec.family === "qwenFixed") return validateQwenFixedSize(args.size);
341    validateSizeFormat(args.size);
342    return normalizeSize(args.size);
343  }
344 
345  if (spec.family === "qwen2") {
346    return getQwen2SizeFromAspectRatio(args.aspectRatio, args.quality);
347  }
348 
349  if (spec.family === "qwenFixed") {
350    return getQwenFixedSizeFromAspectRatio(args.aspectRatio, args.quality);
351  }
352 
353  return getSizeFromAspectRatio(args.aspectRatio, args.quality);
354}
355 
356function buildParameters(
357  family: DashScopeModelFamily,
358  size: string,
359): Record<string, unknown> {
360  const parameters: Record<string, unknown> = {
361    prompt_extend: false,
362    size,
363  };
364 
365  if (family === "qwen2" || family === "qwenFixed") {
366    parameters.watermark = false;
367    parameters.negative_prompt = QWEN_NEGATIVE_PROMPT;
368  }
369 
370  return parameters;
371}
372 
373type DashScopeResponse = {
374  output?: {
375    result_image?: string;
376    choices?: Array<{
377      message?: {
378        content?: Array<{ image?: string }>;
379      };
380    }>;
381  };
382};
383 
384async function extractImageFromResponse(result: DashScopeResponse): Promise<Uint8Array> {
385  let imageData: string | null = null;
386 
387  if (result.output?.result_image) {
388    imageData = result.output.result_image;
389  } else if (result.output?.choices?.[0]?.message?.content) {
390    const content = result.output.choices[0].message.content;
391    for (const item of content) {
392      if (item.image) {
393        imageData = item.image;
394        break;
395      }
396    }
397  }
398 
399  if (!imageData) {
400    console.error("Response:", JSON.stringify(result, null, 2));
401    throw new Error("No image in response");
402  }
403 
404  if (imageData.startsWith("http://") || imageData.startsWith("https://")) {
405    const imgRes = await fetch(imageData);
406    if (!imgRes.ok) throw new Error("Failed to download image");
407    const buf = await imgRes.arrayBuffer();
408    return new Uint8Array(buf);
409  }
410 
411  return Uint8Array.from(Buffer.from(imageData, "base64"));
412}
413 
414export async function generateImage(
415  prompt: string,
416  model: string,
417  args: CliArgs
418): Promise<Uint8Array> {
419  const apiKey = getApiKey();
420  if (!apiKey) throw new Error("DASHSCOPE_API_KEY is required");
421 
422  if (args.referenceImages.length > 0) {
423    throw new Error(
424      "Reference images are not supported with DashScope provider in baoyu-image-gen. Use --provider google with a Gemini multimodal model."
425    );
426  }
427 
428  const spec = getModelSpec(model);
429  const size = resolveSizeForModel(model, args);
430  const url = `${getBaseUrl()}/api/v1/services/aigc/multimodal-generation/generation`;
431 
432  const body = {
433    model,
434    input: {
435      messages: [
436        {
437          role: "user",
438          content: [{ text: prompt }],
439        },
440      ],
441    },
442    parameters: buildParameters(spec.family, size),
443  };
444 
445  console.log(`Generating image with DashScope (${model})...`, { family: spec.family, size });
446 
447  const res = await fetch(url, {
448    method: "POST",
449    headers: {
450      "Content-Type": "application/json",
451      Authorization: `Bearer ${apiKey}`,
452    },
453    body: JSON.stringify(body),
454  });
455 
456  if (!res.ok) {
457    const err = await res.text();
458    throw new Error(`DashScope API error (${res.status}): ${err}`);
459  }
460 
461  const result = await res.json() as DashScopeResponse;
462  return extractImageFromResponse(result);
463}
464