/* SPDX-License-Identifier: GPL-2.0-or-later */
#define _GNU_SOURCE 1
#include "../gpu.h"
#include <assert.h>
#include <dirent.h>

static int descriptors(void)
{
    DIR *dir = opendir("/proc/self/fd");
    struct dirent *entry;
    int count = 0;
    assert(dir);
    while ((entry = readdir(dir)))
        if (entry->d_name[0] != '.') count++;
    closedir(dir);
    return count;
}

static void unknown(gpu_sample s)
{
    assert(s.busy == -1 && s.used == -1 && s.total == -1 && s.temp == LLONG_MIN);
}

static gpu_sample test_command(const char *path, char *const args[])
{
    char text[256];
    return gpu_command(path, args, text, sizeof(text)) ? gpu_csv(text) : gpu_unknown;
}

static void *cache_writer(void *arg)
{
    size_t index = *(size_t *)arg;
    int i;
    for (i = 0; i < 20000; i++)
        gpu_publish(index, (gpu_sample){i % 101, i, i + 1, (long long)index});
    return NULL;
}

static void cache_concurrency(void)
{
    pthread_t writers[2];
    size_t indices[2] = {0, 1}, i;
    int repeat;
    gpu_count = 2;
    for (i = 0; i < 2; i++) {
        gpu_publish(i, (gpu_sample){0, 0, 1, (long long)i});
        assert(!pthread_create(&writers[i], NULL, cache_writer, &indices[i]));
    }
    for (repeat = 0; repeat < 20000; repeat++) {
        gpu_frame frame = gpu_snapshot();
        assert(frame.count == 2);
        for (i = 0; i < 2; i++) {
            assert(frame.samples[i].total == frame.samples[i].used + 1);
            assert(frame.samples[i].busy == frame.samples[i].used % 101);
            assert(frame.samples[i].temp == (long long)i);
        }
    }
    for (i = 0; i < 2; i++) assert(!pthread_join(writers[i], NULL));
}

static void parent_death(void)
{
    pid_t child, grandchild = 0, result;
    long long deadline;
    char path[96];
    int status;
    assert(!prctl(PR_SET_CHILD_SUBREAPER, 1));
    child = fork();
    assert(child >= 0);
    if (!child) {
        char *const args[] = {"sleep", "10", NULL};
        test_command("/usr/bin/sleep", args);
        _exit(0);
    }
    snprintf(path, sizeof(path), "/proc/%ld/task/%ld/children", (long)child, (long)child);
    deadline = gpu_milliseconds() + 500;
    do {
        FILE *file = fopen(path, "r");
        long id;
        if (file) {
            if (fscanf(file, "%ld", &id) == 1) grandchild = (pid_t)id;
            fclose(file);
        }
        if (!grandchild) poll(NULL, 0, 5);
    } while (!grandchild && gpu_milliseconds() < deadline);
    assert(grandchild > 0);
    assert(!kill(child, SIGKILL));
    assert(waitpid(child, &status, 0) == child);
    deadline = gpu_milliseconds() + 1000;
    do {
        result = waitpid(grandchild, &status, WNOHANG);
        if (!result) poll(NULL, 0, 5);
    } while (!result && gpu_milliseconds() < deadline);
    assert(result == grandchild && WIFSIGNALED(status) && WTERMSIG(status) == SIGKILL);
    assert(!prctl(PR_SET_CHILD_SUBREAPER, 0));
}

