1import path from "node:path";
2import { readFile } from "node:fs/promises";
3import type { CliArgs } from "../types";
4 
5const DEFAULT_MODEL = "google/nano-banana-pro";
6const SYNC_WAIT_SECONDS = 60;
7const POLL_INTERVAL_MS = 2000;
8const MAX_POLL_MS = 300_000;
9 
10export function getDefaultModel(): string {
11  return process.env.REPLICATE_IMAGE_MODEL || DEFAULT_MODEL;
12}
13 
14function getApiToken(): string | null {
15  return process.env.REPLICATE_API_TOKEN || null;
16}
17 
18function getBaseUrl(): string {
19  const base = process.env.REPLICATE_BASE_URL || "https://api.replicate.com";
20  return base.replace(/\/+$/g, "");
21}
22 
23export function parseModelId(model: string): { owner: string; name: string; version: string | null } {
24  const [ownerName, version] = model.split(":");
25  const parts = ownerName!.split("/");
26  if (parts.length !== 2 || !parts[0] || !parts[1]) {
27    throw new Error(
28      `Invalid Replicate model format: "${model}". Expected "owner/name" or "owner/name:version".`
29    );
30  }
31  return { owner: parts[0], name: parts[1], version: version || null };
32}
33 
34export function buildInput(prompt: string, args: CliArgs, referenceImages: string[]): Record<string, unknown> {
35  const input: Record<string, unknown> = { prompt };
36 
37  if (args.aspectRatio) {
38    input.aspect_ratio = args.aspectRatio;
39  } else if (referenceImages.length > 0) {
40    input.aspect_ratio = "match_input_image";
41  }
42 
43  if (args.n > 1) {
44    input.number_of_images = args.n;
45  }
46 
47  if (args.quality === "normal") {
48    input.resolution = "1K";
49  } else if (args.quality === "2k") {
50    input.resolution = "2K";
51  }
52 
53  input.output_format = "png";
54 
55  if (referenceImages.length > 0) {
56    input.image_input = referenceImages;
57  }
58 
59  return input;
60}
61 
62async function readImageAsDataUrl(p: string): Promise<string> {
63  const buf = await readFile(p);
64  const ext = path.extname(p).toLowerCase();
65  let mimeType = "image/png";
66  if (ext === ".jpg" || ext === ".jpeg") mimeType = "image/jpeg";
67  else if (ext === ".gif") mimeType = "image/gif";
68  else if (ext === ".webp") mimeType = "image/webp";
69  return `data:${mimeType};base64,${buf.toString("base64")}`;
70}
71 
72type PredictionResponse = {
73  id: string;
74  status: string;
75  output: unknown;
76  error: string | null;
77  urls?: { get?: string };
78};
79 
80async function createPrediction(
81  apiToken: string,
82  model: { owner: string; name: string; version: string | null },
83  input: Record<string, unknown>,
84  sync: boolean
85): Promise<PredictionResponse> {
86  const baseUrl = getBaseUrl();
87 
88  let url: string;
89  const body: Record<string, unknown> = { input };
90 
91  if (model.version) {
92    url = `${baseUrl}/v1/predictions`;
93    body.version = model.version;
94  } else {
95    url = `${baseUrl}/v1/models/${model.owner}/${model.name}/predictions`;
96  }
97 
98  const headers: Record<string, string> = {
99    Authorization: `Bearer ${apiToken}`,
100    "Content-Type": "application/json",
101  };
102 
103  if (sync) {
104    headers["Prefer"] = `wait=${SYNC_WAIT_SECONDS}`;
105  }
106 
107  const res = await fetch(url, {
108    method: "POST",
109    headers,
110    body: JSON.stringify(body),
111  });
112 
113  if (!res.ok) {
114    const err = await res.text();
115    throw new Error(`Replicate API error (${res.status}): ${err}`);
116  }
117 
118  return (await res.json()) as PredictionResponse;
119}
120 
121async function pollPrediction(apiToken: string, getUrl: string): Promise<PredictionResponse> {
122  const start = Date.now();
123 
124  while (Date.now() - start < MAX_POLL_MS) {
125    const res = await fetch(getUrl, {
126      headers: { Authorization: `Bearer ${apiToken}` },
127    });
128 
129    if (!res.ok) {
130      const err = await res.text();
131      throw new Error(`Replicate poll error (${res.status}): ${err}`);
132    }
133 
134    const prediction = (await res.json()) as PredictionResponse;
135 
136    if (prediction.status === "succeeded") return prediction;
137    if (prediction.status === "failed" || prediction.status === "canceled") {
138      throw new Error(`Replicate prediction ${prediction.status}: ${prediction.error || "unknown error"}`);
139    }
140 
141    await new Promise((r) => setTimeout(r, POLL_INTERVAL_MS));
142  }
143 
144  throw new Error(`Replicate prediction timed out after ${MAX_POLL_MS / 1000}s`);
145}
146 
147export function extractOutputUrl(prediction: PredictionResponse): string {
148  const output = prediction.output;
149 
150  if (typeof output === "string") return output;
151 
152  if (Array.isArray(output)) {
153    const first = output[0];
154    if (typeof first === "string") return first;
155  }
156 
157  if (output && typeof output === "object" && "url" in output) {
158    const url = (output as Record<string, unknown>).url;
159    if (typeof url === "string") return url;
160  }
161 
162  throw new Error(`Unexpected Replicate output format: ${JSON.stringify(output)}`);
163}
164 
165async function downloadImage(url: string): Promise<Uint8Array> {
166  const res = await fetch(url);
167  if (!res.ok) throw new Error(`Failed to download image from Replicate: ${res.status}`);
168  const buf = await res.arrayBuffer();
169  return new Uint8Array(buf);
170}
171 
172export async function generateImage(
173  prompt: string,
174  model: string,
175  args: CliArgs
176): Promise<Uint8Array> {
177  const apiToken = getApiToken();
178  if (!apiToken) throw new Error("REPLICATE_API_TOKEN is required. Get one at https://replicate.com/account/api-tokens");
179 
180  const parsedModel = parseModelId(model);
181 
182  const refDataUrls: string[] = [];
183  for (const refPath of args.referenceImages) {
184    refDataUrls.push(await readImageAsDataUrl(refPath));
185  }
186 
187  const input = buildInput(prompt, args, refDataUrls);
188 
189  console.log(`Generating image with Replicate (${model})...`);
190 
191  let prediction = await createPrediction(apiToken, parsedModel, input, true);
192 
193  if (prediction.status !== "succeeded") {
194    if (!prediction.urls?.get) {
195      throw new Error("Replicate prediction did not return a poll URL");
196    }
197    console.log("Waiting for prediction to complete...");
198    prediction = await pollPrediction(apiToken, prediction.urls.get);
199  }
200 
201  console.log("Generation completed.");
202 
203  const outputUrl = extractOutputUrl(prediction);
204  return downloadImage(outputUrl);
205}
206