1import type { CliArgs, Quality } from "../types";
2 
3type ZaiModelFamily = "glm" | "legacy";
4 
5type ZaiRequestBody = {
6  model: string;
7  prompt: string;
8  quality: "hd" | "standard";
9  size: string;
10};
11 
12type ZaiResponse = {
13  data?: Array<{ url?: string }>;
14};
15 
16const DEFAULT_MODEL = "glm-image";
17const GLM_MAX_PIXELS = 2 ** 22;
18const LEGACY_MAX_PIXELS = 2 ** 21;
19const GLM_SIZE_STEP = 32;
20const LEGACY_SIZE_STEP = 16;
21 
22const GLM_RECOMMENDED_SIZES: Record<string, string> = {
23  "1:1": "1280x1280",
24  "3:2": "1568x1056",
25  "2:3": "1056x1568",
26  "4:3": "1472x1088",
27  "3:4": "1088x1472",
28  "16:9": "1728x960",
29  "9:16": "960x1728",
30};
31 
32const LEGACY_RECOMMENDED_SIZES: Record<string, string> = {
33  "1:1": "1024x1024",
34  "9:16": "768x1344",
35  "3:4": "864x1152",
36  "16:9": "1344x768",
37  "4:3": "1152x864",
38  "2:1": "1440x720",
39  "1:2": "720x1440",
40};
41 
42export function getDefaultModel(): string {
43  return process.env.ZAI_IMAGE_MODEL || process.env.BIGMODEL_IMAGE_MODEL || DEFAULT_MODEL;
44}
45 
46function getApiKey(): string | null {
47  return process.env.ZAI_API_KEY || process.env.BIGMODEL_API_KEY || null;
48}
49 
50export function buildZaiUrl(): string {
51  const base = (process.env.ZAI_BASE_URL || process.env.BIGMODEL_BASE_URL || "https://api.z.ai/api/paas/v4")
52    .replace(/\/+$/g, "");
53  if (base.endsWith("/images/generations")) return base;
54  if (base.endsWith("/api/paas/v4")) return `${base}/images/generations`;
55  if (base.endsWith("/v4")) return `${base}/images/generations`;
56  return `${base}/api/paas/v4/images/generations`;
57}
58 
59export function getModelFamily(model: string): ZaiModelFamily {
60  return model.trim().toLowerCase() === "glm-image" ? "glm" : "legacy";
61}
62 
63export function parseAspectRatio(ar: string): { width: number; height: number } | null {
64  const match = ar.match(/^(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)$/);
65  if (!match) return null;
66  const width = Number(match[1]);
67  const height = Number(match[2]);
68  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
69    return null;
70  }
71  return { width, height };
72}
73 
74export function parseSize(size: string): { width: number; height: number } | null {
75  const match = size.trim().match(/^(\d+)\s*[xX*]\s*(\d+)$/);
76  if (!match) return null;
77  const width = parseInt(match[1]!, 10);
78  const height = parseInt(match[2]!, 10);
79  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
80    return null;
81  }
82  return { width, height };
83}
84 
85function formatSize(width: number, height: number): string {
86  return `${width}x${height}`;
87}
88 
89function roundToStep(value: number, step: number): number {
90  return Math.max(step, Math.round(value / step) * step);
91}
92 
93function getRatioValue(ar: string): number | null {
94  const parsed = parseAspectRatio(ar);
95  if (!parsed) return null;
96  return parsed.width / parsed.height;
97}
98 
99function findClosestRatioKey(ar: string, candidates: string[]): string | null {
100  const targetRatio = getRatioValue(ar);
101  if (targetRatio == null) return null;
102 
103  let bestKey: string | null = null;
104  let bestDiff = Infinity;
105  for (const candidate of candidates) {
106    const candidateRatio = getRatioValue(candidate);
107    if (candidateRatio == null) continue;
108    const diff = Math.abs(candidateRatio - targetRatio);
109    if (diff < bestDiff) {
110      bestDiff = diff;
111      bestKey = candidate;
112    }
113  }
114 
115  return bestDiff <= 0.05 ? bestKey : null;
116}
117 
118function getTargetPixels(quality: Quality): number {
119  return quality === "normal" ? 1024 * 1024 : 1536 * 1536;
120}
121 
122function fitToPixelBudget(
123  width: number,
124  height: number,
125  targetPixels: number,
126  maxPixels: number,
127  step: number,
128): { width: number; height: number } {
129  let nextWidth = width;
130  let nextHeight = height;
131  const pixels = nextWidth * nextHeight;
132 
133  if (pixels > maxPixels) {
134    const scale = Math.sqrt(maxPixels / pixels);
135    nextWidth *= scale;
136    nextHeight *= scale;
137  } else {
138    const scale = Math.sqrt(targetPixels / pixels);
139    nextWidth *= scale;
140    nextHeight *= scale;
141  }
142 
143  let roundedWidth = roundToStep(nextWidth, step);
144  let roundedHeight = roundToStep(nextHeight, step);
145  let roundedPixels = roundedWidth * roundedHeight;
146 
147  while (roundedPixels > maxPixels && (roundedWidth > step || roundedHeight > step)) {
148    if (roundedWidth >= roundedHeight && roundedWidth > step) {
149      roundedWidth -= step;
150    } else if (roundedHeight > step) {
151      roundedHeight -= step;
152    } else {
153      break;
154    }
155    roundedPixels = roundedWidth * roundedHeight;
156  }
157 
158  return { width: roundedWidth, height: roundedHeight };
159}
160 
161function validateCustomSize(
162  size: string,
163  family: ZaiModelFamily,
164): string {
165  const parsed = parseSize(size);
166  if (!parsed) {
167    throw new Error("Z.AI --size must be in WxH format, for example 1280x1280.");
168  }
169 
170  const widthStep = family === "glm" ? GLM_SIZE_STEP : LEGACY_SIZE_STEP;
171  const minEdge = family === "glm" ? 1024 : 512;
172  const maxPixels = family === "glm" ? GLM_MAX_PIXELS : LEGACY_MAX_PIXELS;
173 
174  if (parsed.width < minEdge || parsed.width > 2048 || parsed.height < minEdge || parsed.height > 2048) {
175    throw new Error(
176      family === "glm"
177        ? "GLM-image custom size requires width and height between 1024 and 2048."
178        : "Z.AI legacy image models require width and height between 512 and 2048."
179    );
180  }
181 
182  if (parsed.width % widthStep !== 0 || parsed.height % widthStep !== 0) {
183    throw new Error(
184      family === "glm"
185        ? "GLM-image custom size requires width and height divisible by 32."
186        : "Z.AI legacy image models require width and height divisible by 16."
187    );
188  }
189 
190  if (parsed.width * parsed.height > maxPixels) {
191    throw new Error(
192      family === "glm"
193        ? "GLM-image custom size must not exceed 2^22 total pixels."
194        : "Z.AI legacy image size must not exceed 2^21 total pixels."
195    );
196  }
197 
198  return formatSize(parsed.width, parsed.height);
199}
200 
201export function resolveSizeForModel(
202  model: string,
203  args: Pick<CliArgs, "size" | "aspectRatio" | "quality">,
204): string {
205  const family = getModelFamily(model);
206  const quality = args.quality === "normal" ? "normal" : "2k";
207 
208  if (args.size) {
209    return validateCustomSize(args.size, family);
210  }
211 
212  const recommended = family === "glm" ? GLM_RECOMMENDED_SIZES : LEGACY_RECOMMENDED_SIZES;
213  const defaultSize = family === "glm" ? "1280x1280" : "1024x1024";
214 
215  if (!args.aspectRatio) return defaultSize;
216 
217  const recommendedRatio = findClosestRatioKey(args.aspectRatio, Object.keys(recommended));
218  if (recommendedRatio) {
219    return recommended[recommendedRatio]!;
220  }
221 
222  const parsedRatio = parseAspectRatio(args.aspectRatio);
223  if (!parsedRatio) return defaultSize;
224 
225  const targetPixels = getTargetPixels(quality);
226  const maxPixels = family === "glm" ? GLM_MAX_PIXELS : LEGACY_MAX_PIXELS;
227  const step = family === "glm" ? GLM_SIZE_STEP : LEGACY_SIZE_STEP;
228  const fit = fitToPixelBudget(
229    parsedRatio.width,
230    parsedRatio.height,
231    targetPixels,
232    maxPixels,
233    step,
234  );
235  return formatSize(fit.width, fit.height);
236}
237 
238function getZaiQuality(quality: CliArgs["quality"]): "hd" | "standard" {
239  return quality === "normal" ? "standard" : "hd";
240}
241 
242export function validateArgs(_model: string, args: CliArgs): void {
243  if (args.referenceImages.length > 0) {
244    throw new Error("Z.AI GLM-image currently supports text-to-image only in baoyu-image-gen. Remove --ref or choose another provider.");
245  }
246 
247  if (args.n > 1) {
248    throw new Error("Z.AI image generation currently returns a single image per request in baoyu-image-gen.");
249  }
250}
251 
252export function buildRequestBody(
253  prompt: string,
254  model: string,
255  args: CliArgs,
256): ZaiRequestBody {
257  validateArgs(model, args);
258  return {
259    model,
260    prompt,
261    quality: getZaiQuality(args.quality),
262    size: resolveSizeForModel(model, args),
263  };
264}
265 
266export async function extractImageFromResponse(result: ZaiResponse): Promise<Uint8Array> {
267  const url = result.data?.[0]?.url;
268  if (!url) {
269    throw new Error("No image URL in Z.AI response");
270  }
271 
272  const imageResponse = await fetch(url);
273  if (!imageResponse.ok) {
274    throw new Error(`Failed to download image from Z.AI: ${imageResponse.status}`);
275  }
276 
277  return new Uint8Array(await imageResponse.arrayBuffer());
278}
279 
280export async function generateImage(
281  prompt: string,
282  model: string,
283  args: CliArgs,
284): Promise<Uint8Array> {
285  const apiKey = getApiKey();
286  if (!apiKey) {
287    throw new Error("ZAI_API_KEY is required. Get one from https://docs.z.ai/.");
288  }
289 
290  const response = await fetch(buildZaiUrl(), {
291    method: "POST",
292    headers: {
293      "Content-Type": "application/json",
294      Authorization: `Bearer ${apiKey}`,
295    },
296    body: JSON.stringify(buildRequestBody(prompt, model, args)),
297  });
298 
299  if (!response.ok) {
300    const err = await response.text();
301    throw new Error(`Z.AI API error (${response.status}): ${err}`);
302  }
303 
304  const result = (await response.json()) as ZaiResponse;
305  return extractImageFromResponse(result);
306}
307