int main(int argc, char **argv)
{
    gpu_sample s;
    char text[256], memory[32];
    char *const sleeper[] = {"sleep", "10", NULL};
    char *const failure[] = {"false", NULL};
    long long start;
    int status;
    pid_t result;
    int fd_count, repeat;
    assert(gpu_page_first(0, 5000) == 0);
    assert(gpu_page_first(2, 5000) == 0);
    assert(gpu_page_first(3, 5000) == 2);
    assert(gpu_page_first(3, 10000) == 0);
    for (repeat = 0; repeat < GPU_MAX; repeat += 2)
        assert(gpu_page_first(GPU_MAX, (unsigned long long)repeat / 2 * 5000) == (size_t)repeat);
    gpu_page = 0;
    gpu_page_due = 5000;
    assert(gpu_display_first(5, 0) == 0);
    assert(gpu_display_first(5, 10000) == 2);
    assert(gpu_display_first(5, 20000) == 4);
    assert(gpu_display_first(5, 30000) == 0);
    assert(gpu_display_first(5, 30001) == 0);
    assert(gpu_selection_valid("auto") && gpu_selection_valid("off"));
    assert(gpu_selection_valid("0000:AB:1f.7"));
    assert(!gpu_selection_valid("0000:01:20.0"));
    assert(!gpu_selection_valid("0000:01:00.8"));
    assert(!gpu_selection_valid("-001:01:00.0"));
    assert(!gpu_selection_valid("../../passwd"));
    assert(!gpu_selection_valid("0000:01:00.0;id"));
    assert(gpu_number(" 0\n", 0, 100, -1) == 0);
    assert(gpu_number("100", 0, 100, -1) == 100);
    assert(gpu_number("101", 0, 100, -1) == -1);
    assert(gpu_number("12oops", 0, 100, -1) == -1);
    assert(gpu_number("", 0, 100, -1) == -1);
    assert(gpu_number("N/A", 0, 100, -1) == -1);
    assert(gpu_number("999999999999999999999999", 0, LLONG_MAX, -1) == -1);
    strcpy(text, "100, 12288, 16384, -5\n");
    s = gpu_csv(text);
    assert(s.busy == 100 && s.used == 12884901888LL &&
           s.total == 17179869184LL && s.temp == -5);
    strcpy(text, "N/A, [Not Supported], 16384, 0\n");
    s = gpu_csv(text);
    assert(s.busy == -1 && s.used == -1 && s.total == 17179869184LL && s.temp == 0);
    strcpy(text, "0, 0, 0, N/A\n");
    s = gpu_csv(text);
    assert(s.busy == 0 && s.used == 0 && s.total == -1 && s.temp == LLONG_MIN);
    strcpy(text, "0, 20, 10, 35\n");
    assert(gpu_csv(text).used == -1);
    strcpy(text, "0, 1, 2"); unknown(gpu_csv(text));
    strcpy(text, "0, 1, 2, 3, 4"); unknown(gpu_csv(text));
    strcpy(text, "N/A, N/A, N/A, N/A"); unknown(gpu_csv(text));
    gpu_memory(memory, sizeof(memory), -1); assert(!strcmp(memory, "N/A"));
    gpu_memory(memory, sizeof(memory), 19 * 1048576); assert(!strcmp(memory, "19.0MiB"));
    gpu_memory(memory, sizeof(memory), 17179869184LL); assert(!strcmp(memory, "16.0GiB"));
    gpu_memory(memory, sizeof(memory), LLONG_MAX); assert(strlen(memory) <= 11);
    assert(!gpu_detect("off") && !gpu_count);
    assert(!gpu_detect("ffff:ff:1f.7") && gpu_count == 1 && !gpu_devices[0].backend);
    gpu_cache[0] = s;
    gpu_stamp[0] = gpu_milliseconds() - 6000;
    unknown(gpu_snapshot().samples[0]);
    {
        gpu_frame frame = {.count = 3};
        gpu_count = 3;
        for (repeat = 0; repeat < 3; repeat++) {
            snprintf(gpu_devices[repeat].bus, 32, "0000:%02x:00.0", repeat + 1);
            gpu_devices[repeat].backend = repeat == 1 ? 1 : 2;
            frame.samples[repeat] = gpu_unknown;
        }
        frame.samples[1] = (gpu_sample){77, 1, 2, 42};
        strcpy(text, "00000000:03:00.0, 90, 4, 8, 55\n00000000:01:00.0, N/A, 1, 2, N/A\n");
        gpu_nvidia_parse(text, &frame);
        assert(frame.samples[0].busy == -1 && frame.samples[0].used == 1048576);
        assert(frame.samples[1].busy == 77 && frame.samples[1].used == 1);
        assert(frame.samples[2].busy == 90 && frame.samples[2].temp == 55);
        strcpy(text, "00000000:03:00.0, 1, 4, 8, 55\n0000:03:00.0, 2, 4, 8, 55\n");
        gpu_nvidia_parse(text, &frame);
        unknown(frame.samples[2]);
        assert(!gpu_bus_equal("0000:01:00.0junk", "0000:01:00.0"));
        assert(gpu_bus_equal("00000000:AB:00.0 ", "0000:ab:00.0"));
        assert(!gpu_bus_equal("0001:01:00.0", "0000:01:00.0"));
        for (repeat = 0; repeat < 3; repeat++) gpu_publish(repeat, frame.samples[repeat]);
        gpu_stamp[0] = gpu_milliseconds() - 6000;
        frame = gpu_snapshot();
        unknown(frame.samples[0]);
        assert(frame.samples[1].busy == 77);
        unknown(frame.samples[2]);
        /* Missing AMD path is independent of the populated neighboring cache. */
        snprintf(gpu_devices[0].device, PATH_MAX, "/nonexistent/g15stats-test");
        gpu_publish(0, gpu_amd(&gpu_devices[0]));
        assert(gpu_snapshot().samples[1].busy == 77);
    }

    /* Real processes exercise the deadline and reaping, not fake GPU files. */
    start = gpu_milliseconds();
    unknown(test_command("/usr/bin/sleep", sleeper));
    assert(gpu_milliseconds() - start >= 750);
    assert(gpu_milliseconds() - start < 1500);
    poll(NULL, 0, 50);
    unknown(test_command("/usr/bin/false", failure));
    result = waitpid(-1, &status, WNOHANG);
    assert(result == -1 && errno == ECHILD);
    unknown(test_command("/nonexistent/g15stats-test", failure));
    fd_count = descriptors();
    for (repeat = 0; repeat < 20; repeat++)
        unknown(test_command("/usr/bin/false", failure));
    assert(descriptors() == fd_count);
    {
        char payload[12001], output[GPU_MAX * 256];
        char *const args[] = {"printf", "%s", payload, NULL};
        memset(payload, 'x', sizeof(payload) - 1);
        payload[sizeof(payload) - 1] = 0;
        assert(gpu_command("/usr/bin/printf", args, output, sizeof(output)));
        assert(!strcmp(payload, output));
        assert(!gpu_command("/usr/bin/printf", args, output, 256));
        poll(NULL, 0, 50);
        unknown(test_command("/usr/bin/false", failure));
    }
    parent_death();
    cache_concurrency();

    if (argc == 2) {
        gpu_frame frame;
        size_t i;
        assert(!gpu_detect(argv[1]) && gpu_count);
        assert(!gpu_start());
        start = gpu_milliseconds();
        do {
            frame = gpu_snapshot();
            for (i = 0; i < frame.count; i++) if (frame.samples[i].total <= 0) break;
            if (i == frame.count) break;
            poll(NULL, 0, 10);
        } while (gpu_milliseconds() - start < 1800);
        for (i = 0; i < frame.count; i++) {
            char path[PATH_MAX + 32];
            s = frame.samples[i];
            if (i) assert(strcmp(gpu_devices[i-1].bus, gpu_devices[i].bus) < 0);
            snprintf(path, sizeof(path), "%s/vendor", gpu_devices[i].device);
            FILE *vendor = fopen(path, "r");
            unsigned id;
            assert(vendor && fscanf(vendor, "%x", &id) == 1);
            fclose(vendor);
            assert(id == (gpu_devices[i].backend == 1 ? 0x1002U : 0x10deU));
            printf("LIVE %zu %s vendor=%04x backend=%d utilization=%lld used=%lld total=%lld temp=%lld\n",
                   i+1, gpu_devices[i].bus, id, gpu_devices[i].backend, s.busy, s.used, s.total, s.temp);
            assert(s.busy >= 0 && s.used >= 0 && s.total > 0 && s.temp != LLONG_MIN);
        }
        for (repeat = 0; repeat < 100000; repeat++) {
            frame = gpu_snapshot();
            assert(frame.count == gpu_count);
            for (i = 0; i < frame.count; i++) assert(frame.samples[i].total > 0);
        }
        start = gpu_milliseconds();
        assert(!gpu_stop());
        assert(gpu_milliseconds() - start < 1200);
        puts("PASS all live adapters, sysfs vendor identities, concurrent snapshots and bounded shutdown");
        /* Read-only failure injection: point just the stopped AMD collector at
         * a nonexistent path, then prove the actual NVIDIA peer remains live. */
        if (gpu_count > 1) {
            size_t amd = GPU_MAX, nv = GPU_MAX;
            char saved[PATH_MAX];
            for (i = 0; i < gpu_count; i++) {
                if (gpu_devices[i].backend == 1) amd = i;
                if (gpu_devices[i].backend == 2) nv = i;
            }
            if (amd < GPU_MAX && nv < GPU_MAX) {
                strcpy(saved, gpu_devices[amd].device);
                strcpy(gpu_devices[amd].device, "/nonexistent/g15stats-test");
                for (i = 0; i < gpu_count; i++) gpu_publish(i, gpu_unknown);
                assert(!gpu_start());
                start = gpu_milliseconds();
                do {
                    frame = gpu_snapshot();
                    if (frame.samples[nv].total > 0) break;
                    poll(NULL, 0, 10);
                } while (gpu_milliseconds() - start < 1800);
                assert(frame.samples[nv].total > 0 && frame.samples[nv].busy >= 0);
                unknown(frame.samples[amd]);
                assert(!gpu_stop());
                strcpy(gpu_devices[amd].device, saved);
                puts("PASS failed AMD collector isolated from real live NVIDIA peer");
            }
        }
    }
    puts("PASS GPU parsing, units, missing data, expiry, selection, deadline, reap and parent death");
    return 0;
}