/* SPDX-License-Identifier: GPL-2.0-or-later
 * Read-only Linux GPU telemetry. Included privately by g15stats.c.
 */
#ifndef G15STATS_GPU_H
#define G15STATS_GPU_H
#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <glob.h>
#include <limits.h>
#include <poll.h>
#include <pthread.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/prctl.h>
#include <sys/syscall.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

typedef struct {
    long long busy, used, total, temp; /* percent, bytes, bytes, Celsius */
} gpu_sample;

static const gpu_sample gpu_unknown = {-1, -1, -1, LLONG_MIN};
#define GPU_MAX 64
typedef struct {
    int backend; /* 0 missing, 1 amdgpu, 2 nvidia */
    char device[PATH_MAX], bus[32];
} gpu_adapter;
typedef struct {
    size_t count;
    gpu_sample samples[GPU_MAX];
} gpu_frame;
static gpu_adapter gpu_devices[GPU_MAX]; /* immutable while workers run */
static size_t gpu_count;
static gpu_sample gpu_cache[GPU_MAX];
static long long gpu_stamp[GPU_MAX];
static pthread_mutex_t gpu_lock = PTHREAD_MUTEX_INITIALIZER;
static pthread_t gpu_threads[GPU_MAX];
static size_t gpu_thread_count;
static atomic_int gpu_stopping;
static pid_t gpu_pending = -1; /* NVIDIA worker owns until joined. */
static long long gpu_epoch;
static size_t gpu_page;
static long long gpu_page_due;

static int gpu_selection_valid(const char *selection)
{
    size_t i;
    if (!strcmp(selection, "auto") || !strcmp(selection, "off")) return 1;
    if (strlen(selection) != 12) return 0;
    for (i = 0; i < 12; i++) {
        if (i == 4 || i == 7) { if (selection[i] != ':') return 0; }
        else if (i == 10) { if (selection[i] != '.') return 0; }
        else if (!isxdigit((unsigned char)selection[i])) return 0;
    }
    return selection[8] <= '1' && selection[11] <= '7';
}

static long long gpu_milliseconds(void)
{
    struct timespec now;
    if (clock_gettime(CLOCK_MONOTONIC, &now))
        return 0;
    return (long long)now.tv_sec * 1000 + now.tv_nsec / 1000000;
}

/* Reject partial values, overflow and unsupported textual fields. */
static long long gpu_number(const char *text, long long min, long long max,
                            long long missing)
{
    char *end;
    long long value;
    while (isspace((unsigned char)*text)) text++;
    if (!isdigit((unsigned char)*text) && *text != '-') return missing;
    errno = 0;
    value = strtoll(text, &end, 10);
    if (end == text || errno) return missing;
    while (isspace((unsigned char)*end)) end++;
    return *end || value < min || value > max ? missing : value;
}

static long long gpu_file(const char *path, long long min, long long max,
                          long long missing)
{
    char text[96];
    FILE *file = fopen(path, "re");
    long long value = missing;
    if (!file) return missing;
    if (fgets(text, sizeof(text), file) && !ferror(file)) {
        /* sysfs values fit in one short line; do not accept truncated input. */
        if (fgetc(file) == EOF && !ferror(file))
            value = gpu_number(text, min, max, missing);
    }
    fclose(file);
    return value;
}

/* Select once, in sorted PCI-address order, not unstable DRM card order.
 * Explicit PCI selection keeps an N/A column even for missing devices.
 */
static int gpu_detect(const char *selection)
{
    glob_t paths = {0};
    size_t i;
    gpu_count = 0;
    gpu_epoch = gpu_milliseconds();
    gpu_page = 0;
    gpu_page_due = gpu_epoch + 5000;
    memset(gpu_stamp, 0, sizeof(gpu_stamp));
    for (i = 0; i < GPU_MAX; i++) gpu_cache[i] = gpu_unknown;
    if (!strcmp(selection, "off")) return 0;
    int error = glob("/sys/bus/pci/devices/*", 0, NULL, &paths);
    if (error && error != GLOB_NOMATCH) {
        globfree(&paths);
        return -1;
    }
    for (i = 0; i < paths.gl_pathc; i++) {
        char path[PATH_MAX], driver[PATH_MAX];
        const char *bus = strrchr(paths.gl_pathv[i], '/') + 1;
        ssize_t len;
        int backend;
        if (strcmp(selection, "auto") && strcasecmp(selection, bus)) continue;
        snprintf(path, sizeof(path), "%s/driver", paths.gl_pathv[i]);
        len = readlink(path, driver, sizeof(driver) - 1);
        if (len < 0) continue;
        driver[len] = 0;
        if (!strcmp(strrchr(driver, '/') + 1, "amdgpu"))
            backend = 1;
        else if (!strcmp(strrchr(driver, '/') + 1, "nvidia"))
            backend = 2;
        else continue;
        if (gpu_count == GPU_MAX) {
            globfree(&paths);
            gpu_count = 0;
            return -1; /* Never silently truncate the inventory. */
        }
        gpu_devices[gpu_count].backend = backend;
        snprintf(gpu_devices[gpu_count].device, PATH_MAX, "%s", paths.gl_pathv[i]);
        snprintf(gpu_devices[gpu_count].bus, 32, "%s", bus);
        gpu_count++;
    }
    globfree(&paths);
    if (!gpu_count && strcmp(selection, "auto")) {
        memset(&gpu_devices[0], 0, sizeof(gpu_devices[0]));
        snprintf(gpu_devices[0].bus, 32, "%s", selection);
        gpu_count = 1;
    }
    return 0;
}

