/* cuda_drv_test2.c - minimal CUDA driver-API smoke test (dlopen, no toolkit)
 * with copy-path pause points for live inspection.
 *
 * Same sequence as cuda_drv_test.c, plus:
 *   CDT_CP_PAUSE=<s>  sleep before cuMemcpyHtoD (after cuMemAlloc): the
 *                     context is live, the device pointer is known - probe
 *                     the GPU VA translation while parked.
 *   CDT_HP_PAUSE=<s>  sleep after cuMemcpyHtoD (before DtoH): inspect what
 *                     the HtoD left behind.
 *   CDT_WAIT=<s>      park after the test (existing behaviour).
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <dlfcn.h>
#include <stdint.h>

typedef int CUresult;
typedef int CUdevice;
typedef void *CUcontext;
typedef uint64_t CUdeviceptr;

static void fail_info(CUresult r)
{
    const char *(*errname)(CUresult, const char **) = dlsym(RTLD_DEFAULT, "cuGetErrorName");
    const char *(*errstr)(CUresult, const char **) = dlsym(RTLD_DEFAULT, "cuGetErrorString");
    const char *s = NULL;
    const char *hold = getenv("CDT_HOLD");

    if (errname && errname(r, &s) == 0 && s)
        printf("cuGetErrorName: %s\n", s);
    if (errstr && errstr(r, &s) == 0 && s)
        printf("cuGetErrorString: %s\n", s);
    if (hold && atoi(hold) > 0) {
        printf("CDT_HOLD: sleeping %d s for inspection\n", atoi(hold));
        fflush(stdout);
        sleep((unsigned)atoi(hold));
    }
}

#define CHECK(expr) do { \
    CUresult r = (expr); \
    if (r) { printf("%-24s FAILED rc=%d\n", #expr, r); fail_info(r); return 1; } \
    printf("%-24s OK\n", #expr); \
} while (0)

#define MARK(name) do { printf("-> %s\n", name); } while (0)

static void pause_env(const char *name)
{
    const char *v = getenv(name);

    if (v && atoi(v) > 0) {
        printf("%s: parked %d s\n", name, atoi(v));
        fflush(stdout);
        sleep((unsigned)atoi(v));
    }
}

int main(void)
{
    setvbuf(stdout, NULL, _IONBF, 0);

    void *h = dlopen("libcuda.so.1", RTLD_NOW);
    if (!h) { printf("dlopen libcuda.so.1 failed: %s\n", dlerror()); return 1; }

    CUresult (*cuInit)(unsigned int) = dlsym(h, "cuInit");
    CUresult (*cuDeviceGetCount)(int *) = dlsym(h, "cuDeviceGetCount");
    CUresult (*cuDeviceGet)(CUdevice *, int) = dlsym(h, "cuDeviceGet");
    CUresult (*cuCtxCreate)(CUcontext *, unsigned int, CUdevice) = dlsym(h, "cuCtxCreate_v2");
    CUresult (*cuMemAlloc)(CUdeviceptr *, size_t) = dlsym(h, "cuMemAlloc_v2");
    CUresult (*cuMemcpyHtoD)(CUdeviceptr, const void *, size_t) = dlsym(h, "cuMemcpyHtoD_v2");
    CUresult (*cuMemcpyDtoH)(void *, CUdeviceptr, size_t) = dlsym(h, "cuMemcpyDtoH_v2");
    CUresult (*cuMemFree)(CUdeviceptr) = dlsym(h, "cuMemFree_v2");
    CUresult (*cuCtxDestroy)(CUcontext) = dlsym(h, "cuCtxDestroy_v2");

    if (!cuInit || !cuDeviceGetCount || !cuDeviceGet || !cuCtxCreate ||
        !cuMemAlloc || !cuMemcpyHtoD || !cuMemcpyDtoH || !cuMemFree || !cuCtxDestroy) {
        printf("dlsym failed\n");
        return 1;
    }

    int count = 0;
    CUdevice dev;
    CUcontext ctx = NULL;
    CUdeviceptr dptr = 0;
    unsigned int buf[64];
    unsigned int back[64];
    size_t i;

    MARK("cuInit");
    CHECK(cuInit(0));
    MARK("cuDeviceGetCount");
    CHECK(cuDeviceGetCount(&count));
    printf("device count = %d\n", count);
    if (count < 1) return 1;

    MARK("cuDeviceGet");
    CHECK(cuDeviceGet(&dev, 0));
    printf("device 0 = %d\n", dev);

    MARK("cuCtxCreate");
    {
        const char *pause = getenv("CDT_PAUSE");
        if (pause && atoi(pause) > 0) {
            printf("CDT_PAUSE: sleeping %d s before cuCtxCreate\n", atoi(pause));
            fflush(stdout);
            sleep((unsigned)atoi(pause));
        }
    }
    CHECK(cuCtxCreate(&ctx, 0, dev));
    printf("context = %p\n", ctx);

    MARK("cuMemAlloc");
    CHECK(cuMemAlloc(&dptr, sizeof(buf)));
    printf("device ptr = 0x%llx (%zu bytes)\n", (unsigned long long)dptr, sizeof(buf));

    for (i = 0; i < 64; i++) buf[i] = (unsigned int)(i * 3 + 1);
    pause_env("CDT_CP_PAUSE");
    MARK("cuMemcpyHtoD");
    CHECK(cuMemcpyHtoD(dptr, buf, sizeof(buf)));
    pause_env("CDT_HP_PAUSE");
    memset(back, 0, sizeof(back));
    MARK("cuMemcpyDtoH");
    CHECK(cuMemcpyDtoH(back, dptr, sizeof(buf)));
    {
        size_t matched = 0;
        for (i = 0; i < 64; i++)
            if (buf[i] == back[i]) matched++;
        printf("roundtrip %s matched=%zu/64 back[0..7]=%u,%u,%u,%u,%u,%u,%u,%u expected 1,4,7,10,13,16,19,22\n",
               matched == 64 ? "match" : "MISMATCH", matched,
               back[0], back[1], back[2], back[3],
               back[4], back[5], back[6], back[7]);
    }

    MARK("cuMemFree");
    CHECK(cuMemFree(dptr));
    MARK("cuCtxDestroy");
    CHECK(cuCtxDestroy(ctx));
    printf("ALL OK\n");
    pause_env("CDT_WAIT");
    return 0;
}