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