1import path from "node:path";
2import { readFile } from "node:fs/promises";
3import type { CliArgs } from "../types";
4 
5const DEFAULT_MODEL = "google/gemini-3.1-flash-image-preview";
6const COMMON_ASPECT_RATIOS = [
7  "1:1",
8  "2:3",
9  "3:2",
10  "3:4",
11  "4:3",
12  "4:5",
13  "5:4",
14  "9:16",
15  "16:9",
16  "21:9",
17];
18const GEMINI_EXTENDED_ASPECT_RATIOS = ["1:4", "4:1", "1:8", "8:1"];
19 
20type OpenRouterImageEntry = {
21  image_url?: string | { url?: string | null } | null;
22  imageUrl?: string | { url?: string | null } | null;
23};
24 
25type OpenRouterMessagePart = {
26  type?: string;
27  text?: string;
28  image_url?: string | { url?: string | null } | null;
29  imageUrl?: string | { url?: string | null } | null;
30};
31 
32type OpenRouterResponse = {
33  choices?: Array<{
34    finish_reason?: string | null;
35    native_finish_reason?: string | null;
36    message?: {
37      images?: OpenRouterImageEntry[];
38      content?: string | OpenRouterMessagePart[] | null;
39    };
40  }>;
41};
42 
43export function getDefaultModel(): string {
44  return process.env.OPENROUTER_IMAGE_MODEL || DEFAULT_MODEL;
45}
46 
47function normalizeModelId(model: string): string {
48  return model.trim().toLowerCase().split(":")[0]!;
49}
50 
51function isTextAndImageModel(model: string): boolean {
52  const normalized = normalizeModelId(model);
53  if (normalized === "openrouter/auto") {
54    return true;
55  }
56 
57  if (normalized.startsWith("google/gemini-") && normalized.includes("image")) {
58    return true;
59  }
60 
61  if (normalized.startsWith("openai/gpt-") && normalized.includes("image")) {
62    return true;
63  }
64 
65  return false;
66}
67 
68function getSupportedAspectRatios(model: string): Set<string> {
69  const normalized = normalizeModelId(model);
70  if (normalized !== "google/gemini-3.1-flash-image-preview") {
71    return new Set(COMMON_ASPECT_RATIOS);
72  }
73 
74  return new Set([...COMMON_ASPECT_RATIOS, ...GEMINI_EXTENDED_ASPECT_RATIOS]);
75}
76 
77function getApiKey(): string | null {
78  return process.env.OPENROUTER_API_KEY || null;
79}
80 
81function getBaseUrl(): string {
82  const base = process.env.OPENROUTER_BASE_URL || "https://openrouter.ai/api/v1";
83  return base.replace(/\/+$/g, "");
84}
85 
86function getHeaders(apiKey: string): Record<string, string> {
87  const headers: Record<string, string> = {
88    "Content-Type": "application/json",
89    Authorization: `Bearer ${apiKey}`,
90  };
91 
92  const referer = process.env.OPENROUTER_HTTP_REFERER?.trim();
93  if (referer) {
94    headers["HTTP-Referer"] = referer;
95  }
96 
97  const title = process.env.OPENROUTER_TITLE?.trim();
98  if (title) {
99    headers["X-OpenRouter-Title"] = title;
100    headers["X-Title"] = title;
101  }
102 
103  return headers;
104}
105 
106function parsePixelSize(value: string): { width: number; height: number } | null {
107  const match = value.match(/^(\d+)\s*[xX]\s*(\d+)$/);
108  if (!match) return null;
109 
110  const width = parseInt(match[1]!, 10);
111  const height = parseInt(match[2]!, 10);
112 
113  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
114    return null;
115  }
116 
117  return { width, height };
118}
119 
120function gcd(a: number, b: number): number {
121  let x = Math.abs(a);
122  let y = Math.abs(b);
123  while (y !== 0) {
124    const next = x % y;
125    x = y;
126    y = next;
127  }
128  return x || 1;
129}
130 
131function inferAspectRatio(size: string | null): string | null {
132  if (!size) return null;
133  const parsed = parsePixelSize(size);
134  if (!parsed) return null;
135 
136  const divisor = gcd(parsed.width, parsed.height);
137  return `${parsed.width / divisor}:${parsed.height / divisor}`;
138}
139 
140function inferImageSize(size: string | null): "1K" | "2K" | "4K" | null {
141  if (!size) return null;
142  const parsed = parsePixelSize(size);
143  if (!parsed) return null;
144 
145  const longestEdge = Math.max(parsed.width, parsed.height);
146  if (longestEdge <= 1024) return "1K";
147  if (longestEdge <= 2048) return "2K";
148  return "4K";
149}
150 
151export function getImageSize(args: CliArgs): "1K" | "2K" | "4K" | null {
152  if (args.imageSize) return args.imageSize as "1K" | "2K" | "4K";
153 
154  const inferredFromSize = inferImageSize(args.size);
155  if (inferredFromSize) return inferredFromSize;
156 
157  if (args.quality === "normal") return "1K";
158  if (args.quality === "2k") return "2K";
159  return null;
160}
161 
162export function getAspectRatio(model: string, args: CliArgs): string | null {
163  if (args.aspectRatio) return args.aspectRatio;
164 
165  const inferred = inferAspectRatio(args.size);
166  if (!inferred || !getSupportedAspectRatios(model).has(inferred)) {
167    return null;
168  }
169 
170  return inferred;
171}
172 
173function getModalities(model: string): string[] {
174  return isTextAndImageModel(model) ? ["image", "text"] : ["image"];
175}
176 
177export function validateArgs(model: string, args: CliArgs): void {
178  const requestedAspectRatio = args.aspectRatio || inferAspectRatio(args.size);
179  if (!requestedAspectRatio) {
180    return;
181  }
182 
183  const supported = getSupportedAspectRatios(model);
184  if (supported.has(requestedAspectRatio)) {
185    return;
186  }
187 
188  const requestedValue = args.aspectRatio
189    ? `aspect ratio ${requestedAspectRatio}`
190    : `size ${args.size} (aspect ratio ${requestedAspectRatio})`;
191 
192  throw new Error(
193    `OpenRouter model ${model} does not support ${requestedValue}. Supported values: ${Array.from(supported).join(", ")}`
194  );
195}
196 
197function getMimeType(filename: string): string {
198  const ext = path.extname(filename).toLowerCase();
199  if (ext === ".jpg" || ext === ".jpeg") return "image/jpeg";
200  if (ext === ".webp") return "image/webp";
201  if (ext === ".gif") return "image/gif";
202  return "image/png";
203}
204 
205async function readImageAsDataUrl(filePath: string): Promise<string> {
206  const bytes = await readFile(filePath);
207  return `data:${getMimeType(filePath)};base64,${bytes.toString("base64")}`;
208}
209 
210export function buildContent(
211  prompt: string,
212  referenceImages: string[],
213): string | Array<Record<string, unknown>> {
214  if (referenceImages.length === 0) {
215    return prompt;
216  }
217 
218  const content: Array<Record<string, unknown>> = [{ type: "text", text: prompt }];
219 
220  for (const imageUrl of referenceImages) {
221    content.push({
222      type: "image_url",
223      image_url: { url: imageUrl },
224    });
225  }
226 
227  return content;
228}
229 
230function extractImageUrl(entry: OpenRouterImageEntry | OpenRouterMessagePart): string | null {
231  const value = entry.image_url ?? entry.imageUrl;
232  if (!value) return null;
233  if (typeof value === "string") return value;
234  return value.url ?? null;
235}
236 
237function decodeDataUrl(value: string): Uint8Array | null {
238  const match = value.match(/^data:image\/[^;]+;base64,([A-Za-z0-9+/=]+)$/);
239  if (!match) return null;
240  return Uint8Array.from(Buffer.from(match[1]!, "base64"));
241}
242 
243async function downloadImage(value: string): Promise<Uint8Array> {
244  const inline = decodeDataUrl(value);
245  if (inline) return inline;
246 
247  if (value.startsWith("http://") || value.startsWith("https://")) {
248    const response = await fetch(value);
249    if (!response.ok) {
250      throw new Error(`Failed to download OpenRouter image: ${response.status}`);
251    }
252    const buffer = await response.arrayBuffer();
253    return new Uint8Array(buffer);
254  }
255 
256  return Uint8Array.from(Buffer.from(value, "base64"));
257}
258 
259export async function extractImageFromResponse(result: OpenRouterResponse): Promise<Uint8Array> {
260  const choice = result.choices?.[0];
261  const message = choice?.message;
262 
263  for (const image of message?.images ?? []) {
264    const imageUrl = extractImageUrl(image);
265    if (imageUrl) return downloadImage(imageUrl);
266  }
267 
268  if (Array.isArray(message?.content)) {
269    for (const item of message.content) {
270      const imageUrl = extractImageUrl(item);
271      if (imageUrl) return downloadImage(imageUrl);
272 
273      if (item.type === "text" && item.text) {
274        const inline = decodeDataUrl(item.text);
275        if (inline) return inline;
276      }
277    }
278  } else if (typeof message?.content === "string") {
279    const inline = decodeDataUrl(message.content);
280    if (inline) return inline;
281  }
282 
283  const finishReason =
284    choice?.native_finish_reason || choice?.finish_reason || "unknown";
285  throw new Error(
286    `No image in OpenRouter response (finish_reason=${finishReason})`,
287  );
288}
289 
290export function buildRequestBody(
291  prompt: string,
292  model: string,
293  args: CliArgs,
294  referenceImages: string[],
295): Record<string, unknown> {
296  validateArgs(model, args);
297 
298  const imageConfig: Record<string, string> = {};
299 
300  const imageSize = getImageSize(args);
301  if (imageSize) {
302    imageConfig.image_size = imageSize;
303  }
304 
305  const aspectRatio = getAspectRatio(model, args);
306  if (aspectRatio) {
307    imageConfig.aspect_ratio = aspectRatio;
308  }
309 
310  const body: Record<string, unknown> = {
311    messages: [
312      {
313        role: "user",
314        content: buildContent(prompt, referenceImages),
315      },
316    ],
317    modalities: getModalities(model),
318    stream: false,
319  };
320 
321  if (Object.keys(imageConfig).length > 0) {
322    body.image_config = imageConfig;
323    body.provider = {
324      require_parameters: true,
325    };
326  }
327 
328  return body;
329}
330 
331export async function generateImage(
332  prompt: string,
333  model: string,
334  args: CliArgs
335): Promise<Uint8Array> {
336  const apiKey = getApiKey();
337  if (!apiKey) {
338    throw new Error("OPENROUTER_API_KEY is required. Get one at https://openrouter.ai/settings/keys");
339  }
340 
341  const referenceImages: string[] = [];
342  for (const refPath of args.referenceImages) {
343    referenceImages.push(await readImageAsDataUrl(refPath));
344  }
345 
346  const body = {
347    model,
348    ...buildRequestBody(prompt, model, args, referenceImages),
349  };
350 
351  console.log(
352    `Generating image with OpenRouter (${model})...`,
353    (body.image_config as Record<string, string>),
354  );
355 
356  const response = await fetch(`${getBaseUrl()}/chat/completions`, {
357    method: "POST",
358    headers: getHeaders(apiKey),
359    body: JSON.stringify(body),
360  });
361 
362  if (!response.ok) {
363    const errorText = await response.text();
364    throw new Error(`OpenRouter API error (${response.status}): ${errorText}`);
365  }
366 
367  const result = (await response.json()) as OpenRouterResponse;
368  return extractImageFromResponse(result);
369}
370