/**
 * MVGAL - Multi-Vendor GPU Aggregation Layer for Linux
 * NVML Loader Implementation
 * SPDX-License-Identifier: MIT
 */

#include "nvml_loader.hpp"
#include <dlfcn.h>
#include <iostream>

namespace mvgal {

NvmlLoader::NvmlLoader()
    : m_handle(nullptr),
      m_nvmlInit(nullptr),
      m_nvmlShutdown(nullptr),
      m_nvmlDeviceGetCount(nullptr),
      m_nvmlDeviceGetHandleByIndex(nullptr),
      m_nvmlDeviceGetHandleByPciBusId(nullptr),
      m_nvmlDeviceGetMemoryInfo(nullptr),
      m_nvmlDeviceGetUtilizationRates(nullptr),
      m_nvmlDeviceGetTemperature(nullptr),
      m_nvmlDeviceGetPowerUsage(nullptr),
      m_nvmlDeviceGetClockInfo(nullptr),
      m_nvmlDeviceSetPowerManagementLimit(nullptr)
{
}

NvmlLoader::~NvmlLoader()
{
    unload();
}

NvmlLoader& NvmlLoader::instance()
{
    static NvmlLoader loader;
    return loader;
}

bool NvmlLoader::load()
{
    if (m_handle) return true;

    /* Try common NVML library names */
    const char* libNames[] = {
        "libnvidia-ml.so.1",
        "libnvidia-ml.so"
    };

    for (const char* name : libNames) {
        m_handle = dlopen(name, RTLD_LAZY);
        if (m_handle) break;
    }

    if (!m_handle) return false;

    /* Load functions */
#define LOAD_FUNC(name) \
    m_##name = reinterpret_cast<decltype(m_##name)>(dlsym(m_handle, #name)); \
    if (!m_##name) { unload(); return false; }

    LOAD_FUNC(nvmlInit);
    LOAD_FUNC(nvmlShutdown);
    LOAD_FUNC(nvmlDeviceGetCount);
    LOAD_FUNC(nvmlDeviceGetHandleByIndex);
    LOAD_FUNC(nvmlDeviceGetHandleByPciBusId);
    LOAD_FUNC(nvmlDeviceGetMemoryInfo);
    LOAD_FUNC(nvmlDeviceGetUtilizationRates);
    LOAD_FUNC(nvmlDeviceGetTemperature);
    LOAD_FUNC(nvmlDeviceGetPowerUsage);
    LOAD_FUNC(nvmlDeviceGetClockInfo);
    LOAD_FUNC(nvmlDeviceSetPowerManagementLimit);

#undef LOAD_FUNC

    return true;
}

void NvmlLoader::unload()
{
    if (m_handle) {
        if (m_nvmlShutdown) m_nvmlShutdown();
        dlclose(m_handle);
        m_handle = nullptr;
    }

    m_nvmlInit = nullptr;
    m_nvmlShutdown = nullptr;
    m_nvmlDeviceGetCount = nullptr;
    m_nvmlDeviceGetHandleByIndex = nullptr;
    m_nvmlDeviceGetHandleByPciBusId = nullptr;
    m_nvmlDeviceGetMemoryInfo = nullptr;
    m_nvmlDeviceGetUtilizationRates = nullptr;
    m_nvmlDeviceGetTemperature = nullptr;
    m_nvmlDeviceGetPowerUsage = nullptr;
    m_nvmlDeviceGetClockInfo = nullptr;
    m_nvmlDeviceSetPowerManagementLimit = nullptr;
}

nvmlReturn_t NvmlLoader::init()
{
    if (!m_nvmlInit) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlInit();
}

nvmlReturn_t NvmlLoader::shutdown()
{
    if (!m_nvmlShutdown) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlShutdown();
}

nvmlReturn_t NvmlLoader::deviceGetCount(unsigned int* deviceCount)
{
    if (!m_nvmlDeviceGetCount) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetCount(deviceCount);
}

nvmlReturn_t NvmlLoader::deviceGetHandleByIndex(unsigned int index, nvmlDevice_t* device)
{
    if (!m_nvmlDeviceGetHandleByIndex) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetHandleByIndex(index, device);
}

nvmlReturn_t NvmlLoader::deviceGetHandleByPciBusId(const char* pciBusId, nvmlDevice_t* device)
{
    if (!m_nvmlDeviceGetHandleByPciBusId) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetHandleByPciBusId(pciBusId, device);
}

nvmlReturn_t NvmlLoader::deviceGetMemoryInfo(nvmlDevice_t device, nvmlMemory_t* memory)
{
    if (!m_nvmlDeviceGetMemoryInfo) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetMemoryInfo(device, memory);
}

nvmlReturn_t NvmlLoader::deviceGetUtilizationRates(nvmlDevice_t device, nvmlUtilization_t* utilization)
{
    if (!m_nvmlDeviceGetUtilizationRates) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetUtilizationRates(device, utilization);
}

nvmlReturn_t NvmlLoader::deviceGetTemperature(nvmlDevice_t device, unsigned int sensorType, unsigned int* temp)
{
    if (!m_nvmlDeviceGetTemperature) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetTemperature(device, sensorType, temp);
}

nvmlReturn_t NvmlLoader::deviceGetPowerUsage(nvmlDevice_t device, unsigned int* power)
{
    if (!m_nvmlDeviceGetPowerUsage) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetPowerUsage(device, power);
}

nvmlReturn_t NvmlLoader::deviceGetClockInfo(nvmlDevice_t device, nvmlClockType_t type, unsigned int* clock)
{
    if (!m_nvmlDeviceGetClockInfo) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceGetClockInfo(device, type, clock);
}

nvmlReturn_t NvmlLoader::deviceSetPowerManagementLimit(nvmlDevice_t device, unsigned int limit)
{
    if (!m_nvmlDeviceSetPowerManagementLimit) return NVML_ERROR_LIBRARY_NOT_FOUND;
    return m_nvmlDeviceSetPowerManagementLimit(device, limit);
}

} // namespace mvgal
