1import path from "node:path";
2import { readFile } from "node:fs/promises";
3import { execFileSync } from "node:child_process";
4import type { CliArgs } from "../types";
5 
6const GOOGLE_MULTIMODAL_MODELS = [
7  "gemini-3-pro-image",
8  "gemini-3.1-flash-image",
9  "gemini-3-pro-image-preview",
10  "gemini-3-flash-preview",
11  "gemini-3.1-flash-image-preview",
12];
13const GOOGLE_IMAGEN_MODELS = [
14  "imagen-3.0-generate-002",
15  "imagen-3.0-generate-001",
16];
17 
18export function getDefaultModel(): string {
19  return process.env.GOOGLE_IMAGE_MODEL || "gemini-3-pro-image";
20}
21 
22export function normalizeGoogleModelId(model: string): string {
23  return model.startsWith("models/") ? model.slice("models/".length) : model;
24}
25 
26export function isGoogleMultimodal(model: string): boolean {
27  const normalized = normalizeGoogleModelId(model);
28  return GOOGLE_MULTIMODAL_MODELS.some((m) => normalized.includes(m));
29}
30 
31export function isGoogleImagen(model: string): boolean {
32  const normalized = normalizeGoogleModelId(model);
33  return GOOGLE_IMAGEN_MODELS.some((m) => normalized.includes(m));
34}
35 
36function getGoogleApiKey(): string | null {
37  return process.env.GOOGLE_API_KEY || process.env.GEMINI_API_KEY || null;
38}
39 
40export function getGoogleImageSize(args: CliArgs): "1K" | "2K" | "4K" {
41  if (args.imageSize) return args.imageSize as "1K" | "2K" | "4K";
42  return args.quality === "2k" ? "2K" : "1K";
43}
44 
45function getGoogleBaseUrl(): string {
46  const base =
47    process.env.GOOGLE_BASE_URL || "https://generativelanguage.googleapis.com";
48  return base.replace(/\/+$/g, "");
49}
50 
51export function buildGoogleUrl(pathname: string): string {
52  const base = getGoogleBaseUrl();
53  const cleanedPath = pathname.replace(/^\/+/g, "");
54  if (base.endsWith("/v1beta")) return `${base}/${cleanedPath}`;
55  return `${base}/v1beta/${cleanedPath}`;
56}
57 
58function toModelPath(model: string): string {
59  const modelId = normalizeGoogleModelId(model);
60  return `models/${modelId}`;
61}
62 
63function getHttpProxy(): string | null {
64  return (
65    process.env.https_proxy ||
66    process.env.HTTPS_PROXY ||
67    process.env.http_proxy ||
68    process.env.HTTP_PROXY ||
69    process.env.ALL_PROXY ||
70    null
71  );
72}
73 
74async function postGoogleJsonViaCurl<T>(
75  url: string,
76  apiKey: string,
77  body: unknown,
78): Promise<T> {
79  const proxy = getHttpProxy();
80  const bodyStr = JSON.stringify(body);
81  const args = [
82    "-s",
83    "--connect-timeout",
84    "30",
85    "--max-time",
86    "300",
87    ...(proxy ? ["-x", proxy] : []),
88    url,
89    "-H",
90    "Content-Type: application/json",
91    "-H",
92    `x-goog-api-key: ${apiKey}`,
93    "-d",
94    "@-",
95  ];
96 
97  let result = "";
98  try {
99    result = execFileSync("curl", args, {
100      input: bodyStr,
101      encoding: "utf8",
102      maxBuffer: 100 * 1024 * 1024,
103      timeout: 310000,
104    });
105  } catch (error) {
106    const e = error as { message?: string; stderr?: string | Buffer };
107    const stderrText =
108      typeof e.stderr === "string"
109        ? e.stderr
110        : e.stderr
111          ? e.stderr.toString("utf8")
112          : "";
113    const details = stderrText.trim() || e.message || "curl request failed";
114    throw new Error(`Google API request failed via curl: ${details}`);
115  }
116 
117  const parsed = JSON.parse(result) as any;
118  if (parsed.error) {
119    throw new Error(
120      `Google API error (${parsed.error.code}): ${parsed.error.message}`,
121    );
122  }
123  return parsed as T;
124}
125 
126async function postGoogleJsonViaFetch<T>(
127  url: string,
128  apiKey: string,
129  body: unknown,
130): Promise<T> {
131  const res = await fetch(url, {
132    method: "POST",
133    headers: {
134      "Content-Type": "application/json",
135      "x-goog-api-key": apiKey,
136    },
137    body: JSON.stringify(body),
138  });
139 
140  if (!res.ok) {
141    const err = await res.text();
142    throw new Error(`Google API error (${res.status}): ${err}`);
143  }
144 
145  return (await res.json()) as T;
146}
147 
148async function postGoogleJson<T>(pathname: string, body: unknown): Promise<T> {
149  const apiKey = getGoogleApiKey();
150  if (!apiKey) throw new Error("GOOGLE_API_KEY or GEMINI_API_KEY is required");
151 
152  const url = buildGoogleUrl(pathname);
153  const proxy = getHttpProxy();
154 
155  // When an HTTP proxy is detected, use curl instead of fetch.
156  // Bun's fetch has a known issue where long-lived connections through
157  // HTTP proxies get their sockets closed unexpectedly, causing image
158  // generation requests to fail with "socket connection was closed
159  // unexpectedly". Using curl as the HTTP client works around this.
160  if (proxy) {
161    return postGoogleJsonViaCurl<T>(url, apiKey, body);
162  }
163 
164  return postGoogleJsonViaFetch<T>(url, apiKey, body);
165}
166 
167export function buildPromptWithAspect(
168  prompt: string,
169  ar: string | null,
170  quality: CliArgs["quality"],
171): string {
172  let result = prompt;
173  if (ar) {
174    result += ` Aspect ratio: ${ar}.`;
175  }
176  if (quality === "2k") {
177    result += " High resolution 2048px.";
178  }
179  return result;
180}
181 
182export function addAspectRatioToPrompt(prompt: string, ar: string | null): string {
183  if (!ar) return prompt;
184  return `${prompt} Aspect ratio: ${ar}.`;
185}
186 
187async function readImageAsBase64(
188  p: string,
189): Promise<{ data: string; mimeType: string }> {
190  const buf = await readFile(p);
191  const ext = path.extname(p).toLowerCase();
192  let mimeType = "image/png";
193  if (ext === ".jpg" || ext === ".jpeg") mimeType = "image/jpeg";
194  else if (ext === ".gif") mimeType = "image/gif";
195  else if (ext === ".webp") mimeType = "image/webp";
196  return { data: buf.toString("base64"), mimeType };
197}
198 
199export function extractInlineImageData(response: {
200  candidates?: Array<{
201    content?: { parts?: Array<{ inlineData?: { data?: string } }> };
202  }>;
203}): string | null {
204  for (const candidate of response.candidates || []) {
205    for (const part of candidate.content?.parts || []) {
206      const data = part.inlineData?.data;
207      if (typeof data === "string" && data.length > 0) return data;
208    }
209  }
210  return null;
211}
212 
213export function extractPredictedImageData(response: {
214  predictions?: Array<any>;
215  generatedImages?: Array<any>;
216}): string | null {
217  const candidates = [
218    ...(response.predictions || []),
219    ...(response.generatedImages || []),
220  ];
221  for (const candidate of candidates) {
222    if (!candidate || typeof candidate !== "object") continue;
223    if (typeof candidate.imageBytes === "string") return candidate.imageBytes;
224    if (typeof candidate.bytesBase64Encoded === "string")
225      return candidate.bytesBase64Encoded;
226    if (typeof candidate.data === "string") return candidate.data;
227    const image = candidate.image;
228    if (image && typeof image === "object") {
229      if (typeof image.imageBytes === "string") return image.imageBytes;
230      if (typeof image.bytesBase64Encoded === "string")
231        return image.bytesBase64Encoded;
232      if (typeof image.data === "string") return image.data;
233    }
234  }
235  return null;
236}
237 
238async function generateWithGemini(
239  prompt: string,
240  model: string,
241  args: CliArgs,
242): Promise<Uint8Array> {
243  const promptWithAspect = addAspectRatioToPrompt(prompt, args.aspectRatio);
244  const parts: Array<{
245    text?: string;
246    inlineData?: { data: string; mimeType: string };
247  }> = [];
248  for (const refPath of args.referenceImages) {
249    const { data, mimeType } = await readImageAsBase64(refPath);
250    parts.push({ inlineData: { data, mimeType } });
251  }
252  parts.push({ text: promptWithAspect });
253 
254  const imageConfig: { imageSize: "1K" | "2K" | "4K" } = {
255    imageSize: getGoogleImageSize(args),
256  };
257 
258  console.log("Generating image with Gemini...", imageConfig);
259  const response = await postGoogleJson<{
260    candidates?: Array<{
261      content?: { parts?: Array<{ inlineData?: { data?: string } }> };
262    }>;
263  }>(`${toModelPath(model)}:generateContent`, {
264    contents: [
265      {
266        role: "user",
267        parts,
268      },
269    ],
270    generationConfig: {
271      responseModalities: ["IMAGE"],
272      imageConfig,
273    },
274  });
275  console.log("Generation completed.");
276 
277  const imageData = extractInlineImageData(response);
278  if (imageData) return Uint8Array.from(Buffer.from(imageData, "base64"));
279 
280  throw new Error("No image in response");
281}
282 
283async function generateWithImagen(
284  prompt: string,
285  model: string,
286  args: CliArgs,
287): Promise<Uint8Array> {
288  const fullPrompt = buildPromptWithAspect(
289    prompt,
290    args.aspectRatio,
291    args.quality,
292  );
293  const imageSize = getGoogleImageSize(args);
294  if (imageSize === "4K") {
295    console.error(
296      "Warning: Imagen models do not support 4K imageSize, using 2K instead.",
297    );
298  }
299 
300  const parameters: Record<string, unknown> = {
301    sampleCount: args.n,
302  };
303  if (args.aspectRatio) {
304    parameters.aspectRatio = args.aspectRatio;
305  }
306  if (imageSize === "1K" || imageSize === "2K") {
307    parameters.imageSize = imageSize;
308  } else {
309    parameters.imageSize = "2K";
310  }
311 
312  const response = await postGoogleJson<{
313    predictions?: Array<any>;
314    generatedImages?: Array<any>;
315  }>(`${toModelPath(model)}:predict`, {
316    instances: [
317      {
318        prompt: fullPrompt,
319      },
320    ],
321    parameters,
322  });
323 
324  const imageData = extractPredictedImageData(response);
325  if (imageData) return Uint8Array.from(Buffer.from(imageData, "base64"));
326 
327  throw new Error("No image in response");
328}
329 
330export async function generateImage(
331  prompt: string,
332  model: string,
333  args: CliArgs,
334): Promise<Uint8Array> {
335  if (isGoogleImagen(model)) {
336    if (args.referenceImages.length > 0) {
337      throw new Error(
338        "Reference images are not supported with Imagen models. Use a Gemini multimodal model such as gemini-3-pro-image, gemini-3.1-flash-image, gemini-3-pro-image-preview, gemini-3-flash-preview, or gemini-3.1-flash-image-preview.",
339      );
340    }
341    return generateWithImagen(prompt, model, args);
342  }
343 
344  if (!isGoogleMultimodal(model) && args.referenceImages.length > 0) {
345    throw new Error(
346      "Reference images are only supported with Gemini multimodal models such as gemini-3-pro-image, gemini-3.1-flash-image, gemini-3-pro-image-preview, gemini-3-flash-preview, or gemini-3.1-flash-image-preview.",
347    );
348  }
349 
350  return generateWithGemini(prompt, model, args);
351}
352