const fs = require("fs");
const path = require("path");
const axios = require("axios");
const { S3Client, PutObjectCommand, DeleteObjectsCommand } = require("@aws-sdk/client-s3");

// NanoBanana2 生图接口，文档 https://api.wuyinkeji.com/doc/65
const NANOBANANA_URL = "https://api.wuyinkeji.com/api/async/image_nanoBanana2";
// 全模型通用结果详情接口，文档见 https://api.wuyinkeji.com/doc/47
const RESULT_DETAIL_URL = "https://api.wuyinkeji.com/api/async/detail";
const CONFIG_FILE_NAME = "config_nano_banana.json";

function loadConfig(baseDir) {
  const configPath = path.join(baseDir, CONFIG_FILE_NAME);
  if (!fs.existsSync(configPath)) return {};
  try {
    const raw = fs.readFileSync(configPath, "utf8");
    return JSON.parse(raw) || {};
  } catch (e) {
    console.log("读取配置文件失败，将忽略配置文件。错误:", e.message);
    return {};
  }
}

function ensureApiKey(config) {
  const key = config.api_key || process.env.WUYIN_API_KEY;
  if (!key) {
    throw new Error(
      "未找到 API 密钥，请在 config_nano_banana.json 中填写 api_key，" +
        "或设置环境变量 WUYIN_API_KEY。"
    );
  }
  return key;
}

function ensurePrompt(config) {
  const prompt = config.prompt;
  if (!prompt || !prompt.trim()) {
    throw new Error(
      "未在 config_nano_banana.json 中找到有效的 prompt，请先在文件中填写默认提示词。"
    );
  }
  return prompt.trim();
}

/** 检查是否配置了 R2 图床（用于 01 参考图上传） */
function isR2Configured(config) {
  return (
    config &&
    config.r2_account_id &&
    config.r2_access_key_id &&
    config.r2_secret_access_key &&
    config.r2_bucket &&
    config.r2_public_url
  );
}

/** 创建 R2（S3 兼容）客户端 */
function createR2Client(config) {
  const accountId = config.r2_account_id;
  const endpoint = `https://${accountId}.r2.cloudflarestorage.com`;
  return new S3Client({
    region: "auto",
    endpoint,
    // 对不少 S3 兼容服务更稳：强制使用 path-style
    // 形如 https://<accountId>.r2.cloudflarestorage.com/<bucket>/<key>
    forcePathStyle: true,
    credentials: {
      accessKeyId: config.r2_access_key_id,
      secretAccessKey: config.r2_secret_access_key,
    },
  });
}

/**
 * 将 01 目录中的参考图上传到 R2，返回可公网访问的 URL 列表及本次上传的 key 列表（用于事后删除）
 */
async function uploadRefImagesToR2({ client, bucket, publicBaseUrl, inputDir, maxFiles = 14 }) {
  if (!fs.existsSync(inputDir)) return { urls: [], keys: [] };

  const exts = new Set([".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"]);
  const files = fs
    .readdirSync(inputDir)
    .map((name) => path.join(inputDir, name))
    .filter((p) => exts.has(path.extname(p).toLowerCase()))
    .slice(0, maxFiles);

  const runId = Date.now().toString(36) + "-" + Math.random().toString(36).slice(2, 8);
  const prefix = `ref/${runId}`;
  const urls = [];
  const keys = [];

  for (const filePath of files) {
    const name = path.basename(filePath);
    const key = `${prefix}/${name}`;
    const body = fs.readFileSync(filePath);
    const ext = path.extname(name).toLowerCase();
    const contentType = { ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".png": "image/png", ".webp": "image/webp", ".bmp": "image/bmp", ".gif": "image/gif" }[ext] || "application/octet-stream";

    await client.send(
      new PutObjectCommand({
        Bucket: bucket,
        Key: key,
        Body: body,
        ContentType: contentType,
      })
    );
    const base = publicBaseUrl.replace(/\/$/, "");
    urls.push(`${base}/${key}`);
    keys.push(key);
  }
  return { urls, keys };
}

