// Whether onnxruntime's CUDA provider can run here, and if not, what is missing and how to fix it.
// The provider needs the NVIDIA driver, parts of the CUDA 13 toolkit and cuDNN 9 for CUDA 13; it
// fails late and obscurely without them, so each library is looked for first, by the name the
// provider loads it by (from onnxruntime_providers_cuda.dll's imports and the names it loads).

Cuda_Library :: struct {
  name: string;
  part: Cuda_Part;
}

Cuda_Part :: enum { DRIVER; TOOLKIT; CUDNN; }

#if OS == .WINDOWS {
  CUDA_LIBRARIES :: Cuda_Library.[
    .{"nvcuda.dll",                             .DRIVER},
    .{"cudart64_13.dll",                        .TOOLKIT},
    .{"cublas64_13.dll",                        .TOOLKIT},
    .{"cublasLt64_13.dll",                      .TOOLKIT},
    .{"cufft64_12.dll",                         .TOOLKIT},
    .{"cudnn64_9.dll",                          .CUDNN},
    .{"cudnn_graph64_9.dll",                    .CUDNN},
    .{"cudnn_ops64_9.dll",                      .CUDNN},
    .{"cudnn_cnn64_9.dll",                      .CUDNN},
    .{"cudnn_heuristic64_9.dll",                .CUDNN},
    .{"cudnn_engines_precompiled64_9.dll",      .CUDNN},
    .{"cudnn_engines_runtime_compiled64_9.dll", .CUDNN},
  ];
  CUDA_PROVIDER_LIBRARY :: "onnxruntime_providers_cuda.dll";
} else #if OS == .LINUX {
  CUDA_LIBRARIES :: Cuda_Library.[
    .{"libcuda.so.1",        .DRIVER},
    .{"libcudart.so.13",     .TOOLKIT},
    .{"libcublas.so.13",     .TOOLKIT},
    .{"libcublasLt.so.13",   .TOOLKIT},
    .{"libcufft.so.12",      .TOOLKIT},
    .{"libcudnn.so.9",       .CUDNN},
  ];
  CUDA_PROVIDER_LIBRARY :: "libonnxruntime_providers_cuda.so";
} else {
  CUDA_LIBRARIES :: Cuda_Library.[];
  CUDA_PROVIDER_LIBRARY :: "";
}

// Whether the CUDA provider should be tried, and when not, an explanation for people: what is
// missing, and the steps that fix it. Temp.
ocr_cuda_check :: () -> available: bool, explanation: string {
  #if OS == .MACOS {
    return false, "OCR runs on the CPU: onnxruntime has no CUDA on macOS.";
  }

  provider_path := tprint("%/%", path_strip_filename(ort.library_path), CUDA_PROVIDER_LIBRARY);
  if !file_exists(provider_path) {
    return false, tprint("OCR runs on the CPU: the onnxruntime beside the program is the CPU-only build (% is not there).\nTo use an NVIDIA GPU, run `jai fetch.jai` in modules/onnxruntime (without `cpu`) and rebuild.", CUDA_PROVIDER_LIBRARY);
  }

  missing_driver, missing_toolkit, missing_cudnn: [..] string;
  missing_driver.allocator, missing_toolkit.allocator, missing_cudnn.allocator = temp, temp, temp;
  for CUDA_LIBRARIES {
    if ocr_library_loads(it.name)  continue;
    if it.part == {
      case .DRIVER;   array_add(*missing_driver,  it.name);
      case .TOOLKIT;  array_add(*missing_toolkit, it.name);
      case .CUDNN;    array_add(*missing_cudnn,   it.name);
    }
  }
  if !missing_driver.count && !missing_toolkit.count && !missing_cudnn.count  return true, "";

  builder: String_Builder;
  builder.allocator = temp;
  if missing_driver.count {
    append(*builder, "OCR runs on the CPU: no NVIDIA driver was found (");
    append(*builder, join(..missing_driver, ", ",, temp));
    append(*builder, "), so there is no NVIDIA GPU to use.");
    return false, builder_to_string(*builder,, temp);
  }

  append(*builder, "OCR runs on the CPU. The NVIDIA GPU can run it once these are installed:\n");
  step := 1;
  if missing_toolkit.count {
    print_to_builder(*builder, "\n%. The CUDA Toolkit 13.x, which has % that are missing.\n", step, join(..missing_toolkit, ", ",, temp));
    append(*builder, "   https://developer.nvidia.com/cuda-downloads\n");
    #if OS == .WINDOWS  append(*builder, "   Its installer puts its bin folder on PATH.\n");
    else                append(*builder, "   Then make sure its lib64 folder is on LD_LIBRARY_PATH.\n");
    step += 1;
  }
  if missing_cudnn.count {
    print_to_builder(*builder, "\n%. cuDNN 9.x for CUDA 13, which has % that are missing.\n", step, join(..missing_cudnn, ", ",, temp));
    append(*builder, "   https://developer.nvidia.com/cudnn-downloads\n");
    #if OS == .WINDOWS {
      append(*builder, "   Its installer does not change PATH: add the folder with these DLLs to it yourself,\n");
      append(*builder, "   such as C:\\Program Files\\NVIDIA\\CUDNN\\v9.x\\bin\\13.x\\x64.\n");
    } else {
      append(*builder, "   Then make sure its lib folder is on LD_LIBRARY_PATH.\n");
    }
    step += 1;
  }
  append(*builder, "\nCopying the DLLs beside the program also works. Restart the program afterwards; ");
  append(*builder, "PATH is read when it starts.");
  return false, builder_to_string(*builder,, temp);
}

#scope_module

ocr_library_loads :: (name: string) -> bool {
  #if OS == .WINDOWS {
    // The search order the CUDA provider's own loads use: beside the program, the system folders, PATH.
    module := LoadLibraryW(utf8_to_wide(name,, temp));
    if !module  return false;
    FreeLibrary(module);
    return true;
  } else {
    handle := dlopen(temp_c_string(name), RTLD_LAZY);
    if !handle  return false;
    dlclose(handle);
    return true;
  }
}