#!/usr/bin/env node

import { mkdir, readFile, writeFile } from "node:fs/promises";
import path from "node:path";

const DEFAULT_UPLOAD_URL =
  "https://lms.laitool.cn/lms/FileUpload/FileUpload/01564ca6c907eab817058e5356bb40b57cd3f3171adbbc9ea49be0de448cc68e";

function argValue(name, fallback = undefined) {
  const index = process.argv.indexOf(name);
  if (index === -1 || index + 1 >= process.argv.length) return fallback;
  return process.argv[index + 1];
}

function requireArg(name) {
  const value = argValue(name);
  if (!value) {
    throw new Error(`Missing required argument: ${name}`);
  }
  return value;
}

async function readJson(filePath) {
  return JSON.parse(await readFile(filePath, "utf8"));
}

function sanitizeName(value) {
  return String(value)
    .trim()
    .replace(/[\\/:*?"<>|]+/g, "_")
    .replace(/\s+/g, "_")
    .slice(0, 120);
}

function extensionFromMime(contentType) {
  if (contentType === "image/png") return ".png";
  if (contentType === "image/webp") return ".webp";
  return ".jpg";
}

function extractInlineImage(response) {
  const candidates = response?.candidates ?? [];
  for (const candidate of candidates) {
    const parts = candidate?.content?.parts ?? [];
    for (const part of parts) {
      const inlineData = part.inlineData ?? part.inline_data;
      if (inlineData?.data) {
        return {
          base64: inlineData.data,
          contentType: inlineData.mimeType ?? inlineData.mime_type ?? "image/jpeg",
        };
      }
    }
  }
  throw new Error("Image model response did not contain inline image data.");
}

function extractUploadUrl(resultText) {
  try {
    const json = JSON.parse(resultText);
    return (
      json.url ??
      json.fileUrl ??
      json.fileURL ??
      json.data?.url ??
      json.data?.fileUrl ??
      json.data?.fileURL ??
      json.data?.path ??
      json.result?.url ??
      json.result?.fileUrl ??
      ""
    );
  } catch {
    const match = resultText.match(/https?:\/\/[^\s"'<>]+/);
    return match?.[0] ?? "";
  }
}

async function generateImage({ prompt, model, apiKey, aspectRatio, contentType }) {
  const endpoint = `https://magic666.top/v1beta/models/${encodeURIComponent(
    model,
  )}:generateContent?key=${encodeURIComponent(apiKey)}`;


  const response = await fetch(endpoint, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify({
      contents: [{ role: "user", parts: [{ text: prompt }] }],
      generationConfig: {
        responseModalities: ["IMAGE"],
        ...(aspectRatio ? { aspectRatio } : {}),
      },
    }),
  });

  const responseText = await response.text();
  if (!response.ok) {
    throw new Error(`Image generation failed: ${response.status} ${responseText}`);
  }

  const image = extractInlineImage(JSON.parse(responseText));
  return {
    base64: image.base64,
    contentType: contentType ?? image.contentType,
  };
}

async function uploadImage({ base64, fileName, contentType, uploadUrl }) {
  const originalSize = Buffer.byteLength(Buffer.from(base64, "base64")).toString();
  const response = await fetch(uploadUrl, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    redirect: "follow",
    body: JSON.stringify({
      file: base64,
      fileName,
      contentType,
      metadata: {
        uploadTime: new Date().toISOString(),
        originalSize,
      },
    }),
  });

  const resultText = await response.text();
  if (!response.ok) {
    throw new Error(`Image upload failed: ${response.status} ${resultText}`);
  }

  return {
    raw: resultText,
    url: extractUploadUrl(resultText),
  };
}

async function main() {
  const jobsPath = requireArg("--jobs");
  const outDir = path.resolve(argValue("--out-dir", "outputs/generated_assets"));
  const manifestPath = path.resolve(argValue("--manifest", path.join(outDir, "uploaded_assets.json")));
  // const apiKey = process.env.GEMINI_API_KEY ?? process.env.GOOGLE_API_KEY;

  const apiKey = "sk-QLEqdjgAE457aFsxmvUHdoj0EnyBcFaX23I4DxCJ8bgnrJNw"
  const model = argValue("--model", process.env.NANO_BANANA_MODEL ?? "nano-banana-2");
  const uploadUrl = argValue("--upload-url", process.env.IMAGE_UPLOAD_URL ?? DEFAULT_UPLOAD_URL);

  if (!apiKey) {
    throw new Error("Set GEMINI_API_KEY or GOOGLE_API_KEY before running image generation.");
  }
  const jobs = await readJson(jobsPath);
  if (!Array.isArray(jobs)) {
    throw new Error("Jobs file must be a JSON array.");
  }

  await mkdir(outDir, { recursive: true });
  const results = [];

  for (const [index, job] of jobs.entries()) {
    if (!job.prompt) {
      throw new Error(`Job ${index + 1} is missing prompt.`);
    }

    const generated = await generateImage({
      prompt: job.prompt,
      model: "gemini-3-pro-image-preview",
      apiKey,
      aspectRatio: job.aspectRatio,
      contentType: job.contentType,
    });

    const fileName =
      job.fileName ??
      `${String(index + 1).padStart(3, "0")}_${sanitizeName(job.internalRef ?? job.name ?? "asset")}${extensionFromMime(
        generated.contentType,
      )}`;
    const localPath = path.join(outDir, fileName);
    await writeFile(localPath, Buffer.from(generated.base64, "base64"));

    const uploaded = await uploadImage({
      base64: generated.base64,
      fileName,
      contentType: generated.contentType,
      uploadUrl,
    });

    results.push({
      internalRef: job.internalRef ?? "",
      type: job.type ?? "",
      promptName: job.promptName ?? job.name ?? "",
      fileName,
      localPath,
      contentType: generated.contentType,
      publicUrl: uploaded.url,
      uploadRawResponse: uploaded.raw,
      jimengUploadName: job.jimengUploadName ?? path.parse(fileName).name,
      referenceDuty: job.referenceDuty ?? "",
      status: uploaded.url ? "uploaded" : "uploaded_url_unparsed",
    });
  }

  await writeFile(manifestPath, `${JSON.stringify(results, null, 2)}\n`, "utf8");
  console.log(`Wrote ${results.length} uploaded asset records to ${manifestPath}`);
}

main().catch((error) => {
  console.error(error);
  process.exitCode = 1;
});