/*
 * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
 *
 * SPDX-License-Identifier: MIT
 */

#include <hip/hip_runtime.h>
#include <iostream>
#include <stdlib.h>

bool UNSETENV(std::string var) {
  int result = -1;
#ifdef __unix__
  result = unsetenv(var.c_str());
#else
  result = _putenv((var + '=').c_str());
#endif
  return (result == 0) ? true : false;
}

bool SETENV(std::string var, std::string value) {
  int result = -1;
#ifdef __unix__
  result = setenv(var.c_str(), value.c_str(), 1);
#else
  result = _putenv((var + '=' + value).c_str());
#endif
  return (result == 0) ? true : false;
}

// Expects 1 command line arg, which is the Device Visible String
int main(int argc, char** argv) {
  if (argc != 2) {
    std::cerr << "Invalid number of args passed.\n"
              << "argc : " << argc << std::endl;
    for (int i = 0; i < argc; i++) {
      std::cerr << "  argv[" << i << "] : " << argv[0] << std::endl;
    }
    std::cerr << "The program expects device visibility string i.e. 0,1,2" << std::endl;
    return -1;
  }

  // disable visible_devices env from shell
#ifdef __HIP_PLATFORM_NVIDIA__
  UNSETENV("CUDA_VISIBLE_DEVICES");
  SETENV("CUDA_VISIBLE_DEVICES", argv[1]);
  auto init_res = hipInit(0);
  if (hipSuccess != init_res) {
    std::cerr << "CUDA INIT API returned : " << hipGetErrorString(init_res) << std::endl;
    return -1;
  }
#else
  UNSETENV("ROCR_VISIBLE_DEVICES");
  UNSETENV("HIP_VISIBLE_DEVICES");
  SETENV("ROCR_VISIBLE_DEVICES", argv[1]);
  SETENV("HIP_VISIBLE_DEVICES", argv[1]);
#endif

  int count = 0;
  auto res = hipGetDeviceCount(&count);
  if (hipSuccess != res) {
    std::cerr << "HIP API returned : " << hipGetErrorString(res) << std::endl;
    return -1;
  }

#ifdef __HIP_PLATFORM_NVIDIA__
  UNSETENV("CUDA_VISIBLE_DEVICES");
#else
  UNSETENV("ROCR_VISIBLE_DEVICES");
  UNSETENV("HIP_VISIBLE_DEVICES");
#endif
  return count;
}