/* 
    Script to generate test API json from the ComfyUI_examples repo. 
    Requires the repo to be cloned to the tests-ui directory or specified via the EXAMPLE_REPO_PATH env var.
*/

import chalk from "chalk";
import fs from "fs";
import path from "path";
import { fileURLToPath } from "node:url";
import { getFromPngBuffer } from "@/scripts/metadata/png";
import { getFromFlacBuffer } from "@/scripts/metadata/flac";
import dotenv from "dotenv";
dotenv.config();

const dirname = path.dirname(fileURLToPath(import.meta.url));
const repoPath =
  process.env.EXAMPLE_REPO_PATH || path.resolve(dirname, "ComfyUI_examples");
const workflowsPath = path.resolve(dirname, "workflows", "examples");

if (!fs.existsSync(repoPath)) {
  console.error(
    `ComfyUI_examples repo not found. Please clone this to ${repoPath} or set the EXAMPLE_REPO_PATH env var (see .env_example) and re-run.`
  );
}

if (!fs.existsSync(workflowsPath)) {
  await fs.promises.mkdir(workflowsPath);
}

async function* getFiles(
  dir: string,
  ...exts: string[]
): AsyncGenerator<string, void, void> {
  const dirents = await fs.promises.readdir(dir, { withFileTypes: true });
  for (const dirent of dirents) {
    const res = path.resolve(dir, dirent.name);
    if (dirent.isDirectory()) {
      yield* getFiles(res, ...exts);
    } else if (exts.includes(path.extname(res))) {
      yield res;
    }
  }
}

async function validateMetadata(metadata: Record<string, string>) {
  const check = (prop: "prompt" | "workflow") => {
    const v = metadata?.[prop];
    if (!v) throw `${prop} not found in metadata`;
    try {
      JSON.parse(v);
    } catch (error) {
      throw `${prop} invalid json: ${error.message}`;
    }
    return v;
  };

  return { prompt: check("prompt"), workflow: check("workflow") };
}

async function hasExampleChanged(
  existingFilePath: string,
  exampleJson: string
) {
  return exampleJson !== (await fs.promises.readFile(existingFilePath, "utf8"));
}

// Example images to ignore as they don't contain workflows
const ignore = [
  "unclip_sunset.png",
  "unclip_mountains.png",
  "inpaint_yosemite_inpaint_example.png",
  "controlnet_shark_depthmap.png",
  "controlnet_pose_worship.png",
  "controlnet_pose_present.png",
  "controlnet_input_scribble_example.png",
  "controlnet_house_scribble.png",
];

// Find all existing examples so we can check if any are removed/changed
const existing = new Set(
  (await fs.promises.readdir(workflowsPath, { withFileTypes: true }))
    .filter((d) => d.isFile())
    .map((d) => path.resolve(workflowsPath, d.name))
);

const results = {
  new: [],
  changed: [],
  unchanged: [],
  missing: [],
  failed: [],
};

let total = 0;
for await (const file of getFiles(repoPath, ".png", ".flac")) {
  const cleanedName = path
    .relative(repoPath, file)
    .replaceAll("/", "_")
    .replaceAll("\\", "_");

  if (ignore.includes(cleanedName)) continue;
  total++;

  let metadata: { prompt: string; workflow: string };
  try {
    const { buffer } = await fs.promises.readFile(file);
    switch (path.extname(file)) {
      case ".png":
        metadata = await validateMetadata(getFromPngBuffer(buffer));
        break;
      case ".flac":
        metadata = await validateMetadata(getFromFlacBuffer(buffer));
        break;
    }

    const outPath = path.resolve(workflowsPath, cleanedName + ".json");
    const exampleJson = JSON.stringify(metadata);
    if (existing.has(outPath)) {
      existing.delete(outPath);
      if (await hasExampleChanged(outPath, exampleJson)) {
        results.changed.push(outPath);
      } else {
        // Unchanged, no point in re-saving
        results.unchanged.push(outPath);
        continue;
      }
    } else {
      results.new.push(outPath);
    }

    await fs.promises.writeFile(outPath, exampleJson, "utf8");
  } catch (error) {
    results.failed.push({ file, error });
  }
}

// Any workflows left in the existing set are now missing, these will want checking and manually removing
results.missing.push(...existing);

const c = (v: number, gt0: "red" | "yellow" | "green") =>
  chalk[v > 0 ? gt0 : "gray"](v);

console.log(`Processed ${chalk.green(total)} examples`);
console.log(`  ${chalk.gray(results.unchanged.length)} unchanged`);
console.log(`  ${c(results.changed.length, "yellow")} changed`);
console.log(`  ${c(results.new.length, "green")} new`);
console.log(`  ${c(results.missing.length, "red")} missing`);
console.log(`  ${c(results.failed.length, "red")} failed`);

if (results.missing.length) {
  console.log();
  console.log(
    chalk.red(
      "The following examples are missing and require manual reviewing & removal:"
    )
  );
  for (const m of results.missing) {
    console.log(m);
  }
}

if (results.failed.length) {
  console.log();
  console.log(chalk.red("The following examples failed to extract:"));
  for (const m of results.failed) {
    console.log(m.file);
    console.error(m.error);
    console.log();
  }
}