1import type { CliArgs } from "../types";
2import * as crypto from "node:crypto";
3 
4type JimengSizePreset = "normal" | "2k" | "4k";
5 
6export function getDefaultModel(): string {
7  return process.env.JIMENG_IMAGE_MODEL || "jimeng_t2i_v40";
8}
9 
10function getAccessKey(): string | null {
11  return process.env.JIMENG_ACCESS_KEY_ID || null;
12}
13 
14function getSecretKey(): string | null {
15  return process.env.JIMENG_SECRET_ACCESS_KEY || null;
16}
17 
18function getRegion(): string {
19  return process.env.JIMENG_REGION || "cn-north-1";
20}
21 
22function getBaseUrl(): string {
23  return process.env.JIMENG_BASE_URL || "https://visual.volcengineapi.com";
24}
25 
26function resolveEndpoint(query: Record<string, string>): {
27  url: string;
28  host: string;
29  canonicalUri: string;
30} {
31  let baseUrl: URL;
32  try {
33    baseUrl = new URL(getBaseUrl());
34  } catch {
35    throw new Error(`Invalid JIMENG_BASE_URL: ${getBaseUrl()}`);
36  }
37 
38  baseUrl.search = "";
39  for (const [key, value] of Object.entries(query).sort(([a], [b]) => a.localeCompare(b))) {
40    baseUrl.searchParams.set(key, value);
41  }
42 
43  return {
44    url: baseUrl.toString(),
45    host: baseUrl.host,
46    canonicalUri: baseUrl.pathname || "/",
47  };
48}
49 
50/**
51 * Volcengine HMAC-SHA256 signature generation
52 * Following the official documentation at:
53 * https://www.volcengine.com/docs/85621/1817045
54 */
55function generateSignature(
56  method: string,
57  query: Record<string, string>,
58  headers: Record<string, string>,
59  body: string,
60  accessKey: string,
61  secretKey: string,
62  region: string,
63  service: string,
64  canonicalUri: string
65): string {
66  // 1. Create canonical request
67  // Sort query parameters alphabetically
68  const sortedQuery = Object.entries(query)
69    .sort(([a], [b]) => a.localeCompare(b))
70    .map(([k, v]) => `${encodeURIComponent(k)}=${encodeURIComponent(v)}`)
71    .join("&");
72 
73  // Sort headers alphabetically and create canonical headers
74  const sortedHeaders = Object.entries(headers)
75    .sort(([a], [b]) => a.localeCompare(b))
76    .map(([k, v]) => `${k.toLowerCase()}:${v.trim()}\n`)
77    .join("");
78 
79  const signedHeaders = Object.keys(headers)
80    .sort()
81    .map(k => k.toLowerCase())
82    .join(";");
83 
84  const hashedPayload = crypto.createHash("sha256").update(body, "utf8").digest("hex");
85 
86  const canonicalRequest = [
87    method,
88    canonicalUri,
89    sortedQuery,
90    sortedHeaders,
91    signedHeaders,
92    hashedPayload,
93  ].join("\n");
94 
95  const hashedCanonicalRequest = crypto
96    .createHash("sha256")
97    .update(canonicalRequest, "utf8")
98    .digest("hex");
99 
100  // 2. Create string to sign
101  const algorithm = "HMAC-SHA256";
102  const timestamp = headers["X-Date"] || headers["x-date"];
103  if (!timestamp) {
104    throw new Error("Jimeng signature generation requires an X-Date header.");
105  }
106  const dateStamp = timestamp.slice(0, 8);
107 
108  const credentialScope = `${dateStamp}/${region}/${service}/request`;
109 
110  const stringToSign = [
111    algorithm,
112    timestamp,
113    credentialScope,
114    hashedCanonicalRequest,
115  ].join("\n");
116 
117  // 3. Calculate signature
118  const kDate = crypto
119    .createHmac("sha256", secretKey)
120    .update(dateStamp)
121    .digest();
122 
123  const kRegion = crypto.createHmac("sha256", kDate).update(region).digest();
124  const kService = crypto.createHmac("sha256", kRegion).update(service).digest();
125  const kSigning = crypto.createHmac("sha256", kService).update("request").digest();
126 
127  const signature = crypto
128    .createHmac("sha256", kSigning)
129    .update(stringToSign)
130    .digest("hex");
131 
132  // 4. Create authorization header
133  return `${algorithm} Credential=${accessKey}/${credentialScope}, SignedHeaders=${signedHeaders}, Signature=${signature}`;
134}
135 
136/**
137 * Parse aspect ratio string like "16:9", "1:1", "4:3" into width and height
138 */
139function parseAspectRatio(ar: string): { width: number; height: number } | null {
140  const match = ar.match(/^(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)$/);
141  if (!match) return null;
142  const w = parseFloat(match[1]!);
143  const h = parseFloat(match[2]!);
144  if (w <= 0 || h <= 0) return null;
145  return { width: w, height: h };
146}
147 
148/**
149 * Supported size presets for different quality levels
150 * Based on Volcengine Jimeng documentation
151 */
152const SIZE_PRESETS: Record<string, Record<string, string>> = {
153  normal: {
154    "1:1": "1024x1024",
155    "4:3": "1360x1020",
156    "16:9": "1536x864",
157    "3:2": "1440x960",
158    "21:9": "1920x824",
159  },
160  "2k": {
161    "1:1": "2048x2048",
162    "4:3": "2304x1728",
163    "16:9": "2560x1440",
164    "3:2": "2496x1664",
165    "21:9": "3024x1296",
166  },
167  "4k": {
168    "1:1": "4096x4096",
169    "4:3": "4694x3520",
170    "16:9": "5404x3040",
171    "3:2": "4992x3328",
172    "21:9": "6198x2656",
173  },
174};
175 
176function normalizeDimensions(value: string): string | null {
177  const match = value.trim().match(/^(\d+)\s*[xX*]\s*(\d+)$/);
178  if (!match) return null;
179  return `${match[1]}x${match[2]}`;
180}
181 
182function getClosestPresetSize(ar: string | null, qualityLevel: JimengSizePreset): string {
183  const presets = SIZE_PRESETS[qualityLevel];
184  const defaultSize = presets["1:1"]!;
185 
186  if (!ar) return defaultSize;
187 
188  const parsed = parseAspectRatio(ar);
189  if (!parsed) return defaultSize;
190 
191  const targetRatio = parsed.width / parsed.height;
192  let bestMatch = defaultSize;
193  let bestDiff = Infinity;
194 
195  for (const [ratio, size] of Object.entries(presets)) {
196    const [w, h] = ratio.split(":").map(Number);
197    const presetRatio = w / h;
198    const diff = Math.abs(presetRatio - targetRatio);
199    if (diff < bestDiff) {
200      bestDiff = diff;
201      bestMatch = size;
202    }
203  }
204 
205  return bestMatch;
206}
207 
208function normalizeImageSizePreset(imageSize: string, ar: string | null): string | null {
209  const preset = imageSize.trim().toUpperCase();
210  if (preset === "1K") return getClosestPresetSize(ar, "normal");
211  if (preset === "2K") return getClosestPresetSize(ar, "2k");
212  if (preset === "4K") return getClosestPresetSize(ar, "4k");
213  return normalizeDimensions(imageSize);
214}
215 
216function getImageSize(ar: string | null, quality: CliArgs["quality"], imageSize?: string | null): string {
217  if (imageSize) {
218    const normalizedSize = normalizeImageSizePreset(imageSize, ar);
219    if (normalizedSize) return normalizedSize;
220  }
221 
222  // Default to 2K quality if not specified
223  const qualityLevel: JimengSizePreset = quality === "normal" ? "normal" : "2k";
224  return getClosestPresetSize(ar, qualityLevel);
225}
226 
227/**
228 * Step 1: Submit async task to Volcengine Jimeng API
229 */
230async function submitTask(
231  prompt: string,
232  model: string,
233  size: string,
234  accessKey: string,
235  secretKey: string,
236  region: string
237): Promise<string> {
238  // Query parameters for submit endpoint
239  const query = {
240    Action: "CVSync2AsyncSubmitTask",
241    Version: "2022-08-31",
242  };
243  const endpoint = resolveEndpoint(query);
244 
245  // Request body - Jimeng API expects width/height as separate integers
246  const [width, height] = size.split("x").map(Number);
247  const bodyObj = {
248    req_key: model,
249    prompt,
250    // Use separate width and height parameters instead of size string
251    width: width,
252    height: height,
253    // Optional: seed for reproducibility
254    // seed: Math.floor(Math.random() * 999999),
255  };
256 
257  const body = JSON.stringify(bodyObj);
258 
259  // Headers
260  const timestampHeader = new Date().toISOString().replace(/[:\-]|\.\d{3}/g, "");
261  const headers = {
262    "Content-Type": "application/json",
263    "X-Date": timestampHeader,
264    "Host": endpoint.host,
265  };
266 
267  // Generate signature
268  const authorization = generateSignature(
269    "POST",
270    query,
271    headers,
272    body,
273    accessKey,
274    secretKey,
275    region,
276    "cv",
277    endpoint.canonicalUri
278  );
279 
280  console.error(`Submitting task to Jimeng (${model})...`, { width, height });
281 
282  const res = await fetch(endpoint.url, {
283    method: "POST",
284    headers: {
285      ...headers,
286      "Authorization": authorization,
287    },
288    body,
289  });
290 
291  if (!res.ok) {
292    const err = await res.text();
293    throw new Error(`Jimeng API submit error (${res.status}): ${err}`);
294  }
295 
296  const result = (await res.json()) as {
297    code?: number;
298    message?: string;
299    data?: {
300      task_id?: string;
301    };
302  };
303 
304  // Volcengine API returns code 10000 for success
305  if (result.code !== 10000 || !result.data?.task_id) {
306    console.error("Submit response:", JSON.stringify(result, null, 2));
307    throw new Error(`Failed to submit task: ${result.message || "Unknown error"}`);
308  }
309 
310  return result.data.task_id;
311}
312 
313/**
314 * Step 2: Poll for task result
315 * Returns image data directly as Uint8Array
316 */
317async function pollForResult(
318  taskId: string,
319  model: string,
320  accessKey: string,
321  secretKey: string,
322  region: string
323): Promise<Uint8Array> {
324  const maxAttempts = 60;
325  const pollIntervalMs = 2000;
326 
327  for (let attempt = 0; attempt < maxAttempts; attempt++) {
328    // Query parameters for result endpoint
329    const query = {
330      Action: "CVSync2AsyncGetResult",
331      Version: "2022-08-31",
332    };
333    const endpoint = resolveEndpoint(query);
334 
335    // Request body - include req_key and task_id
336    const bodyObj = {
337      req_key: model,
338      task_id: taskId,
339    };
340 
341    const body = JSON.stringify(bodyObj);
342 
343    // Headers
344    const timestampHeader = new Date().toISOString().replace(/[:\-]|\.\d{3}/g, "");
345    const headers = {
346      "Content-Type": "application/json",
347      "X-Date": timestampHeader,
348      "Host": endpoint.host,
349    };
350 
351    // Generate signature
352    const authorization = generateSignature(
353      "POST",
354      query,
355      headers,
356      body,
357      accessKey,
358      secretKey,
359      region,
360      "cv",
361      endpoint.canonicalUri
362    );
363 
364    const res = await fetch(endpoint.url, {
365      method: "POST",
366      headers: {
367        ...headers,
368        "Authorization": authorization,
369      },
370      body,
371    });
372 
373    if (!res.ok) {
374      const err = await res.text();
375      throw new Error(`Jimeng API poll error (${res.status}): ${err}`);
376    }
377 
378    const result = (await res.json()) as {
379      code?: number;
380      message?: string;
381      data?: {
382        status?: string;
383        image_urls?: string[];
384        binary_data_base64?: string[];
385      };
386    };
387 
388    // Volcengine API returns code 10000 for success
389    if (result.code === 10000 && result.data) {
390      const { status, image_urls, binary_data_base64 } = result.data;
391 
392      // Check for base64 image data (preferred by Jimeng)
393      if (binary_data_base64 && binary_data_base64.length > 0) {
394        console.error("Image received as base64 data");
395        const base64Data = binary_data_base64[0]!;
396        // Convert base64 to Uint8Array
397        const binaryString = Buffer.from(base64Data, "base64").toString("binary");
398        const bytes = new Uint8Array(binaryString.length);
399        for (let i = 0; i < binaryString.length; i++) {
400          bytes[i] = binaryString.charCodeAt(i);
401        }
402        return bytes;
403      }
404 
405      // Fallback to URL format
406      if (status === "done" && image_urls && image_urls.length > 0) {
407        // Download from URL
408        console.error(`Downloading image from ${image_urls[0]}...`);
409        const imgRes = await fetch(image_urls[0]!);
410        if (!imgRes.ok) {
411          throw new Error(`Failed to download image from ${image_urls[0]}`);
412        }
413        const buffer = await imgRes.arrayBuffer();
414        return new Uint8Array(buffer);
415      }
416 
417      if (status === "in_queue" || status === "generating") {
418        console.error(`Task status: ${status} (${attempt + 1}/${maxAttempts})`);
419        await new Promise(resolve => setTimeout(resolve, pollIntervalMs));
420        continue;
421      }
422 
423      if (status === "fail") {
424        throw new Error(`Jimeng task failed: ${result.message || "Generation failed"}`);
425      }
426    }
427 
428    console.error("Poll response:", JSON.stringify(result, null, 2));
429    throw new Error(`Unexpected response during polling: ${result.message || "Unknown error"}`);
430  }
431 
432  throw new Error("Task timeout: image generation took too long");
433}
434 
435export async function generateImage(
436  prompt: string,
437  model: string,
438  args: CliArgs
439): Promise<Uint8Array> {
440  if (args.referenceImages.length > 0) {
441    throw new Error(
442      "Jimeng does not support reference images. Use --provider google, openai, openrouter, or replicate."
443    );
444  }
445 
446  const accessKey = getAccessKey();
447  const secretKey = getSecretKey();
448  const region = getRegion();
449 
450  if (!accessKey || !secretKey) {
451    throw new Error(
452      "JIMENG_ACCESS_KEY_ID and JIMENG_SECRET_ACCESS_KEY are required. " +
453      "Get your credentials from https://console.volcengine.com/iam/keymanage"
454    );
455  }
456 
457  const size = getImageSize(args.aspectRatio, args.quality, args.imageSize);
458 
459  // Step 1: Submit task
460  const taskId = await submitTask(prompt, model, size, accessKey, secretKey, region);
461 
462  // Step 2: Poll for result (returns image data directly)
463  const imageData = await pollForResult(taskId, model, accessKey, secretKey, region);
464 
465  console.error("Image generation complete!");
466  return imageData;
467}
468