/** 从 R2 删除指定 key 的对象 */
async function deleteR2Objects({ client, bucket, keys }) {
  if (!keys || keys.length === 0) return;
  await client.send(
    new DeleteObjectsCommand({
      Bucket: bucket,
      Delete: { Objects: keys.map((Key) => ({ Key })) },
    })
  );
}

async function collectReferenceImages(inputDir, maxFiles = 14) {
  if (!fs.existsSync(inputDir)) return [];

  const exts = new Set([".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"]);
  const files = fs
    .readdirSync(inputDir)
    .map((name) => path.join(inputDir, name))
    .filter((p) => exts.has(path.extname(p).toLowerCase()))
    .slice(0, maxFiles);

  const result = [];
  for (const file of files) {
    const buf = fs.readFileSync(file);
    result.push(buf.toString("base64"));
  }
  return result;
}

async function createTask({ apiKey, prompt, size, aspectRatio, refB64List }) {
  const headers = {
    Authorization: apiKey,
    "Content-Type": "application/json;charset=utf-8;",
  };

  const body = {
    prompt,
    size,
    aspectRatio,
    key: apiKey,
  };
  if (refB64List && refB64List.length > 0) {
    body.urls = refB64List;
  }

  const resp = await axios.post(NANOBANANA_URL, body, {
    headers,
    timeout: 30000,
  });
  const payload = resp.data;

  if (!payload || payload.code !== 200) {
    throw new Error("创建任务失败: " + JSON.stringify(payload));
  }

  const data = payload.data || {};
  const taskId = data.id;
  if (!taskId) {
    throw new Error("返回数据中缺少任务 id: " + JSON.stringify(payload));
  }
  return taskId;
}

function extractImageUrls(dataObj) {
  if (!dataObj) return [];

  const candidates = [];
  if (Array.isArray(dataObj)) {
    for (const v of dataObj) {
      if (v) candidates.push(String(v));
    }
  } else if (typeof dataObj === "object") {
    // 兼容多种字段命名，包括结果详情返回的 result 数组
    for (const key of ["img_url", "img_urls", "image_urls", "urls", "images", "result"]) {
      if (dataObj[key]) {
        const val = dataObj[key];
        if (typeof val === "string") {
          candidates.push(val);
        } else if (Array.isArray(val)) {
          for (const v of val) {
            if (v) candidates.push(String(v));
          }
        }
      }
    }
  }

  const seen = new Set();
  const result = [];
  for (const url of candidates) {
    if (!seen.has(url)) {
      seen.add(url);
      result.push(url);
    }
  }
  return result;
}

async function queryResult({ apiKey, taskId, pollIntervalMs, maxWaitMs }) {
  const headers = {
    Authorization: apiKey,
    "Content-Type": "application/x-www-form-urlencoded;charset=utf-8;",
  };
  const params = { key: apiKey, id: taskId };
  const start = Date.now();

  while (true) {
    const resp = await axios.get(RESULT_DETAIL_URL, {
      params,
      headers,
      timeout: 30000,
    });
    const payload = resp.data;

    if (!payload || payload.code !== 200) {
      const msg = (payload && payload.msg) || "未知错误";
      throw new Error(`查询任务失败: code=${payload && payload.code}, msg=${msg}`);
    }

    const data = payload.data || {};
    const status = data.status;

    if (status === 2) {
      const urls = extractImageUrls(data);
      return urls;
    }
    if (status === 3) {
      const reason = data.message || payload.msg || "未知原因";
      throw new Error("任务生成失败: " + reason);
    }

    const urls = extractImageUrls(data);
    if (urls.length > 0) return urls;

    if (Date.now() - start > maxWaitMs) {
      throw new Error("等待任务结果超时，请稍后在控制台或结果查询接口自行确认。");
    }

    await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
  }
}

