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