import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { generateImage } from "ai";
import index from "./index.html";

type Size = `${number}x${number}`;

async function generateWithReference({
	baseURL,
	apiKey,
	model,
	prompt,
	size,
	referenceImages,
}: {
	baseURL: string;
	apiKey: string;
	model: string;
	prompt: string;
	size: Size;
	referenceImages: string[];
}): Promise<string[]> {
	const form = new FormData();
	form.append("model", model);
	form.append("prompt", prompt);
	form.append("size", size);

	for (let i = 0; i < referenceImages.length; i++) {
		const dataUrl = referenceImages[i];
		if (!dataUrl) continue;
		const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/);
		if (!match) continue;
		const mime = match[1]!;
		const b64 = match[2]!;
		const bytes = Buffer.from(b64, "base64");
		const ext = mime.split("/")[1] ?? "png";
		form.append("image", new Blob([bytes], { type: mime }), `ref-${i}.${ext}`);
	}

	const url = `${baseURL.replace(/\/+$/, "")}/images/edits`;
	const res = await fetch(url, {
		method: "POST",
		headers: { Authorization: `Bearer ${apiKey}` },
		body: form,
	});

	if (!res.ok) {
		const text = await res.text().catch(() => "");
		throw new Error(`Upstream ${res.status}: ${text || res.statusText}`);
	}

	const data = (await res.json()) as {
		data?: Array<{ b64_json?: string }>;
	};

	return (data.data ?? [])
		.map((item) => (item.b64_json ? `data:image/png;base64,${item.b64_json}` : null))
		.filter((s): s is string => s !== null);
}

const server = Bun.serve({
	hostname: "0.0.0.0",
	routes: {
		"/": index,
		"/api/generate": {
			POST: async (req) => {
				try {
					const { baseURL, apiKey, model, prompt, size, referenceImages } =
						(await req.json()) as {
							baseURL?: string;
							apiKey?: string;
							model?: string;
							prompt?: string;
							size?: Size;
							referenceImages?: string[];
						};

					if (!baseURL || !apiKey || !model || !prompt) {
						return Response.json(
							{ error: "baseURL, apiKey, model, prompt are required" },
							{ status: 400 },
						);
					}

					if (Array.isArray(referenceImages) && referenceImages.length > 0) {
						const images = await generateWithReference({
							baseURL,
							apiKey,
							model,
							prompt,
							size: size ?? "1024x1024",
							referenceImages,
						});
						return Response.json({ images });
					}

					const provider = createOpenAICompatible({
						name: "custom",
						apiKey,
						baseURL,
					});

					const { images } = await generateImage({
						model: provider.imageModel(model),
						prompt,
						size: size || "1024x1024",
					});

					const out = images.map((img) => {
						const mediaType = (img as { mediaType?: string }).mediaType ?? "image/png";
						const base64 = (img as { base64?: string }).base64;
						return base64 ? `data:${mediaType};base64,${base64}` : null;
					}).filter(Boolean);

					return Response.json({ images: out });
				} catch (err) {
					const message = err instanceof Error ? err.message : String(err);
					console.error("[generate] error:", err);
					return Response.json({ error: message }, { status: 500 });
				}
			},
		},
	},
	development: {
		hmr: true,
		console: true,
	},
});

console.log(`Listening on ${server.url}`);