/**
 * MVGAL - Multi-Vendor GPU Aggregation Layer for Linux
 * NVML Loader - Dynamic loading of NVIDIA Management Library
 * SPDX-License-Identifier: MIT
 */

#ifndef MVGAL_RUNTIME_NVML_LOADER_HPP
#define MVGAL_RUNTIME_NVML_LOADER_HPP

#include <string>
#include <vector>

namespace mvgal {

/**
 * NVML structures (subset)
 */
struct nvmlMemory_t {
    unsigned long long total;
    unsigned long long free;
    unsigned long long used;
};

struct nvmlUtilization_t {
    unsigned int gpu;
    unsigned int memory;
};

typedef enum nvmlReturn_enum {
    NVML_SUCCESS = 0,
    NVML_ERROR_UNINITIALIZED = 1,
    NVML_ERROR_INVALID_ARGUMENT = 2,
    NVML_ERROR_NOT_SUPPORTED = 3,
    NVML_ERROR_NO_PERMISSION = 4,
    NVML_ERROR_ALREADY_INITIALIZED = 5,
    NVML_ERROR_NOT_FOUND = 6,
    NVML_ERROR_INSUFFICIENT_SIZE = 7,
    NVML_ERROR_INSUFFICIENT_POWER = 8,
    NVML_ERROR_DRIVER_NOT_LOADED = 9,
    NVML_ERROR_TIMEOUT = 10,
    NVML_ERROR_IRQ_ISSUE = 11,
    NVML_ERROR_LIBRARY_NOT_FOUND = 12,
    NVML_ERROR_FUNCTION_NOT_FOUND = 13,
    NVML_ERROR_CORRUPTED_INFOROM = 14,
    NVML_ERROR_GPU_IS_LOST = 15,
    NVML_ERROR_RESET_REQUIRED = 16,
    NVML_ERROR_OPERATING_SYSTEM = 17,
    NVML_ERROR_LIB_RM_VERSION_MISMATCH = 18,
    NVML_ERROR_IN_USE = 19,
    NVML_ERROR_MEMORY = 20,
    NVML_ERROR_NO_DATA = 21,
    NVML_ERROR_VGPU_ECC_NOT_SUPPORTED = 22,
    NVML_ERROR_UNKNOWN = 999
} nvmlReturn_t;

typedef enum nvmlClockType_enum {
    NVML_CLOCK_GRAPHICS = 0,
    NVML_CLOCK_SM = 1,
    NVML_CLOCK_MEM = 2,
    NVML_CLOCK_VIDEO = 3
} nvmlClockType_t;

typedef struct nvmlDevice_st* nvmlDevice_t;

/**
 * NVML Loader class
 */
class NvmlLoader {
public:
    static NvmlLoader& instance();

    bool load();
    void unload();
    bool isLoaded() const { return m_handle != nullptr; }

    /* NVML API wrappers */
    nvmlReturn_t init();
    nvmlReturn_t shutdown();
    nvmlReturn_t deviceGetCount(unsigned int* deviceCount);
    nvmlReturn_t deviceGetHandleByIndex(unsigned int index, nvmlDevice_t* device);
    nvmlReturn_t deviceGetHandleByPciBusId(const char* pciBusId, nvmlDevice_t* device);
    nvmlReturn_t deviceGetMemoryInfo(nvmlDevice_t device, nvmlMemory_t* memory);
    nvmlReturn_t deviceGetUtilizationRates(nvmlDevice_t device, nvmlUtilization_t* utilization);
    nvmlReturn_t deviceGetTemperature(nvmlDevice_t device, unsigned int sensorType, unsigned int* temp);
    nvmlReturn_t deviceGetPowerUsage(nvmlDevice_t device, unsigned int* power);
    nvmlReturn_t deviceGetClockInfo(nvmlDevice_t device, nvmlClockType_t type, unsigned int* clock);
    nvmlReturn_t deviceSetPowerManagementLimit(nvmlDevice_t device, unsigned int limit);

private:
    NvmlLoader();
    ~NvmlLoader();

    void* m_handle;

    /* Function pointers */
    nvmlReturn_t (*m_nvmlInit)();
    nvmlReturn_t (*m_nvmlShutdown)();
    nvmlReturn_t (*m_nvmlDeviceGetCount)(unsigned int*);
    nvmlReturn_t (*m_nvmlDeviceGetHandleByIndex)(unsigned int, nvmlDevice_t*);
    nvmlReturn_t (*m_nvmlDeviceGetHandleByPciBusId)(const char*, nvmlDevice_t*);
    nvmlReturn_t (*m_nvmlDeviceGetMemoryInfo)(nvmlDevice_t, nvmlMemory_t*);
    nvmlReturn_t (*m_nvmlDeviceGetUtilizationRates)(nvmlDevice_t, nvmlUtilization_t*);
    nvmlReturn_t (*m_nvmlDeviceGetTemperature)(nvmlDevice_t, unsigned int, unsigned int*);
    nvmlReturn_t (*m_nvmlDeviceGetPowerUsage)(nvmlDevice_t, unsigned int*);
    nvmlReturn_t (*m_nvmlDeviceGetClockInfo)(nvmlDevice_t, nvmlClockType_t, unsigned int*);
    nvmlReturn_t (*m_nvmlDeviceSetPowerManagementLimit)(nvmlDevice_t, unsigned int);
};

} // namespace mvgal

#endif // MVGAL_RUNTIME_NVML_LOADER_HPP
