#!/usr/bin/env node
/**
 * Pre-flight for the GPU box. Run `npm run check` before the first render to
 * see whether ComfyUI is up, new enough, and holding the right weights.
 */
import { readFileSync, existsSync } from "node:fs";

for (const file of [".env.local", ".env"]) {
  if (!existsSync(file)) continue;
  for (const line of readFileSync(file, "utf8").split("\n")) {
    const match = line.match(/^\s*([A-Z0-9_]+)\s*=\s*(.*)\s*$/);
    if (match && !process.env[match[1]]) process.env[match[1]] = match[2].replace(/^["']|["']$/g, "");
  }
}

const url = (process.env.COMFY_URL ?? "http://127.0.0.1:8188").replace(/\/$/, "");
const token = process.env.COMFY_TOKEN;
const headers = token ? { Authorization: `Bearer ${token}` } : {};
const wantCheckpoint = process.env.YUE2_CHECKPOINT ?? "yue2_3b_int8_convrot.safetensors";
const wantEncoder = process.env.SHEETSAGE_ENCODER ?? "sheetsage2_bf16.safetensors";

const REQUIRED = [
  "YuE2GenerateABC", "YuE2GenerateMusic", "EmptyYuE2LatentAudio",
  "SheetSage2AudioToABC", "AudioEncoderLoader", "SaveAudioAdvanced", "PreviewAny",
];

const ok = (m) => console.log("  ok    " + m);
const warn = (m) => console.log("  note  " + m);
const bad = (m) => console.log("  fail  " + m);

console.log("\nChecking " + url + (token ? " with a token" : "") + "\n");

let info;
try {
  const res = await fetch(url + "/object_info", { headers });
  if (!res.ok) {
    throw new Error(
      res.status === 401 || res.status === 403
        ? "answered " + res.status + ", so COMFY_TOKEN is missing or wrong"
        : "answered " + res.status,
    );
  }
  info = await res.json();
  ok("ComfyUI is reachable");
} catch (error) {
  bad("cannot reach ComfyUI (" + error.message + ")");
  console.log("\n  Start it, or point COMFY_URL at the right host. To work on the");
  console.log("  interface without a GPU: COMFY_MOCK=1 npm run dev\n");
  process.exit(1);
}

const missing = REQUIRED.filter((n) => !(n in info));
if (missing.length) bad("missing nodes: " + missing.join(", ") + "  (needs ComfyUI v0.35.0 or newer)");
else ok("all YuE2 nodes present");

const combo = (node, input) => {
  const spec = info[node]?.input?.required?.[input];
  return Array.isArray(spec) && Array.isArray(spec[0]) ? spec[0] : [];
};

const checkpoints = combo("CheckpointLoaderSimple", "ckpt_name");
const yue = checkpoints.filter((f) => /yue2/i.test(f));
if (!checkpoints.length) bad("no checkpoints at all in ComfyUI/models/checkpoints/");
else if (!yue.length) bad("no YuE2 checkpoint found. Put yue2_3b_int8_convrot.safetensors in ComfyUI/models/checkpoints/");
else if (!yue.includes(wantCheckpoint)) warn("YUE2_CHECKPOINT is " + wantCheckpoint + ", but ComfyUI has: " + yue.join(", "));
else ok("checkpoint " + wantCheckpoint);

const encoders = combo("AudioEncoderLoader", "audio_encoder_name");
if (!encoders.length) warn("no audio encoder found. Reference covers need sheetsage2_bf16.safetensors in ComfyUI/models/audio_encoders/");
else if (!encoders.includes(wantEncoder)) warn("SHEETSAGE_ENCODER is " + wantEncoder + ", but ComfyUI has: " + encoders.join(", "));
else ok("transcriber " + wantEncoder);

console.log("");