import { readFile } from "node:fs/promises";
import { extname, isAbsolute, join, resolve } from "node:path";

import { createOpenAICompatible } from "@ai-sdk/openai-compatible";
import { streamText } from "ai";

import type { Config } from "@/config";
import type { ImageJob } from "@/jobs";

type GeneratedImage = {
	bytes: Uint8Array;
	mediaType: string;
};

type InputImagePart = {
	type: "image";
	image: Uint8Array;
	mediaType: string;
};

const MARKDOWN_IMAGE_REGEX = /!\[[^\]]*\]\((https?:\/\/[^\s)]+)\)/;
const INPUT_IMAGES_DIR = "images";

function createProvider(config: Config) {
	return createOpenAICompatible({
		name: "flow2api",
		baseURL: config.baseURL,
		apiKey: config.apiKey,
	});
}

function mediaTypeFromPath(filePath: string): string {
	switch (extname(filePath).toLowerCase()) {
		case ".png":
			return "image/png";
		case ".jpg":
		case ".jpeg":
			return "image/jpeg";
		case ".webp":
			return "image/webp";
		case ".gif":
			return "image/gif";
		default:
			return "image/png";
	}
}

async function buildImageParts(
	imagePaths: string[],
): Promise<InputImagePart[]> {
	return Promise.all(
		imagePaths.map(async (imagePath) => ({
			type: "image" as const,
			image: new Uint8Array(
				await readFile(
					isAbsolute(imagePath)
						? imagePath
						: resolve(join(INPUT_IMAGES_DIR, imagePath)),
				),
			),
			mediaType: mediaTypeFromPath(imagePath),
		})),
	);
}

function writeSectionTitle(title: string, alreadyPrinted: boolean): void {
	if (alreadyPrinted) {
		return;
	}

	process.stdout.write(`\n${title}:\n\n`);
}

function extractImageUrl(text: string): string | null {
	const match = text.match(MARKDOWN_IMAGE_REGEX);
	return match?.[1] ?? null;
}

async function downloadImage(url: string): Promise<GeneratedImage> {
	const response = await fetch(url);

	if (!response.ok) {
		throw new Error(
			`Failed to download generated image: ${response.status} ${response.statusText}`,
		);
	}

	const mediaType =
		response.headers.get("content-type")?.split(";")[0]?.trim() || "image/png";

	return {
		bytes: new Uint8Array(await response.arrayBuffer()),
		mediaType,
	};
}

export async function streamPrompt(
	config: Config,
	job: Pick<ImageJob, "prompt" | "images">,
): Promise<GeneratedImage> {
	const provider = createProvider(config);
	const imageParts = await buildImageParts(job.images ?? []);
	const result = streamText({
		model: provider(config.model),
		messages: [
			{
				role: "user",
				content: [{ type: "text", text: job.prompt }, ...imageParts],
			},
		],
	});

	let reasoningStarted = false;
	let textStarted = false;
	let textOutput = "";

	for await (const chunk of result.fullStream) {
		if (chunk.type === "reasoning-delta") {
			writeSectionTitle("Thoughts", reasoningStarted);
			reasoningStarted = true;
			process.stdout.write(chunk.text);
			continue;
		}

		if (chunk.type === "text-delta") {
			writeSectionTitle("Response", textStarted);
			textStarted = true;
			textOutput += chunk.text;
			process.stdout.write(chunk.text);
		}
	}

	const imageUrl = extractImageUrl(textOutput);

	if (!imageUrl) {
		throw new Error("The streamed response did not include an image URL.");
	}

	process.stdout.write("\n\nDownloading image...\n");

	return downloadImage(imageUrl);
}
