#!/usr/bin/env node
import { existsSync } from "node:fs";
import fs from "node:fs/promises";
import path from "node:path";
import process from "node:process";

const MODEL = "gemini-2.0-flash-exp";
const API_BASE = "https://generativelanguage.googleapis.com/v1beta/models";
const DEFAULT_OUTPUT_DIR = path.resolve("assets");

const WEB_STYLE_GUIDANCE = [
  "This image is intended for use in a web application.",
  "Ensure the result is clean, modern, and production-ready.",
  "Use crisp edges and consistent lighting suitable for UI/web contexts.",
  "Output a PNG with transparency where appropriate.",
].join(" ");

// --- Argument parsing ---

const args = process.argv.slice(2);

function getFlag(name) {
  const i = args.indexOf(`--${name}`);
  if (i === -1 || i + 1 >= args.length) return undefined;
  return args[i + 1];
}

function hasFlag(name) {
  return args.includes(`--${name}`);
}

if (hasFlag("help") || !args.length) {
  console.log(`
Usage: node generate-image.mjs --prompt "<prompt>" [options]

Options:
  --prompt    Image generation prompt (required)
  --output    Output file path (default: assets/generated-<timestamp>.png)
  --style     Additional style guidance to append to the prompt
  --raw       Skip default web style guidance, use prompt as-is
  --verbose   Show request/response details
  --help      Show this help message

Examples:
  node generate-image.mjs --prompt "a flat icon of a shopping cart, white on transparent"
  node generate-image.mjs --prompt "hero banner with abstract gradient" --output assets/hero.png
  node generate-image.mjs --prompt "login page illustration" --style "minimalist, pastel colors"
`);
  process.exit(0);
}

// --- Load env ---

const ENV_FILE = path.resolve(process.cwd(), ".env.tools");

async function loadEnvFile(targetPath) {
  if (!existsSync(targetPath)) return;
  const content = await fs.readFile(targetPath, "utf8");
  for (const rawLine of content.split(/\r?\n/)) {
    const line = rawLine.trim();
    if (!line || line.startsWith("#")) continue;
    const equals = line.indexOf("=");
    if (equals === -1) continue;
    const key = line.slice(0, equals).trim();
    if (!key) continue;
    let value = line.slice(equals + 1).trim();
    if (
      (value.startsWith('"') && value.endsWith('"')) ||
      (value.startsWith("'") && value.endsWith("'"))
    ) {
      value = value.slice(1, -1);
    }
    if (!(key in process.env)) {
      process.env[key] = value;
    }
  }
}

await loadEnvFile(ENV_FILE);

const apiKey = process.env.GEMINI_API_KEY;
if (!apiKey) {
  console.error("GEMINI_API_KEY not found in the environment (check .env.tools).");
  process.exit(1);
}

// --- Build prompt ---

const explicitPrompt = getFlag("prompt");
if (!explicitPrompt) {
  console.error("Missing --prompt flag. Run with --help for usage.");
  process.exit(1);
}

const styleExtra = getFlag("style");
const raw = hasFlag("raw");
const verbose = hasFlag("verbose");

let finalPrompt = explicitPrompt;
if (!raw) {
  finalPrompt = `${WEB_STYLE_GUIDANCE}\n\n${explicitPrompt}`;
}
if (styleExtra) {
  finalPrompt += `\nStyle: ${styleExtra}`;
}

// --- Output path ---

const outputFlag = getFlag("output");
const outputPath = outputFlag
  ? path.resolve(outputFlag)
  : path.join(DEFAULT_OUTPUT_DIR, `generated-${Date.now()}.png`);

// --- API call ---

async function generateImage() {
  const url = `${API_BASE}/${MODEL}:generateContent?key=${apiKey}`;

  const body = {
    contents: [
      {
        parts: [{ text: finalPrompt }],
      },
    ],
    generationConfig: {
      responseModalities: ["IMAGE", "TEXT"],
    },
  };

  if (verbose) {
    console.log(`Model: ${MODEL}`);
    console.log(`Prompt: ${finalPrompt}`);
    console.log(`Output: ${outputPath}`);
    console.log("Calling Gemini API...");
  }

  const response = await fetch(url, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });

  if (!response.ok) {
    const text = await response.text();
    throw new Error(`Gemini API error (${response.status}): ${text}`);
  }

  const payload = await response.json();

  if (verbose) {
    const textParts = payload.candidates?.[0]?.content?.parts?.filter((p) => p.text) || [];
    if (textParts.length) {
      console.log("Model response text:", textParts.map((p) => p.text).join("\n"));
    }
  }

  // Find the image part in the response
  const imagePart = payload.candidates?.[0]?.content?.parts?.find(
    (p) => p.inlineData?.mimeType?.startsWith("image/")
  );

  if (!imagePart?.inlineData?.data) {
    // Log what we got for debugging
    const parts = payload.candidates?.[0]?.content?.parts || [];
    const partTypes = parts.map((p) =>
      p.text ? "text" : p.inlineData ? `inlineData(${p.inlineData.mimeType})` : "unknown"
    );
    throw new Error(
      `No image data in Gemini response. Got parts: [${partTypes.join(", ")}]. ` +
        "The model may not support image generation, or the prompt was rejected."
    );
  }

  return Buffer.from(imagePart.inlineData.data, "base64");
}

// --- Main ---

async function main() {
  const outputDir = path.dirname(outputPath);
  if (!existsSync(outputDir)) {
    await fs.mkdir(outputDir, { recursive: true });
  }

  const imageBuffer = await generateImage();
  await fs.writeFile(outputPath, imageBuffer);

  console.log(`Generated image saved to: ${outputPath}`);
  console.log("Inspect the image and rerun with a refined prompt if needed.");
}

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