1import path from "node:path";
2import { readFile } from "node:fs/promises";
3import type { CliArgs } from "../types";
4 
5export function getDefaultModel(): string {
6  return process.env.OPENAI_IMAGE_MODEL || "gpt-image-1.5";
7}
8 
9type OpenAIImageResponse = { data: Array<{ url?: string; b64_json?: string }> };
10 
11export function parseAspectRatio(ar: string): { width: number; height: number } | null {
12  const match = ar.match(/^(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)$/);
13  if (!match) return null;
14  const w = parseFloat(match[1]!);
15  const h = parseFloat(match[2]!);
16  if (w <= 0 || h <= 0) return null;
17  return { width: w, height: h };
18}
19 
20type SizeMapping = {
21  square: string;
22  landscape: string;
23  portrait: string;
24};
25 
26export function getOpenAISize(
27  model: string,
28  ar: string | null,
29  quality: CliArgs["quality"]
30): string {
31  const isDalle3 = model.includes("dall-e-3");
32  const isDalle2 = model.includes("dall-e-2");
33 
34  if (isDalle2) {
35    return "1024x1024";
36  }
37 
38  const sizes: SizeMapping = isDalle3
39    ? {
40        square: "1024x1024",
41        landscape: "1792x1024",
42        portrait: "1024x1792",
43      }
44    : {
45        square: "1024x1024",
46        landscape: "1536x1024",
47        portrait: "1024x1536",
48      };
49 
50  if (!ar) return sizes.square;
51 
52  const parsed = parseAspectRatio(ar);
53  if (!parsed) return sizes.square;
54 
55  const ratio = parsed.width / parsed.height;
56 
57  if (Math.abs(ratio - 1) < 0.1) return sizes.square;
58  if (ratio > 1.5) return sizes.landscape;
59  if (ratio < 0.67) return sizes.portrait;
60  return sizes.square;
61}
62 
63export async function generateImage(
64  prompt: string,
65  model: string,
66  args: CliArgs
67): Promise<Uint8Array> {
68  const baseURL = process.env.OPENAI_BASE_URL || "https://api.openai.com/v1";
69  const apiKey = process.env.OPENAI_API_KEY;
70 
71  if (!apiKey) {
72    throw new Error(
73      "OPENAI_API_KEY is required. Codex/ChatGPT desktop login does not automatically grant OpenAI Images API access to this script."
74    );
75  }
76 
77  if (process.env.OPENAI_IMAGE_USE_CHAT === "true") {
78    return generateWithChatCompletions(baseURL, apiKey, prompt, model);
79  }
80 
81  const size = args.size || getOpenAISize(model, args.aspectRatio, args.quality);
82 
83  if (args.referenceImages.length > 0) {
84    if (model.includes("dall-e-2") || model.includes("dall-e-3")) {
85      throw new Error(
86        "Reference images with OpenAI in this skill require GPT Image models. Use --model gpt-image-1.5 (or another gpt-image model)."
87      );
88    }
89    return generateWithOpenAIEdits(baseURL, apiKey, prompt, model, size, args.referenceImages, args.quality);
90  }
91 
92  return generateWithOpenAIGenerations(baseURL, apiKey, prompt, model, size, args.quality);
93}
94 
95async function generateWithChatCompletions(
96  baseURL: string,
97  apiKey: string,
98  prompt: string,
99  model: string
100): Promise<Uint8Array> {
101  const res = await fetch(`${baseURL}/chat/completions`, {
102    method: "POST",
103    headers: {
104      "Content-Type": "application/json",
105      Authorization: `Bearer ${apiKey}`,
106    },
107    body: JSON.stringify({
108      model,
109      messages: [{ role: "user", content: prompt }],
110    }),
111  });
112 
113  if (!res.ok) {
114    const err = await res.text();
115    throw new Error(`OpenAI API error: ${err}`);
116  }
117 
118  const result = (await res.json()) as { choices: Array<{ message: { content: string } }> };
119  const content = result.choices[0]?.message?.content ?? "";
120 
121  const match = content.match(/data:image\/[^;]+;base64,([A-Za-z0-9+/=]+)/);
122  if (match) {
123    return Uint8Array.from(Buffer.from(match[1]!, "base64"));
124  }
125 
126  throw new Error("No image found in chat completions response");
127}
128 
129async function generateWithOpenAIGenerations(
130  baseURL: string,
131  apiKey: string,
132  prompt: string,
133  model: string,
134  size: string,
135  quality: CliArgs["quality"]
136): Promise<Uint8Array> {
137  const body: Record<string, any> = { model, prompt, size };
138 
139  if (model.includes("dall-e-3")) {
140    body.quality = quality === "2k" ? "hd" : "standard";
141  }
142 
143  const res = await fetch(`${baseURL}/images/generations`, {
144    method: "POST",
145    headers: {
146      "Content-Type": "application/json",
147      Authorization: `Bearer ${apiKey}`,
148    },
149    body: JSON.stringify(body),
150  });
151 
152  if (!res.ok) {
153    const err = await res.text();
154    throw new Error(`OpenAI API error: ${err}`);
155  }
156 
157  const result = (await res.json()) as OpenAIImageResponse;
158  return extractImageFromResponse(result);
159}
160 
161async function generateWithOpenAIEdits(
162  baseURL: string,
163  apiKey: string,
164  prompt: string,
165  model: string,
166  size: string,
167  referenceImages: string[],
168  quality: CliArgs["quality"]
169): Promise<Uint8Array> {
170  const form = new FormData();
171  form.append("model", model);
172  form.append("prompt", prompt);
173  form.append("size", size);
174 
175  if (model.includes("gpt-image")) {
176    form.append("quality", quality === "2k" ? "high" : "medium");
177  }
178 
179  for (const refPath of referenceImages) {
180    const bytes = await readFile(refPath);
181    const filename = path.basename(refPath);
182    const mimeType = getMimeType(filename);
183    const blob = new Blob([bytes], { type: mimeType });
184    form.append("image[]", blob, filename);
185  }
186 
187  const res = await fetch(`${baseURL}/images/edits`, {
188    method: "POST",
189    headers: {
190      Authorization: `Bearer ${apiKey}`,
191    },
192    body: form,
193  });
194 
195  if (!res.ok) {
196    const err = await res.text();
197    throw new Error(`OpenAI edits API error: ${err}`);
198  }
199 
200  const result = (await res.json()) as OpenAIImageResponse;
201  return extractImageFromResponse(result);
202}
203 
204export function getMimeType(filename: string): string {
205  const ext = path.extname(filename).toLowerCase();
206  if (ext === ".jpg" || ext === ".jpeg") return "image/jpeg";
207  if (ext === ".webp") return "image/webp";
208  if (ext === ".gif") return "image/gif";
209  return "image/png";
210}
211 
212export async function extractImageFromResponse(result: OpenAIImageResponse): Promise<Uint8Array> {
213  const img = result.data[0];
214 
215  if (img?.b64_json) {
216    return Uint8Array.from(Buffer.from(img.b64_json, "base64"));
217  }
218 
219  if (img?.url) {
220    const imgRes = await fetch(img.url);
221    if (!imgRes.ok) throw new Error("Failed to download image");
222    const buf = await imgRes.arrayBuffer();
223    return new Uint8Array(buf);
224  }
225 
226  throw new Error("No image in response");
227}
228