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, outputDir] = process.argv.slice(2);

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

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

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

// Exponential backoff retry logic
async function fetchWithRetry(url, options, maxRetries = 5) {
  for (let attempt = 0; attempt < maxRetries; attempt++) {
    try {
      const response = await fetch(url, options);
      const text = await response.text();

      // Success
      if (response.ok) {
        return { response, text };
      }

      // Rate limit (429) or server error (500+) - retry with backoff
      if (response.status === 429 || response.status >= 500) {
        if (attempt < maxRetries - 1) {
          // Calculate exponential backoff with jitter
          const baseDelay = Math.pow(2, attempt) * 1000; // 1s, 2s, 4s, 8s, 16s
          const jitter = Math.random() * 1000; // Add up to 1s random jitter
          const delay = baseDelay + jitter;

          console.log(
            `Rate limit hit (attempt ${attempt + 1}/${maxRetries}). Retrying in ${Math.round(delay / 1000)}s...`
          );
          await new Promise((resolve) => setTimeout(resolve, delay));
          continue;
        }
      }

      // Non-retryable error or final attempt failed
      return { response, text };
    } catch (err) {
      // Network error - retry with backoff
      if (attempt < maxRetries - 1) {
        const baseDelay = Math.pow(2, attempt) * 1000;
        const jitter = Math.random() * 1000;
        const delay = baseDelay + jitter;

        console.log(
          `Network error (attempt ${attempt + 1}/${maxRetries}). Retrying in ${Math.round(delay / 1000)}s...`
        );
        await new Promise((resolve) => setTimeout(resolve, delay));
        continue;
      }
      throw err;
    }
  }
}

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, text } = await fetchWithRetry(
    ENDPOINT,
    {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "x-goog-api-key": "PLACEHOLDER_TOKEN",
      },
      body: JSON.stringify(body),
    },
    5 // Max 5 retries
  );

  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
  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);
});