static gpu_sample gpu_amd(const gpu_adapter *adapter)
{
    gpu_sample sample = gpu_unknown;
    char path[PATH_MAX + 64];
    glob_t paths = {0};
    snprintf(path, sizeof(path), "%s/gpu_busy_percent", adapter->device);
    sample.busy = gpu_file(path, 0, 100, -1);
    snprintf(path, sizeof(path), "%s/mem_info_vram_used", adapter->device);
    sample.used = gpu_file(path, 0, LLONG_MAX, -1);
    snprintf(path, sizeof(path), "%s/mem_info_vram_total", adapter->device);
    sample.total = gpu_file(path, 1, LLONG_MAX, -1);
    snprintf(path, sizeof(path), "%s/hwmon/hwmon*/temp1_input", adapter->device);
    if (!glob(path, 0, NULL, &paths) && paths.gl_pathc) {
        long long temp = gpu_file(paths.gl_pathv[0], -273000, 1000000, LLONG_MIN);
        if (temp != LLONG_MIN) sample.temp = temp / 1000;
    }
    globfree(&paths);
    if (sample.total >= 0 && sample.used > sample.total) sample.used = -1;
    return sample;
}

static gpu_sample gpu_csv(char *text)
{
    gpu_sample sample = gpu_unknown;
    char *fields[4], *cursor = text;
    int i;
    for (i = 0; i < 4; i++) {
        fields[i] = strsep(&cursor, ",");
        if (!fields[i]) return sample;
    }
    if (cursor) return sample;
    sample.busy = gpu_number(fields[0], 0, 100, -1);
    sample.used = gpu_number(fields[1], 0, LLONG_MAX / 1048576, -1);
    sample.total = gpu_number(fields[2], 1, LLONG_MAX / 1048576, -1);
    sample.temp = gpu_number(fields[3], -273, 1000, LLONG_MIN);
    if (sample.used >= 0) sample.used *= 1048576;
    if (sample.total >= 0) sample.total *= 1048576;
    if (sample.total >= 0 && sample.used > sample.total) sample.used = -1;
    return sample;
}

/* Only the collector calls this. Never wait for a driver in the draw thread.
 * A killed but uninterruptible child remains owned, preventing process buildup.
 */
static int gpu_above_stdio(int fd)
{
    int copy;
    if (fd < 0 || fd >= 3) return fd;
    copy = fcntl(fd, F_DUPFD_CLOEXEC, 3);
    close(fd);
    return copy;
}

