1import path from "node:path";
2import { readFile } from "node:fs/promises";
3import type { CliArgs } from "../types";
4import { getOpenAISize, extractImageFromResponse } from "./openai.ts";
5 
6type OpenAIImageResponse = { data: Array<{ url?: string; b64_json?: string }> };
7type AzureEndpoint = {
8  resourceBaseURL: string;
9  deployment: string | null;
10};
11 
12const DEFAULT_AZURE_API_VERSION = "2025-04-01-preview";
13const AZURE_EDIT_IMAGE_EXTENSIONS = new Set([".png", ".jpg", ".jpeg"]);
14 
15export function parseAzureBaseURL(url: string): AzureEndpoint {
16  const parsed = new URL(url);
17  const trimmedPath = parsed.pathname.replace(/\/+$/, "");
18  const deploymentMatch = trimmedPath.match(/^(.*?)(?:\/openai)?\/deployments\/([^/]+)$/);
19 
20  if (deploymentMatch) {
21    parsed.pathname = `${deploymentMatch[1] || ""}/openai`;
22    return {
23      resourceBaseURL: parsed.toString().replace(/\/+$/, ""),
24      deployment: decodeURIComponent(deploymentMatch[2]!),
25    };
26  }
27 
28  parsed.pathname = trimmedPath.endsWith("/openai") ? trimmedPath : `${trimmedPath}/openai`;
29  return {
30    resourceBaseURL: parsed.toString().replace(/\/+$/, ""),
31    deployment: null,
32  };
33}
34 
35export function getDefaultModel(): string {
36  const explicitDeployment = process.env.AZURE_OPENAI_DEPLOYMENT?.trim();
37  if (explicitDeployment) return explicitDeployment;
38 
39  const baseURL = process.env.AZURE_OPENAI_BASE_URL;
40  if (baseURL) {
41    try {
42      const { deployment } = parseAzureBaseURL(baseURL);
43      if (deployment) return deployment;
44    } catch {
45      // Ignore invalid URLs here so the required-env check can raise the user-facing error later.
46    }
47  }
48 
49  return process.env.AZURE_OPENAI_IMAGE_MODEL || "gpt-image-1.5";
50}
51 
52function getEndpoint(): AzureEndpoint {
53  const url = process.env.AZURE_OPENAI_BASE_URL;
54  if (!url) {
55    throw new Error(
56      "AZURE_OPENAI_BASE_URL is required. Set it to your Azure resource or deployment endpoint, e.g.: https://your-resource.openai.azure.com or https://your-resource.openai.azure.com/openai/deployments/your-deployment"
57    );
58  }
59  return parseAzureBaseURL(url);
60}
61 
62function getApiKey(): string {
63  const key = process.env.AZURE_OPENAI_API_KEY;
64  if (!key) {
65    throw new Error(
66      "AZURE_OPENAI_API_KEY is required. Get it from Azure Portal → your OpenAI resource → Keys and Endpoint."
67    );
68  }
69  return key;
70}
71 
72function getApiVersion(): string {
73  return process.env.AZURE_API_VERSION || DEFAULT_AZURE_API_VERSION;
74}
75 
76function getDeployment(model: string): string {
77  const deployment = model.trim();
78  if (!deployment) {
79    throw new Error(
80      "Azure deployment name is required. Use --model <deployment>, AZURE_OPENAI_DEPLOYMENT, AZURE_OPENAI_IMAGE_MODEL, or embed the deployment in AZURE_OPENAI_BASE_URL."
81    );
82  }
83  return deployment;
84}
85 
86function buildURL(deployment: string, pathSuffix: string): string {
87  const { resourceBaseURL } = getEndpoint();
88  return `${resourceBaseURL}/deployments/${encodeURIComponent(deployment)}${pathSuffix}?api-version=${getApiVersion()}`;
89}
90 
91function authHeaders(): Record<string, string> {
92  return { "api-key": getApiKey() };
93}
94 
95function getAzureQuality(quality: CliArgs["quality"]): "medium" | "high" {
96  return quality === "2k" ? "high" : "medium";
97}
98 
99export function validateArgs(_model: string, args: CliArgs): void {
100  for (const refPath of args.referenceImages) {
101    const ext = path.extname(refPath).toLowerCase();
102    if (!AZURE_EDIT_IMAGE_EXTENSIONS.has(ext)) {
103      throw new Error(
104        `Azure OpenAI reference images must be PNG or JPG/JPEG. Unsupported file: ${refPath}`
105      );
106    }
107  }
108}
109 
110export async function generateImage(
111  prompt: string,
112  model: string,
113  args: CliArgs
114): Promise<Uint8Array> {
115  const deployment = getDeployment(model);
116  const size = args.size || getOpenAISize(model, args.aspectRatio, args.quality);
117 
118  if (args.referenceImages.length > 0) {
119    return generateWithAzureEdits(prompt, deployment, size, args.referenceImages, args.quality);
120  }
121 
122  return generateWithAzureGenerations(prompt, deployment, size, args.quality);
123}
124 
125async function generateWithAzureGenerations(
126  prompt: string,
127  deployment: string,
128  size: string,
129  quality: CliArgs["quality"]
130): Promise<Uint8Array> {
131  const body: Record<string, any> = {
132    prompt,
133    size,
134    n: 1,
135    quality: getAzureQuality(quality),
136  };
137 
138  const res = await fetch(buildURL(deployment, "/images/generations"), {
139    method: "POST",
140    headers: {
141      "Content-Type": "application/json",
142      ...authHeaders(),
143    },
144    body: JSON.stringify(body),
145  });
146 
147  if (!res.ok) {
148    const err = await res.text();
149    throw new Error(`Azure OpenAI API error: ${err}`);
150  }
151 
152  const result = (await res.json()) as OpenAIImageResponse;
153  return extractImageFromResponse(result);
154}
155 
156async function generateWithAzureEdits(
157  prompt: string,
158  deployment: string,
159  size: string,
160  referenceImages: string[],
161  quality: CliArgs["quality"]
162): Promise<Uint8Array> {
163  const form = new FormData();
164  form.append("prompt", prompt);
165  form.append("size", size);
166  form.append("n", "1");
167  form.append("quality", getAzureQuality(quality));
168 
169  for (const refPath of referenceImages) {
170    const bytes = await readFile(refPath);
171    const filename = path.basename(refPath);
172    const mimeType = path.extname(filename).toLowerCase() === ".png" ? "image/png" : "image/jpeg";
173    const blob = new Blob([bytes], { type: mimeType });
174    form.append("image[]", blob, filename);
175  }
176 
177  const res = await fetch(buildURL(deployment, "/images/edits"), {
178    method: "POST",
179    headers: {
180      ...authHeaders(),
181    },
182    body: form,
183  });
184 
185  if (!res.ok) {
186    const err = await res.text();
187    throw new Error(`Azure OpenAI edits API error: ${err}`);
188  }
189 
190  const result = (await res.json()) as OpenAIImageResponse;
191  return extractImageFromResponse(result);
192}
193