import fs from "fs";
import path from "path";

const ENDPOINT =
  "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent";

const [prompt, imagePath] = process.argv.slice(2);

if (!prompt) {
  console.error("Error: prompt is required");
  process.exit(1);
}

const isEdit = imagePath && imagePath.trim();

async function main() {
  console.log(isEdit ? `Editing image: "${imagePath}"` : `Creating image: "${prompt}"`);

  // Build request parts
  const parts = [{ text: prompt }];

  // Add source image if editing
  if (isEdit) {
    const imageBuffer = fs.readFileSync(imagePath);
    const base64Image = imageBuffer.toString("base64");
    const ext = path.extname(imagePath).toLowerCase();
    const mimeType =
      ext === ".png"
        ? "image/png"
        : ext === ".webp"
          ? "image/webp"
          : ext === ".gif"
            ? "image/gif"
            : "image/jpeg";

    parts.push({
      inline_data: {
        mime_type: mimeType,
        data: base64Image,
      },
    });
    console.log(`Loaded source image (${mimeType})`);
  }

  const body = { contents: [{ parts }] };

  console.log("Calling Gemini API...");

  const response = await fetch(ENDPOINT, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "x-goog-api-key": "PLACEHOLDER_TOKEN",
    },
    body: JSON.stringify(body),
  });

  const text = await response.text();
  if (!response.ok) {
    console.error(`API error (${response.status}): ${text}`);
    process.exit(1);
  }

  const data = JSON.parse(text);

  // Extract image from response (handle both casing styles)
  const responseParts = data?.candidates?.[0]?.content?.parts || [];
  const imagePart = responseParts.find(
    (x) => (x.inlineData && x.inlineData.data) || (x.inline_data && x.inline_data.data)
  );

  if (!imagePart) {
    console.error("No image returned from API");
    console.error(JSON.stringify(data, null, 2));
    process.exit(1);
  }

  const inline = imagePart.inlineData || imagePart.inline_data;

  // Save the generated image
  const outputDir = "./documents/nanobanana";
  fs.mkdirSync(outputDir, { recursive: true });

  // Generate filename: slug from prompt for create, original-edited for edit
  let outputName;
  if (isEdit) {
    const originalName = path.basename(imagePath, path.extname(imagePath));
    const timestamp = Date.now();
    outputName = `${originalName}-edited-${timestamp}`;
  } else {
    const slug = prompt
      .toLowerCase()
      .replace(/[^a-z0-9]+/g, "-")
      .replace(/^-|-$/g, "")
      .slice(0, 50);
    outputName = slug || "image";
  }

  const outputPath = path.join(outputDir, `${outputName}.png`);
  fs.writeFileSync(outputPath, Buffer.from(inline.data, "base64"));

  console.log(`✓ Saved: ${outputPath}`);
  console.log(
    JSON.stringify({
      success: true,
      path: outputPath,
      prompt: prompt,
      isEdit: !!isEdit,
    })
  );
}

main().catch((err) => {
  console.error("Failed:", err.message);
  process.exit(1);
});