// SPDX-License-Identifier: Apache-2.0

#include <dcgm_agent.h>

#include <stdio.h>

static int list_supported_gpus(dcgmHandle_t handle)
{
    unsigned int gpuIds[DCGM_MAX_NUM_DEVICES] = { 0 };
    int count                                  = 0;
    dcgmReturn_t result = dcgmGetAllSupportedDevices(handle, gpuIds, &count);

    if (result != DCGM_ST_OK)
    {
        fprintf(stderr, "dcgmGetAllSupportedDevices failed: %s\n", errorString(result));
        return 1;
    }

    if (count == 0)
    {
        puts("DCGM reports no supported GPUs.");
        return 0;
    }

    printf("DCGM reports %d supported GPU%s:", count, count == 1 ? "" : "s");
    for (int i = 0; i < count; ++i)
    {
        printf(" %u", gpuIds[i]);
    }
    putchar('\n');
    return 0;
}

int main(void)
{
    dcgmHandle_t handle = 0;
    dcgmReturn_t result = dcgmInit();

    if (result != DCGM_ST_OK)
    {
        fprintf(stderr, "dcgmInit failed: %s\n", errorString(result));
        return 1;
    }

    result = dcgmStartEmbedded(DCGM_OPERATION_MODE_AUTO, &handle);
    if (result != DCGM_ST_OK)
    {
        fprintf(stderr, "dcgmStartEmbedded failed: %s\n", errorString(result));
        result = dcgmShutdown();
        if (result != DCGM_ST_OK)
        {
            fprintf(stderr, "dcgmShutdown failed: %s\n", errorString(result));
        }
        return 1;
    }

    int exitCode = list_supported_gpus(handle);

    result = dcgmStopEmbedded(handle);
    if (result != DCGM_ST_OK)
    {
        fprintf(stderr, "dcgmStopEmbedded failed: %s\n", errorString(result));
        exitCode = 1;
    }

    result = dcgmShutdown();
    if (result != DCGM_ST_OK)
    {
        fprintf(stderr, "dcgmShutdown failed: %s\n", errorString(result));
        exitCode = 1;
    }

    return exitCode;
}