import type { AgentTool } from "@jaeswift/jae-agent-core";
import type { ToolResultMessage } from "@jaeswift/jae-ai";
import { Type } from "@sinclair/typebox";
import { html } from "lit";
import { Image } from "lucide";
import { getAppStorage } from "../storage/app-storage.js";
import { registerToolRenderer, renderHeader } from "./renderer-registry.js";
import type { ToolRenderer, ToolRenderResult } from "./types.js";

const imageGenSchema = Type.Object({
	prompt: Type.String({ description: "Image generation prompt describing what to create" }),
	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)" })),
});

export interface ImageGenDetails {
	dataUrl?: string;
	model: string;
	prompt: string;
	width: number;
	height: number;
	error?: string;
}

interface ImageGenParams {
	prompt: string;
	model?: string;
	width?: number;
	height?: number;
	steps?: number;
}

export const imageGenTool: AgentTool<typeof imageGenSchema, ImageGenDetails> = {
	name: "generate_image",
	label: "Generate Image",
	description: "Generate an image using Venice AI image models. Displays inline in chat.",
	parameters: imageGenSchema,
	async execute(toolCallId, params, signal) {
		const { prompt, model = "fluently-xl", width = 1024, height = 1024, steps = 20 } = params;
		const apiKey = await getAppStorage().providerKeys.get("venice");
		if (!apiKey) {
			return {
				content: [
					{
						type: "text",
						text: "Error: Venice API key not set. Add it in Settings > Providers & Models > Venice.",
					},
				],
				details: { model, prompt, width, height, error: "No API key" },
			};
		}
		const res = await fetch("https://api.venice.ai/api/v1/image/generate", {
			method: "POST",
			headers: { "Content-Type": "application/json", Authorization: `Bearer ${apiKey}` },
			body: JSON.stringify({ model, prompt, width, height, steps, return_binary: false, safe_mode: false }),
			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: { model, prompt, width, height, error: err },
			};
		}
		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: { model, prompt, width, height, error: "No image data" },
			};
		}
		const dataUrl = `data:image/png;base64,${b64}`;
		return {
			content: [{ type: "text", text: `Image generated successfully. Model: ${model}, Size: ${width}x${height}` }],
			details: { dataUrl, model, prompt, width, height },
		};
	},
};

class ImageGenRenderer implements ToolRenderer<ImageGenParams, ImageGenDetails> {
	render(
		params: ImageGenParams | undefined,
		result: ToolResultMessage<ImageGenDetails> | undefined,
	): ToolRenderResult {
		const state = result ? (result.isError ? "error" : "complete") : "inprogress";
		if (result?.details?.dataUrl) {
			const d = result.details;
			return {
				content: html`
          <div class="flex flex-col gap-3">
            ${renderHeader(state, Image, "Image Generated")}
            <img src=${d.dataUrl} alt=${d.prompt}
              class="rounded-lg max-w-full border border-border shadow-sm"
              style="max-height:512px;object-fit:contain;"
            />
            <div class="flex items-center justify-between text-xs text-muted-foreground">
              <span>${d.model} · ${d.width}×${d.height}</span>
              <a href=${d.dataUrl} download="generated.png" class="text-primary hover:underline">⬇ Download</a>
            </div>
            <div class="text-xs text-foreground italic">${d.prompt}</div>
          </div>`,
				isCustom: false,
			};
		}
		return {
			content: renderHeader(state, Image, `Generating image: ${params?.prompt?.slice(0, 50) ?? "..."}`),
			isCustom: false,
		};
	}
}

registerToolRenderer("generate_image", new ImageGenRenderer());

export function createImageGenTool(): AgentTool<typeof imageGenSchema, ImageGenDetails> {
	return imageGenTool;
}
