1import { readFile } from "node:fs/promises";
2import path from "node:path";
3import type { CliArgs } from "../types";
4 
5const DEFAULT_MODEL = "agnes-image-2.1-flash";
6const DEFAULT_BASE_URL = "https://apihub.agnes-ai.com/v1";
7const DEFAULT_SIZE = "1024x1024";
8 
9type AgnesResponse = {
10  created?: number;
11  data: Array<{ url?: string; b64_json?: string }>;
12};
13 
14export function getDefaultModel(): string {
15  return process.env.AGNES_IMAGE_MODEL || DEFAULT_MODEL;
16}
17 
18function getApiKey(): string {
19  const key = process.env.AGNES_API_KEY;
20  if (!key) {
21    throw new Error("AGNES_API_KEY is required. Get one from https://apihub.agnes-ai.com.");
22  }
23  return key;
24}
25 
26function getBaseUrl(): string {
27  return (process.env.AGNES_BASE_URL || DEFAULT_BASE_URL).replace(/\/+$/, "");
28}
29 
30export function parseAspectRatio(ar: string): { width: number; height: number } | null {
31  const match = ar.match(/^(\d+(?:\.\d+)?):(\d+(?:\.\d+)?)$/);
32  if (!match) return null;
33  const w = parseFloat(match[1]!);
34  const h = parseFloat(match[2]!);
35  if (w <= 0 || h <= 0) return null;
36  return { width: w, height: h };
37}
38 
39export function snapDim(n: number): number {
40  return Math.max(32, Math.round(n / 32) * 32);
41}
42 
43export function resolveSize(args: Pick<CliArgs, "size" | "aspectRatio">): string {
44  if (args.size) return args.size;
45 
46  if (args.aspectRatio) {
47    const parsed = parseAspectRatio(args.aspectRatio);
48    if (parsed) {
49      if (parsed.width === 1 && parsed.height === 1) return "1024x1024";
50      const maxEdge = 2048;
51      const scale = Math.max(1, Math.floor(maxEdge / Math.max(parsed.width, parsed.height)));
52      const width = parsed.width * scale;
53      const height = parsed.height * scale;
54      return `${snapDim(width)}x${snapDim(height)}`;
55    }
56  }
57 
58  return DEFAULT_SIZE;
59}
60 
61function isRemoteUrl(refPath: string): boolean {
62  return /^https?:\/\//i.test(refPath);
63}
64 
65export async function resolveReferenceImages(
66  referenceImages: string[]
67): Promise<string[]> {
68  const result: string[] = [];
69  for (const refPath of referenceImages) {
70    if (isRemoteUrl(refPath)) {
71      result.push(refPath);
72      continue;
73    }
74    const bytes = await readFile(refPath);
75    const ext = path.extname(refPath).toLowerCase();
76    let mime = "image/png";
77    if (ext === ".jpg" || ext === ".jpeg") mime = "image/jpeg";
78    else if (ext === ".webp") mime = "image/webp";
79    else if (ext === ".gif") mime = "image/gif";
80    const b64 = Buffer.from(bytes).toString("base64");
81    result.push(`data:${mime};base64,${b64}`);
82  }
83  return result;
84}
85 
86export function validateArgs(_model: string, args: CliArgs): void {
87  if (args.n > 1) {
88    throw new Error("Agnes image generation currently returns a single image per request. Set --n 1 or omit --n.");
89  }
90}
91 
92export function getDefaultOutputExtension(_model: string, args: CliArgs): string {
93  return args.responseFormat === "url" ? ".txt" : ".png";
94}
95 
96export function buildRequestBody(
97  prompt: string,
98  model: string,
99  args: Pick<CliArgs, "size" | "aspectRatio" | "referenceImages">
100): Record<string, unknown> {
101  const body: Record<string, unknown> = {
102    model,
103    prompt,
104    size: resolveSize(args),
105  };
106 
107  if (args.referenceImages.length > 0) {
108    body.image = args.referenceImages;
109  }
110 
111  body.extra_body = { response_format: "url" };
112 
113  return body;
114}
115 
116export async function extractImageFromResponse(result: AgnesResponse): Promise<Uint8Array> {
117  const img = result.data[0];
118 
119  if (img?.b64_json) {
120    return Uint8Array.from(Buffer.from(img.b64_json, "base64"));
121  }
122 
123  if (img?.url) {
124    const imgRes = await fetch(img.url);
125    if (!imgRes.ok) throw new Error(`Failed to download image from Agnes: ${imgRes.status}`);
126    return new Uint8Array(await imgRes.arrayBuffer());
127  }
128 
129  throw new Error("No image in Agnes response");
130}
131 
132export async function generateImage(
133  prompt: string,
134  model: string,
135  args: CliArgs
136): Promise<Uint8Array> {
137  const apiKey = getApiKey();
138  const baseUrl = getBaseUrl();
139 
140  const referenceImages = await resolveReferenceImages(args.referenceImages);
141 
142  const body = buildRequestBody(prompt, model, { ...args, referenceImages });
143 
144  const controller = new AbortController();
145  const timeout = setTimeout(() => controller.abort(), 120_000);
146 
147  try {
148    const res = await fetch(`${baseUrl}/images/generations`, {
149      method: "POST",
150      headers: {
151        "Content-Type": "application/json",
152        Authorization: `Bearer ${apiKey}`,
153      },
154      body: JSON.stringify(body),
155      signal: controller.signal,
156    });
157 
158    if (!res.ok) {
159      const err = await res.text();
160      throw new Error(`Agnes API error (${res.status}): ${err}`);
161    }
162 
163    const result = (await res.json()) as AgnesResponse;
164 
165    if (args.responseFormat === "url") {
166      const url = result.data[0]?.url;
167      if (!url) throw new Error("No URL in Agnes response");
168      return new Uint8Array(Buffer.from(url, "utf-8"));
169    }
170 
171    return extractImageFromResponse(result);
172  } finally {
173    clearTimeout(timeout);
174  }
175}
176