// Downloads the pinned onnxruntime release for this OS, checks its SHA-256, and puts the libraries
// in <os>/ and the C API header in include/. The binaries are not in git; run this after cloning.
//
//   jai fetch.jai            the GPU build (CUDA 13) on Windows and Linux; the CPU build on macOS
//   jai fetch.jai - cpu      the CPU-only build: smaller, and enough when there is no NVIDIA GPU
//
// Uses curl and tar, which Windows 10 and later ship with, and certutil / sha256sum / shasum.

ORT_VERSION :: "1.30.0";

Release :: struct { gpu: bool; archive: string; sha256: string; }

#if OS == .WINDOWS {
  OS_FOLDER :: "windows";
  RELEASES :: Release.[
    .{true,  "onnxruntime-win-x64-gpu_cuda13-1.30.0.zip", "8fa4b08359af682cd605892cb59077049700b640128bb93fd2c7776cf9f55bdc"},
    .{false, "onnxruntime-win-x64-1.30.0.zip",            "c6ba983baf5681af108599675d2a89c2d145512d02de28aed0bff177cd0ba949"},
  ];
  LIBRARY_SUFFIX :: ".dll";
} else #if OS == .LINUX {
  OS_FOLDER :: "linux";
  RELEASES :: Release.[
    .{true,  "onnxruntime-linux-x64-gpu_cuda13-1.30.0.tgz", "382d79133112388cf94ce5855789b7c9bef12bef76a08b6b277e5a317213adcd"},
    .{false, "onnxruntime-linux-x64-1.30.0.tgz",            "a5ed5a3cac51fbb2e90da632ae43d19212faaa20e76484e62bcb7c23ddb3b3fd"},
  ];
  LIBRARY_SUFFIX :: ".so";
} else #if OS == .MACOS {
  OS_FOLDER :: "macos";
  RELEASES :: Release.[
    .{false, "onnxruntime-osx-arm64-1.30.0.tgz", "6ebb5062a934537c352937821f9fe9718e7de1a2db1122a93dd363ffd53a7012"},
  ];
  LIBRARY_SUFFIX :: ".dylib";
}

#run,stallable { if !fetch()  exit(1); }

fetch :: () -> bool {
  set_build_options_dc(.{do_output = false});
  args := get_build_options().compile_time_command_line;
  want_gpu := !array_find(args, "cpu");

  release: Release;
  for RELEASES  if it.gpu == want_gpu || RELEASES.count == 1 { release = it; break; }
  if !release.archive { log_error("No % build of onnxruntime for this OS.", ifx want_gpu then "GPU" else "CPU"); return false; }

  here := trim_right(path_strip_filename(#file), "/\\");
  set_working_directory(here);
  download_folder := tprint("%/download", here);   // absolute paths throughout
  make_directory_if_it_does_not_exist(download_folder);
  archive_path := tprint("%/%", download_folder, release.archive);

  if file_exists(archive_path) && sha256_of(archive_path) == release.sha256 {
    log("% is already downloaded.", release.archive);
  } else {
    url := tprint("https://github.com/microsoft/onnxruntime/releases/download/v%/%", ORT_VERSION, release.archive);
    log("Downloading %...", url);
    result := run_command("curl", "-L", "--fail", "--progress-bar", "-o", archive_path, url);
    if result.exit_code != 0 { log_error("curl failed (exit code %).", result.exit_code); return false; }
    hash := sha256_of(archive_path);
    if hash != release.sha256 {
      log_error("% has SHA-256 %, expected %; it was not used.", archive_path, hash, release.sha256);
      file_delete(archive_path);
      return false;
    }
  }

  extract_folder := tprint("%/extract", download_folder);
  remove_folder(extract_folder);
  make_directory_if_it_does_not_exist(extract_folder);
  // Windows' own tar reads .zip; another tar earlier on PATH (Git's) does not.
  tar := ifx OS == .WINDOWS then "C:/Windows/System32/tar.exe" else "tar";
  result := run_command(tar, "-xf", archive_path, "-C", extract_folder);
  if result.exit_code != 0 { log_error("tar failed (exit code %).", result.exit_code); return false; }
  defer remove_folder(extract_folder);

  root := tprint("%/%", extract_folder, path_strip_extension(path_strip_extension(release.archive)));   // .tgz, or .zip
  if !file_exists(tprint("%/lib", root))  root = tprint("%/%", extract_folder, path_strip_extension(release.archive));

  // Replace whatever build was here before, so a GPU and a CPU build never mix.
  library_folder := tprint("%/%", here, OS_FOLDER);
  remove_folder(library_folder);
  make_directory_if_it_does_not_exist(library_folder);
  copied := 0;
  for file_list(tprint("%/lib", root)) {
    name := path_filename(it);
    if !contains(name, LIBRARY_SUFFIX)  continue;   // .so.1.30.0 and the like included
    if !copy_file(it, tprint("%/%", library_folder, name)) { log_error("Could not copy %.", name); return false; }
    copied += 1;
  }
  make_directory_if_it_does_not_exist(tprint("%/include", here));
  for file_list(tprint("%/include", root)) {
    if !copy_file(it, tprint("%/include/%", here, path_filename(it))) { log_error("Could not copy %.", it); return false; }
  }
  write_entire_file(tprint("%/VERSION.txt", library_folder), tprint("%\n%\n", ORT_VERSION, release.archive));
  log("Put % libraries in %/% and the headers in %/include (onnxruntime %, %).", copied, here, OS_FOLDER, here, ORT_VERSION, ifx release.gpu then "GPU" else "CPU");
  return true;
}

// With the OS's own tools: File's delete_directory hands Windows a relative path it cannot resolve,
// and its deletion was seen racing the copies that followed.
remove_folder :: (path: string) {
  if !file_exists(path)  return;
  #if OS == .WINDOWS  run_command("cmd", "/c", "rmdir", "/s", "/q", replace(path, "/", "\\"));
  else                run_command("rm", "-rf", path);
}

// Lowercase hex, or "" when it could not be worked out.
sha256_of :: (path: string) -> string {
  #if OS == .WINDOWS {
    result, output := run_command("certutil", "-hashfile", path, "SHA256", capture_and_return_output = true);
    if result.exit_code != 0  return "";
    lines := split(output, "\n");
    if lines.count < 2  return "";
    return to_lower_copy(replace(trim(lines[1]), " ", ""));   // line 2 is the hash
  } else {
    #if OS == .MACOS  command := string.["shasum", "-a", "256", path];
    else              command := string.["sha256sum", path];
    result, output := run_command(..command, capture_and_return_output = true);
    if result.exit_code != 0  return "";
    return to_lower_copy(trim(split(output, " ")[0]));
  }
}

#import "Basic";
#import "Compiler";
#import "File";
#import "File_Utilities";
#import "Process";
#import "String";