static int gpu_command(const char *executable, char *const args[], char *text, size_t capacity)
{
    int pipes[2], status, nullfd;
    pid_t pid, parent = getpid(), result;
    size_t used = 0;
    long long deadline;
    long maxfd = sysconf(_SC_OPEN_MAX);
    char *const env[] = {"LC_ALL=C", NULL};
    if (gpu_pending > 0) {
        result = waitpid(gpu_pending, &status, WNOHANG);
        if (!result || (result < 0 && errno == EINTR)) return 0;
        gpu_pending = -1;
    }
    if (pipe2(pipes, O_CLOEXEC | O_NONBLOCK)) return 0;
    pipes[0] = gpu_above_stdio(pipes[0]);
    pipes[1] = gpu_above_stdio(pipes[1]);
    if (pipes[0] < 0 || pipes[1] < 0) {
        if (pipes[0] >= 0) close(pipes[0]);
        if (pipes[1] >= 0) close(pipes[1]);
        return 0;
    }
    /* Only the reader is nonblocking. A large all-card response must not
     * make nvidia-smi fail with EAGAIN when the pipe temporarily fills. */
    if (fcntl(pipes[1], F_SETFL, 0) < 0) {
        close(pipes[0]);
        close(pipes[1]);
        return 0;
    }
    nullfd = gpu_above_stdio(open("/dev/null", O_RDWR | O_CLOEXEC));
    if (nullfd < 0) { close(pipes[0]); close(pipes[1]); return 0; }
    pid = fork();
    if (!pid) {
        /* Async-signal-safe child path: no stdio, allocation or shell. */
        if (prctl(PR_SET_PDEATHSIG, SIGKILL) || getppid() != parent) _exit(127);
        if (dup2(nullfd, 0) < 0 || dup2(pipes[1], 1) < 0 || dup2(nullfd, 2) < 0)
            _exit(127);
#ifdef SYS_close_range
        if (syscall(SYS_close_range, 3U, ~0U, 0) < 0)
#endif
        {
            long fd;
            if (maxfd < 0) _exit(127);
            for (fd = 3; fd < maxfd; fd++) close((int)fd);
        }
        execve(executable, args, env);
        _exit(127);
    }
    close(nullfd);
    close(pipes[1]);
    if (pid < 0) { close(pipes[0]); return 0; }
    deadline = gpu_milliseconds() + 750;
    for (;;) {
        ssize_t n = read(pipes[0], text + used, capacity - 1 - used);
        if (n > 0) used += (size_t)n;
        result = waitpid(pid, &status, WNOHANG);
        if (result == pid) {
            /* Drain the final bytes after exit before parsing. */
                        while (used < capacity - 1 &&
                                     (n = read(pipes[0], text + used, capacity - 1 - used)) > 0)
                used += (size_t)n;
            close(pipes[0]);
            text[used] = 0;
            return used < capacity - 1 && WIFEXITED(status) && !WEXITSTATUS(status);
        }
        if (result < 0 && errno != EINTR) { close(pipes[0]); return 0; }
        if (atomic_load(&gpu_stopping) || gpu_milliseconds() >= deadline || used == capacity - 1) break;
        /* Avoid spinning on POLLHUP if stdout closes before the process exits. */
        poll(NULL, 0, 20);
    }
    kill(pid, SIGKILL);
    gpu_pending = pid;
    close(pipes[0]);
    return 0;
}

/* Match numeric PCI identity: NVIDIA commonly prints an eight-digit domain.
 * Malformed/duplicate rows must never contaminate another adapter's sample. */
static int gpu_bus_equal(const char *a, const char *b)
{
    unsigned ad, ab, as, af, bd, bb, bs, bf;
    int n = 0, m = 0;
    if (sscanf(a, "%8x:%2x:%2x.%1x %n", &ad, &ab, &as, &af, &n) != 4 || a[n]) return 0;
    if (sscanf(b, "%8x:%2x:%2x.%1x %n", &bd, &bb, &bs, &bf, &m) != 4 || b[m]) return 0;
    return ad == bd && ab == bb && as == bs && af == bf;
}

static void gpu_nvidia_parse(char *text, gpu_frame *frame)
{
    char *line;
    unsigned char seen[GPU_MAX] = {0};
    size_t i;
    while ((line = strsep(&text, "\n"))) {
        char *bus = strsep(&line, ",");
        if (!line) continue;
        for (i = 0; i < gpu_count; i++) {
            if (gpu_devices[i].backend != 2 || !gpu_bus_equal(bus, gpu_devices[i].bus)) continue;
            frame->samples[i] = seen[i] ? gpu_unknown : gpu_csv(line);
            seen[i] = 1;
            break;
        }
    }
}

static void gpu_nvidia(gpu_frame *frame)
{
    char text[GPU_MAX * 256];
    char *const args[] = {"nvidia-smi",
        "--query-gpu=pci.bus_id,utilization.gpu,memory.used,memory.total,temperature.gpu",
        "--format=csv,noheader,nounits", NULL};
    if (gpu_command("/usr/bin/nvidia-smi", args, text, sizeof(text)))
        gpu_nvidia_parse(text, frame);
}

static void gpu_publish(size_t i, gpu_sample sample)
{
    pthread_mutex_lock(&gpu_lock);
    gpu_cache[i] = sample;
    gpu_stamp[i] = gpu_milliseconds();
    pthread_mutex_unlock(&gpu_lock);
}

