import { existsSync, mkdirSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import type { AgentTool } from "@jaeswift/jae-agent-core";
import { type Static, Type } from "@sinclair/typebox";

const imageGenSchema = Type.Object({
	prompt: Type.String({ description: "Image generation prompt" }),
	model: Type.Optional(Type.String({ description: "Venice image model (default: fluently-xl)" })),
	width: Type.Optional(Type.Number({ description: "Width in pixels (default: 1024)" })),
	height: Type.Optional(Type.Number({ description: "Height in pixels (default: 1024)" })),
	steps: Type.Optional(Type.Number({ description: "Inference steps (default: 20)" })),
	output_dir: Type.Optional(Type.String({ description: "Directory to save image (default: ./images)" })),
});

export type ImageGenToolInput = Static<typeof imageGenSchema>;

export interface ImageGenToolDetails {
	path: string;
	model: string;
	prompt: string;
	width: number;
	height: number;
}

export const imageGenTool: AgentTool<typeof imageGenSchema, ImageGenToolDetails> = {
	name: "generate_image",
	label: "Generate Image",
	description:
		"Generate an image using Venice AI image models and save it to disk. Uses VENICE_API_KEY or OPENAI_API_KEY env var.",
	parameters: imageGenSchema,
	async execute(toolCallId, params, signal) {
		const {
			prompt,
			model = "fluently-xl",
			width = 1024,
			height = 1024,
			steps = 20,
			output_dir = "./images",
		} = params;

		const apiKey = process.env.VENICE_API_KEY || process.env.OPENAI_API_KEY;
		if (!apiKey) {
			return {
				content: [{ type: "text", text: "Error: VENICE_API_KEY or OPENAI_API_KEY environment variable not set." }],
				details: { path: "", model, prompt, width, height },
			};
		}

		const body = JSON.stringify({
			model,
			prompt,
			width,
			height,
			steps,
			return_binary: false,
			safe_mode: false,
		});

		const res = await fetch("https://api.venice.ai/api/v1/image/generate", {
			method: "POST",
			headers: {
				"Content-Type": "application/json",
				Authorization: `Bearer ${apiKey}`,
			},
			body,
			signal: signal ?? AbortSignal.timeout(60000),
		});

		if (!res.ok) {
			const err = await res.text();
			return {
				content: [{ type: "text", text: `Image generation failed (${res.status}): ${err}` }],
				details: { path: "", model, prompt, width, height },
			};
		}

		const data = (await res.json()) as any;
		const b64 = data?.images?.[0];
		if (!b64) {
			return {
				content: [{ type: "text", text: "No image returned from Venice API." }],
				details: { path: "", model, prompt, width, height },
			};
		}

		if (!existsSync(output_dir)) mkdirSync(output_dir, { recursive: true });
		const filename = `gen-${Date.now()}.png`;
		const filepath = join(output_dir, filename);
		writeFileSync(filepath, Buffer.from(b64, "base64"));

		return {
			content: [
				{
					type: "text",
					text: `Image saved to: ${filepath}\nModel: ${model}\nPrompt: ${prompt}\nSize: ${width}x${height}`,
				},
			],
			details: { path: filepath, model, prompt, width, height },
		};
	},
};
