// What gpu-bench says about the machine: the computer, CPU and memory, and the NVIDIA GPU through
// NVML (the driver's management library, nvml.dll / libnvidia-ml.so.1), which also gives the GPU's
// memory, load, temperature, power and clock while the models run. Without NVML (no NVIDIA driver)
// the GPU lines say so and the samples are empty.

System_Info :: struct {
  machine: string;
  os: string;
  cpu: string;
  logical_cores: s64;
  ram_gb: float64;
  gpu: string;              // "" without NVML
  vram_gb: float64;
  driver: string;
  cuda_driver: string;      // the CUDA version the driver supports, such as 13.0
  power_limit_w: float64;   // the board's enforced limit: a laptop's GPU is often far below the desktop part's
  max_sm_clock_mhz: s64;
  max_mem_clock_mhz: s64;
  gpu_error: string;        // why there is no GPU information
  power_source: string;     // a laptop's: "AC" or "battery", which can change its GPU's speed several times over
}

// One sample of the GPU while a model runs.
Gpu_Sample :: struct {
  memory_used_mb: float64;   // the whole device's, every process's
  util_percent: float64;
  temperature_c: float64;
  power_w: float64;
  sm_clock_mhz: float64;
}

system_info :: () -> System_Info {
  info: System_Info;
  info.logical_cores = get_number_of_processors(.ALL_LOGICAL);
  #if OS == .WINDOWS {
    name: [256] u16;
    size: u32 = name.count;
    if GetComputerNameW(name.data, *size)  info.machine = wide_to_utf8(name.data, cast(s32) size);
    v: OSVERSIONINFOW;
    v.dwOSVersionInfoSize = size_of(OSVERSIONINFOW);
    RtlGetVersion(*v);
    info.os = sprint("Windows %.%.%", v.dwMajorVersion, v.dwMinorVersion, v.dwBuildNumber);
    buffer: [256] u8;
    bytes: u32 = buffer.count;
    if RegGetValueA(HKEY_LOCAL_MACHINE, "HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0", "ProcessorNameString", RRF_RT_REG_SZ, null, buffer.data, *bytes) == 0
      info.cpu = copy_string(trim(to_string(buffer.data)));
    memory: MEMORYSTATUSEX;
    memory.dwLength = size_of(MEMORYSTATUSEX);
    if GlobalMemoryStatusEx(*memory)  info.ram_gb = memory.ullTotalPhys / (1024.0 * 1024 * 1024);
    power: SYSTEM_POWER_STATUS;
    if GetSystemPowerStatus(*power) && power.BatteryFlag != 128 {   // 128: no battery
      if power.ACLineStatus == 1       info.power_source = "AC";
      else if power.ACLineStatus == 0  info.power_source = sprint("battery (% percent)", power.BatteryLifePercent);
    }
  } else #if OS == .LINUX {
    host, host_ok := read_entire_file("/etc/hostname", log_errors = false);
    if host_ok  info.machine = trim(host);
    release, release_ok := read_entire_file("/proc/sys/kernel/osrelease", log_errors = false);
    info.os = ifx release_ok then sprint("Linux %", trim(release)) else "Linux";
    cpuinfo, cpu_ok := read_entire_file("/proc/cpuinfo", log_errors = false);
    if cpu_ok  for split(cpuinfo, "\n",, temp)  if begins_with(it, "model name") { info.cpu = trim(advance(it, find_index_from_left(it, #char ":") + 1)); break; }
    meminfo, mem_ok := read_entire_file("/proc/meminfo", log_errors = false);
    if mem_ok  for split(meminfo, "\n",, temp)  if begins_with(it, "MemTotal:") { info.ram_gb = string_to_int(trim(trim(advance(it, 9)), " kB")) / (1024.0 * 1024); break; }
  } else {
    info.os = tprint("%", OS);
  }
  if !info.machine  info.machine = "unknown";

  if !nvml_start() { info.gpu_error = nvml.error; return info; }
  text: [96] u8;
  if nvml.DeviceGetName(nvml.device, text.data, text.count) == 0  info.gpu = copy_string(to_string(text.data));
  if nvml.SystemGetDriverVersion(text.data, text.count) == 0  info.driver = copy_string(to_string(text.data));
  cuda: s32;
  if nvml.SystemGetCudaDriverVersion(*cuda) == 0  info.cuda_driver = sprint("%.%", cuda / 1000, (cuda % 1000) / 10);
  vram: Nvml_Memory;
  if nvml.DeviceGetMemoryInfo(nvml.device, *vram) == 0  info.vram_gb = vram.total / (1024.0 * 1024 * 1024);
  limit: u32;
  if nvml.DeviceGetEnforcedPowerLimit(nvml.device, *limit) == 0  info.power_limit_w = limit / 1000.0;
  clock: u32;
  if nvml.DeviceGetMaxClockInfo(nvml.device, NVML_CLOCK_SM, *clock) == 0   info.max_sm_clock_mhz = clock;
  if nvml.DeviceGetMaxClockInfo(nvml.device, NVML_CLOCK_MEM, *clock) == 0  info.max_mem_clock_mhz = clock;
  return info;
}

// The GPU now; false without NVML.
gpu_sample :: () -> Gpu_Sample, ok: bool {
  s: Gpu_Sample;
  if !nvml.started  return s, false;
  memory: Nvml_Memory;
  if nvml.DeviceGetMemoryInfo(nvml.device, *memory) == 0  s.memory_used_mb = memory.used / (1024.0 * 1024);
  util: Nvml_Utilization;
  if nvml.DeviceGetUtilizationRates(nvml.device, *util) == 0  s.util_percent = cast(float64) util.gpu;
  value: u32;
  if nvml.DeviceGetTemperature(nvml.device, NVML_TEMPERATURE_GPU, *value) == 0  s.temperature_c = cast(float64) value;
  if nvml.DeviceGetPowerUsage(nvml.device, *value) == 0  s.power_w = value / 1000.0;
  if nvml.DeviceGetClockInfo(nvml.device, NVML_CLOCK_SM, *value) == 0  s.sm_clock_mhz = cast(float64) value;
  return s, true;
}

// The CUDA runtime's and cuDNN's versions, as loaded into this process by onnxruntime's CUDA provider,
// with where each came from; "" for one not loaded (before a GPU engine is made).
cuda_library_versions :: () -> cudart: string, cudnn: string {
  cudart, cudnn: string;
  #if OS == .WINDOWS {
    module := GetModuleHandleA("cudart64_13.dll");
    if module {
      get := cast(#type (version: *s32) -> s32 #c_call) GetProcAddress(module, "cudaRuntimeGetVersion");
      version: s32;
      if get && get(*version) == 0  cudart = sprint("%.% (%)", version / 1000, (version % 1000) / 10, module_path(module));
    }
    module = GetModuleHandleA("cudnn64_9.dll");
    if module {
      get := cast(#type () -> u64 #c_call) GetProcAddress(module, "cudnnGetVersion");
      if get {
        version := get();   // 9.x: major * 10000 + minor * 100 + patch
        cudnn = sprint("%.%.% (%)", version / 10000, (version % 10000) / 100, version % 100, module_path(module));
      }
    }
  }
  return cudart, cudnn;
}

#scope_file

#if OS == .WINDOWS {
  SYSTEM_POWER_STATUS :: struct {
    ACLineStatus, BatteryFlag, BatteryLifePercent, SystemStatusFlag: u8;
    BatteryLifeTime, BatteryFullLifeTime: u32;
  }
  GetSystemPowerStatus :: (status: *SYSTEM_POWER_STATUS) -> BOOL #foreign kernel32;
  kernel32 :: #system_library "kernel32";

  module_path :: (module: HANDLE) -> string {
    name: [1024] u16;
    count := GetModuleFileNameW(module, name.data, name.count);
    return wide_to_utf8(name.data, count,, temp);
  }
}

Nvml_Device :: *void;
Nvml_Memory :: struct { total, free, used: u64; }
Nvml_Utilization :: struct { gpu, memory: u32; }
NVML_CLOCK_SM  :: 1;
NVML_CLOCK_MEM :: 2;
NVML_TEMPERATURE_GPU :: 0;

nvml: struct {
  started: bool;
  error: string;
  device: Nvml_Device;

  Init: #type () -> s32 #c_call;
  DeviceGetHandleByIndex: #type (index: u32, device: *Nvml_Device) -> s32 #c_call;
  DeviceGetName: #type (device: Nvml_Device, name: *u8, length: u32) -> s32 #c_call;
  SystemGetDriverVersion: #type (version: *u8, length: u32) -> s32 #c_call;
  SystemGetCudaDriverVersion: #type (version: *s32) -> s32 #c_call;
  DeviceGetMemoryInfo: #type (device: Nvml_Device, memory: *Nvml_Memory) -> s32 #c_call;
  DeviceGetUtilizationRates: #type (device: Nvml_Device, util: *Nvml_Utilization) -> s32 #c_call;
  DeviceGetTemperature: #type (device: Nvml_Device, sensor: s32, value: *u32) -> s32 #c_call;
  DeviceGetPowerUsage: #type (device: Nvml_Device, milliwatts: *u32) -> s32 #c_call;
  DeviceGetEnforcedPowerLimit: #type (device: Nvml_Device, milliwatts: *u32) -> s32 #c_call;
  DeviceGetClockInfo: #type (device: Nvml_Device, type: s32, mhz: *u32) -> s32 #c_call;
  DeviceGetMaxClockInfo: #type (device: Nvml_Device, type: s32, mhz: *u32) -> s32 #c_call;
}

// NVML, and its first GPU: the one onnxruntime's CUDA provider uses (device 0) on a machine with one.
nvml_start :: () -> bool {
  if nvml.started  return true;
  #if OS == .WINDOWS {
    library := LoadLibraryA("nvml.dll");
    if !library  library = LoadLibraryA("C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvml.dll");
    if !library { nvml.error = "nvml.dll was not found: no NVIDIA driver."; return false; }
    find :: (name: *u8) -> *void #expand { return GetProcAddress(library, name); }
  } else #if OS == .LINUX {
    library := dlopen("libnvidia-ml.so.1", RTLD_NOW);
    if !library { nvml.error = "libnvidia-ml.so.1 was not found: no NVIDIA driver."; return false; }
    find :: (name: *u8) -> *void #expand { return dlsym(library, name); }
  } else {
    nvml.error = "No NVML on this OS.";
    return false;
  }
  nvml.Init                        = xx find("nvmlInit_v2");
  nvml.DeviceGetHandleByIndex      = xx find("nvmlDeviceGetHandleByIndex_v2");
  nvml.DeviceGetName               = xx find("nvmlDeviceGetName");
  nvml.SystemGetDriverVersion      = xx find("nvmlSystemGetDriverVersion");
  nvml.SystemGetCudaDriverVersion  = xx find("nvmlSystemGetCudaDriverVersion");
  nvml.DeviceGetMemoryInfo         = xx find("nvmlDeviceGetMemoryInfo");
  nvml.DeviceGetUtilizationRates   = xx find("nvmlDeviceGetUtilizationRates");
  nvml.DeviceGetTemperature        = xx find("nvmlDeviceGetTemperature");
  nvml.DeviceGetPowerUsage         = xx find("nvmlDeviceGetPowerUsage");
  nvml.DeviceGetEnforcedPowerLimit = xx find("nvmlDeviceGetEnforcedPowerLimit");
  nvml.DeviceGetClockInfo          = xx find("nvmlDeviceGetClockInfo");
  nvml.DeviceGetMaxClockInfo       = xx find("nvmlDeviceGetMaxClockInfo");
  if !nvml.Init || !nvml.DeviceGetHandleByIndex || !nvml.DeviceGetName || !nvml.SystemGetDriverVersion || !nvml.SystemGetCudaDriverVersion
     || !nvml.DeviceGetMemoryInfo || !nvml.DeviceGetUtilizationRates || !nvml.DeviceGetTemperature || !nvml.DeviceGetPowerUsage
     || !nvml.DeviceGetEnforcedPowerLimit || !nvml.DeviceGetClockInfo || !nvml.DeviceGetMaxClockInfo {
    nvml.error = "NVML lacks a function gpu-bench needs: an old driver?";
    return false;
  }
  status := nvml.Init();
  if status != 0 { nvml.error = tprint("NVML did not start (error %).", status); return false; }
  status = nvml.DeviceGetHandleByIndex(0, *nvml.device);
  if status != 0 { nvml.error = tprint("NVML found no GPU (error %).", status); return false; }
  nvml.started = true;
  return true;
}

#if OS == .WINDOWS {
  #import "Windows";
  #import "Windows_Utf8";
  #import "Windows_Registry";
} else #if OS == .LINUX {
  #import "POSIX";
}