/* Each AMD adapter has an independent worker; one NVIDIA worker queries all
 * NVIDIA cards within one deadline. A stalled sysfs driver cannot stall peers. */
static void *gpu_watch(void *arg)
{
    size_t index = (gpu_adapter *)arg - gpu_devices, i;
    while (!atomic_load(&gpu_stopping)) {
        if (gpu_devices[index].backend == 1) {
            gpu_publish(index, gpu_amd(&gpu_devices[index]));
        } else {
            gpu_frame frame = {.count = gpu_count};
            for (i = 0; i < gpu_count; i++) frame.samples[i] = gpu_unknown;
            gpu_nvidia(&frame);
            pthread_mutex_lock(&gpu_lock);
            for (i = 0; i < gpu_count; i++) if (gpu_devices[i].backend == 2) {
                gpu_cache[i] = frame.samples[i];
                gpu_stamp[i] = gpu_milliseconds();
            }
            pthread_mutex_unlock(&gpu_lock);
        }
        for (i = 0; i < 100 && !atomic_load(&gpu_stopping); i++) poll(NULL, 0, 20);
    }
    return NULL;
}

static int gpu_start(void)
{
    size_t i;
    int nvidia = 0, failed = 0;
    atomic_store(&gpu_stopping, 0);
    for (i = 0; i < gpu_count; i++) {
        if (!gpu_devices[i].backend || (gpu_devices[i].backend == 2 && nvidia++)) continue;
        if (pthread_create(&gpu_threads[gpu_thread_count], NULL, gpu_watch, &gpu_devices[i])) failed = 1;
        else gpu_thread_count++;
    }
    return failed;
}

/* One shared shutdown deadline, not one per adapter. A stuck kernel read may
 * outlive this call; its static storage remains valid until process exit.
 * Never redetect/restart after a failed stop. */
static int gpu_stop(void)
{
    struct timespec deadline;
    long long end = gpu_milliseconds() + 1000;
    size_t i;
    int failed = 0;
    atomic_store(&gpu_stopping, 1);
    clock_gettime(CLOCK_REALTIME, &deadline);
    deadline.tv_sec++;
    for (i = 0; i < gpu_thread_count; i++) {
        if (pthread_timedjoin_np(gpu_threads[i], NULL, &deadline)) {
            pthread_detach(gpu_threads[i]);
            failed = 1;
        }
    }
    gpu_thread_count = 0;
    if (!failed && gpu_pending > 0) {
        int status;
        pid_t result;
        do {
            result = waitpid(gpu_pending, &status, WNOHANG);
            if (result == gpu_pending || (result < 0 && errno == ECHILD)) {
                gpu_pending = -1;
                break;
            }
            poll(NULL, 0, 5);
        } while (gpu_milliseconds() < end);
        if (gpu_pending > 0) failed = 1;
    }
    return failed;
}

static gpu_frame gpu_snapshot(void)
{
    gpu_frame frame = {.count = gpu_count};
    size_t i;
    pthread_mutex_lock(&gpu_lock);
    for (i = 0; i < gpu_count; i++)
        frame.samples[i] = !gpu_stamp[i] || gpu_milliseconds() - gpu_stamp[i] > 5000 ? gpu_unknown : gpu_cache[i];
    pthread_mutex_unlock(&gpu_lock);
    return frame;
}

static size_t gpu_page_first(size_t count, unsigned long long elapsed)
{
    return count > 2 ? (elapsed / 5000 % ((count + 1) / 2)) * 2 : 0;
}

/* Main thread only. Advance at most one pair per actual draw, so -r 10 (or
 * leaving the summary screen) cannot alias the cycle and hide half the GPUs. */
static size_t gpu_display_first(size_t count, long long now)
{
    if (count > 2 && now >= gpu_page_due) {
        gpu_page = (gpu_page + 1) % ((count + 1) / 2);
        gpu_page_due = now + 5000;
    }
    return gpu_page_first(count, gpu_page * 5000ULL);
}

/* Each line fits the 52-pixel column (13 small-font character cells).
 * Separate used/total rows retain useful precision, even on large cards.
 */
static void gpu_memory(char *text, size_t size, long long bytes)
{
    const char *unit = "MiB";
    double value = (double)bytes / 1048576;
    if (bytes < 0) { snprintf(text, size, "N/A"); return; }
    if (value >= 1024) { value /= 1024; unit = "GiB"; }
    if (value >= 1024) { value /= 1024; unit = "TiB"; }
    if (value >= 1024) { value /= 1024; unit = "PiB"; }
    snprintf(text, size, "%.1f%s", value, unit);
}
#endif