// Downloads model files into models/ and verifies their sha256.
// Files already present with a matching hash are skipped.
import { join } from "node:path";

interface ModelFile {
  path: string; // relative to repo root
  url: string;
  sha256: string;
}

export const MODELS: ModelFile[] = [
  {
    path: "models/demucs/htdemucs_6s.onnx",
    url: "https://huggingface.co/Kani95/htdemucs-6s-ort/resolve/main/htdemucs_6s.onnx",
    sha256: "35cb61d4d35173ff1d16df7393f41f05651ffdfe13d3c10e06abc26708bd6933",
  },
  {
    path: "models/basic-pitch/nmp.onnx",
    url: "https://github.com/spotify/basic-pitch/raw/main/basic_pitch/saved_models/icassp_2022/nmp.onnx",
    sha256: "2c3c1d144bfa61ad236e92e169c13535c880469a12a047d4e73451f2c059a0ec",
  },
];

const ROOT = join(import.meta.dir, "..");

async function sha256Of(file: string): Promise<string> {
  const hasher = new Bun.CryptoHasher("sha256");
  hasher.update(new Uint8Array(await Bun.file(file).arrayBuffer()));
  return hasher.digest("hex");
}

export async function fetchModels(root = ROOT): Promise<void> {
  for (const m of MODELS) {
    const dest = join(root, m.path);
    if (await Bun.file(dest).exists()) {
      const have = await sha256Of(dest);
      if (have === m.sha256) {
        console.log(`ok       ${m.path}`);
        continue;
      }
      console.log(`mismatch ${m.path} (have ${have}), re-downloading`);
    }
    console.log(`fetch    ${m.url}`);
    const res = await fetch(m.url, { redirect: "follow" });
    if (!res.ok) throw new Error(`download failed for ${m.url}: HTTP ${res.status}`);
    const bytes = new Uint8Array(await res.arrayBuffer());
    const hasher = new Bun.CryptoHasher("sha256");
    hasher.update(bytes);
    const got = hasher.digest("hex");
    if (got !== m.sha256) {
      throw new Error(`sha256 mismatch for ${m.path}: expected ${m.sha256}, got ${got}`);
    }
    await Bun.write(dest, bytes);
    console.log(`saved    ${m.path}`);
  }
}

if (import.meta.main) {
  await fetchModels();
}