#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <dirent.h>
#include <pthread.h>
#include <sys/stat.h>

#include <fcntl.h>
#include <string.h>

#define MAX_PATH_SIZE  1024
#define MAX_PF_NUM     20
#define VF_PATH   64
#define VF_LEN    10
#define DEVICE_DIR     "/sys/bus/pci/devices"
#define STREQ(a, b) (strcmp(a, b) == 0)

typedef struct SriovEnableThread {
    pthread_t thread;
    char bdf[VF_PATH];
    char totalVfs[VF_LEN];
} SriovEnableThread;

static int FileWriteString(const char *path, const char *str)
{
    int fd = open(path, O_WRONLY | O_TRUNC);
    if (fd == -1) {
        return -1;
    }

    if (write(fd, str, strlen(str)) < 0) {
        close(fd);
        return -1;
    }

    if (close(fd) != 0) {
        return -1;
    }

    return 0;
}

static int FileReadString(const char *file, char *buf, int len)
{
    int fd;
    ssize_t sz;

    fd = open(file, O_RDONLY);
    if (fd < 0) {
        return -1;
    }

    sz = read(fd, buf, len);
    close(fd);
    if (sz < 0) {
        return -1;
    }

    if (sz == len) {
        printf("String len '%d' is too short to save %s\n", len, file);
        buf[len - 1] = '\0';
        return -1;
    }

    buf[sz] = '\0';
    return sz;
}

static void *DeviceSriovEnable(void *arg)
{
    char numVfsPath[VF_PATH] = { 0 };
    SriovEnableThread *sriovThread = (SriovEnableThread *)arg;
    const char *bdf = sriovThread->bdf;
    const char *totalVfs = sriovThread->totalVfs;

    printf("echo %s > /sys/bus/pci/devices/%s/sriov_numvfs\n", totalVfs, bdf);

    if (snprintf(numVfsPath, sizeof(numVfsPath) - 1,
                   "%s/%s/sriov_numvfs", DEVICE_DIR, bdf) < 0) {
        printf("Get device sriov_numvfs fail\n");
        pthread_exit(NULL);
    }

    if (FileWriteString(numVfsPath, totalVfs)) {
        printf("enable '%s' sriov fail\n", bdf);
        pthread_exit(NULL);
    }
    pthread_exit(NULL);
}

static int DeviceCheckSriovEnable(const char *bdf)
{
    char path[VF_PATH] = { 0 };
    int ret;

    ret = snprintf(path, sizeof(path) - 1,
                     "%s/%s/sriov_totalvfs", DEVICE_DIR, bdf);
    if (ret < 0) {
        return ret;
    }

    if (access(path, R_OK) != 0) {
        return 0;
    }

    return 1;
}

static int DeviceReadValue(const char *bdf, const char *pattern, char *buffer, size_t len)
{
    char path[VF_PATH] = {0};
    int ret;

    ret = snprintf(path, sizeof(path) - 1, "%s/%s/%s", DEVICE_DIR,
                     bdf, pattern);
    if (ret < 0) {
        printf("Set devicePath fail\n");
        return ret;
    }

    if (FileReadString(path, buffer, len) < 0) {
        printf("Read data from devicePath %s fail\n", path);
        return -1;
    }

    return 0;
}

void main(int argc, char **argv)
{
    int ret;
    DIR* dirp;
    int i = 0;
    struct dirent *dp = NULL;
    char buffer[VF_LEN] = { 0 };
    int sriovThreadLen = sizeof(SriovEnableThread) * MAX_PF_NUM;

    dirp = opendir(DEVICE_DIR);
    if (!dirp) {
        printf("Can not find %s\n", DEVICE_DIR);
        return;
    }

    SriovEnableThread *sriovThread = (SriovEnableThread *)malloc(sriovThreadLen);
    if (!sriovThread) {
        printf("Can not alloc sriovThread\n");
        closedir(dirp);
        return;
    }
    memset(sriovThread, 0, sriovThreadLen);

    while ((dp = readdir(dirp)) != NULL) {
        if (STREQ(dp->d_name, ".") || STREQ(dp->d_name, "..")) {
            continue;
        }

        if (DeviceCheckSriovEnable(dp->d_name) <= 0) {
            continue;
        }

        if (DeviceReadValue(dp->d_name, "sriov_totalvfs", buffer, sizeof(buffer)) < 0) {
            continue;
        }

        if (i >= MAX_PF_NUM) {
            printf("pf num is exceed max %d\n", MAX_PF_NUM);
            break;
        }

        strcpy(sriovThread[i].bdf, dp->d_name);
        strcpy(sriovThread[i].totalVfs, buffer);
        ret = pthread_create(&sriovThread[i].thread, NULL, DeviceSriovEnable,
                             (void *)&sriovThread[i]);
        if (ret) {
            printf("create sriov thread %d for %s failed, ret : %d\n", i, sriovThread[i].bdf, ret);
            break;
        }
        i++;
    }

    closedir(dirp);
    /* wait until all sriov enable thread finish */
    for (int j = 0; j < i; j++) {
        pthread_join(sriovThread[j].thread, NULL);
    }

    printf("sriov enable finish\n");
    free(sriovThread);
}