async function downloadImages({ urls, outputDir, taskId }) {
  if (!fs.existsSync(outputDir)) {
    fs.mkdirSync(outputDir, { recursive: true });
  }

  for (let i = 0; i < urls.length; i++) {
    const url = urls[i];
    try {
      const resp = await axios.get(url, {
        responseType: "arraybuffer",
        timeout: 60000,
      });

      let ext = ".jpg";
      for (const cand of [".png", ".jpeg", ".jpg", ".webp", ".bmp", ".gif"]) {
        if (url.toLowerCase().includes(cand)) {
          ext = cand;
          break;
        }
      }

      const filename = path.join(outputDir, `${taskId}_${i + 1}${ext}`);
      fs.writeFileSync(filename, resp.data);
      console.log("已保存:", filename);
    } catch (e) {
      console.log("下载失败:", url, "| 错误:", e.message);
    }
  }
}

async function main() {
  const baseDir = __dirname;
  const config = loadConfig(baseDir);

  const apiKey = ensureApiKey(config);
  const prompt = ensurePrompt(config);

  const size = config.size || "1K";
  const aspectRatio = config.aspectRatio || "auto";
  const pollIntervalMs = (config.poll_interval_seconds || 5) * 1000;
  const maxWaitMs = (config.max_wait_seconds || 300) * 1000;

  const inputDir = path.join(baseDir, config.input_dir || "01");
  const outputDir = path.join(baseDir, config.output_dir || "save");

  let combinedUrls = [];
  let r2KeysToDelete = [];
  const useR2 = isR2Configured(config);

  if (useR2) {
    // 使用 R2 图床：将 01 目录的图片上传到 R2，用返回的 URL 作为参考图
    const client = createR2Client(config);
    const { urls: r2Urls, keys } = await uploadRefImagesToR2({
      client,
      bucket: config.r2_bucket,
      publicBaseUrl: config.r2_public_url,
      inputDir,
      maxFiles: 14,
    });
    r2KeysToDelete = keys;
    if (r2Urls.length > 0) {
      console.log(`已上传 ${r2Urls.length} 张参考图到 R2 图床。`);
    }
    combinedUrls = [...r2Urls];
  } else {
    // 未配置 R2：本地 01 转为 base64
    const refB64List = await collectReferenceImages(inputDir);
    if (refB64List.length > 0) {
      console.log(`已从目录 ${inputDir} 读取 ${refB64List.length} 张本地参考图（base64）。`);
    }
    combinedUrls = [...refB64List];
  }

  // 配置中的参考图 URL（直接透传给接口）
  let extraUrls = [];
  if (typeof config.reference_urls === "string") {
    extraUrls = [config.reference_urls];
  } else if (Array.isArray(config.reference_urls)) {
    extraUrls = config.reference_urls.filter((u) => !!u).map(String);
  }
  combinedUrls = [...combinedUrls, ...extraUrls];

  if (extraUrls.length > 0) {
    console.log(`已从配置文件读取 ${extraUrls.length} 个参考图 URL。`);
  }
  if (combinedUrls.length === 0) {
    console.log(`未找到任何参考图，将仅根据提示词生成。`);
  }

  try {
    console.log("正在创建 NanoBanana2 任务...");
    const taskId = await createTask({
      apiKey,
      prompt,
      size,
      aspectRatio,
      refB64List: combinedUrls,
    });
    console.log("任务已创建，任务 id:", taskId);

    console.log("开始轮询任务结果...");
    const urls = await queryResult({
      apiKey,
      taskId,
      pollIntervalMs,
      maxWaitMs,
    });

    if (!urls || urls.length === 0) {
      console.log("任务完成，但未在结果中找到图片 URL，请登录速创API控制台或检查结果详情接口。");
      return;
    }

    console.log(`任务完成，共获取到 ${urls.length} 个图片 URL，开始下载...`);
    await downloadImages({ urls, outputDir, taskId });
    console.log("全部处理完成。");
  } finally {
    if (useR2 && r2KeysToDelete.length > 0) {
      try {
        const client = createR2Client(config);
        await deleteR2Objects({ client, bucket: config.r2_bucket, keys: r2KeysToDelete });
        console.log("已删除本次上传的 R2 参考图。");
      } catch (e) {
        console.warn("删除 R2 参考图失败:", e.message);
      }
    }
  }
}

main().catch((err) => {
  console.error("执行出错:", err.message || err);
  process.exit(1);
});