#include <vector>
#include <string>
#include <stdint.h>

#include "spdlog/spdlog.h"

#include "acl/acl.h"


std::vector<int> get_npu_ids() {
  std::vector<int> npu_ids;
  uint32_t device_count = 0;
  aclError error_code = aclrtGetDeviceCount(&device_count);
  if (error_code != 0) {
    spdlog::error("Failed to get NPU device count, error code: {}", error_code);
    throw std::runtime_error("Failed to get NPU device count");
  }

  const char *env_available_npu = getenv("ASCEND_RT_VISIBLE_DEVICES");
  if (env_available_npu) {
    std::string npu_str(env_available_npu);
    size_t start = 0;
    while (start < npu_str.size()) {
      size_t next = npu_str.find(',', start);
      if (next == std::string::npos) {
        next = npu_str.size();
      }
      npu_ids.push_back(std::stoi(npu_str.substr(start, next - start)));
      start = next + 1;
      if (npu_ids.size() >= device_count) {
        break;
      }
    }
  } else {
    for (uint32_t i = 0; i < device_count; ++i) {
      npu_ids.push_back(static_cast<int>(i));
    }
  }
  return npu_ids;
}