// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "CudaJitCache.h"

#include <cmath>
#include <cstdlib>
#include <filesystem>

#ifdef __linux__
#include <cerrno>
#include <cstring>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/vfs.h>
#include <unistd.h>
#endif

namespace {
    // Why UsePrivateCudaJitCacheIfNeeded moved the cache, for LogCudaJitCache; empty when it did not.
    std::string private_cache_reason;
}

#ifdef __linux__
namespace {
    // statfs f_type of the filesystems a home directory is likely to be on. A network one is not wrong in
    // itself, but each cuFFT plan opens and locks the cache about ten times: at 1 ms a round trip that is
    // 15-20 ms a plan, at 5 ms as slow as no cache.
    struct FsType {
        long magic;
        const char *name;
    };

    constexpr FsType fs_types[] = {
        {0x6969, "nfs"}, {0x5346414F, "afs"}, {0x6B414653, "afs"}, {0x47504653, "gpfs"},
        {0x0BD00BD0, "lustre"}, {0x00C36400, "ceph"}, {0x19830326, "beegfs"}, {0xFF534D42, "cifs"},
        {0xFE534D42, "smb2"}, {0x517B, "smb"}, {0x01021997, "9p"}, {0x65735546, "fuse"},
        {0x58465342, "xfs"}, {0xEF53, "ext4"}, {0x9123683E, "btrfs"}, {0x01021994, "tmpfs"},
        {0x794C7630, "overlay"}, {0x2FC12FC1, "zfs"},
    };
}
#endif

CudaJitCache InspectCudaJitCache() {
    CudaJitCache c;
    const char *path = std::getenv("CUDA_CACHE_PATH");
    const char *disable = std::getenv("CUDA_CACHE_DISABLE");
    const char *maxsize = std::getenv("CUDA_CACHE_MAXSIZE");

    if (path != nullptr && *path != '\0')
        c.path = path;
#ifdef _WIN32
    else
        c.path = "%APPDATA%\\NVIDIA\\ComputeCache";
#else
    else if (const char *home = std::getenv("HOME"); home != nullptr && *home != '\0')
        c.path = std::string(home) + "/.nv/ComputeCache";
#endif

    if (disable != nullptr && std::string(disable) == "1")
        c.problem = "CUDA_CACHE_DISABLE=1 turns it off";
    else if (maxsize != nullptr && std::strtoull(maxsize, nullptr, 10) < 64 * 1024 * 1024)
        c.problem = fmt::format("CUDA_CACHE_MAXSIZE={} keeps too little of it", maxsize);
    else if (c.path.empty())
        c.problem = "neither CUDA_CACHE_PATH nor HOME is set (a systemd service without User= has no HOME)";

#ifdef __linux__
    if (c.path.empty())
        return c;
    // The driver makes the directory when it is missing, so the nearest one that exists must be writable.
    std::error_code ec;
    std::filesystem::path dir = c.path;
    while (!std::filesystem::exists(dir, ec) && dir.has_parent_path() && dir != dir.parent_path())
        dir = dir.parent_path();

    struct statfs s{};
    if (statfs(dir.c_str(), &s) == 0) {
        c.filesystem = fmt::format("{:#x}", static_cast<unsigned long>(s.f_type));
        for (const auto &t: fs_types)
            if (static_cast<unsigned long>(s.f_type) == static_cast<unsigned long>(t.magic))
                c.filesystem = t.name;
    }

    if (!c.problem.empty())
        return c;
    if (access(dir.c_str(), W_OK | X_OK) != 0) {
        c.problem = fmt::format("{} is not writable ({})", dir.string(), std::strerror(errno));
        return c;
    }
    // Each lookup opens the index read-write and read-locks it with fcntl: where either fails (a read-only
    // cache, a filesystem without POSIX locks) the driver compiles every time, though the cache is there.
    const std::string index = c.path + "/index";
    if (std::filesystem::exists(index, ec)) {
        const int fd = open(index.c_str(), O_RDWR);
        if (fd < 0) {
            c.problem = fmt::format("its index cannot be opened for writing ({})", std::strerror(errno));
            return c;
        }
        struct flock lock{};
        lock.l_type = F_RDLCK;
        lock.l_whence = SEEK_SET;
        if (fcntl(fd, F_SETLK, &lock) != 0 && errno != EAGAIN && errno != EACCES)
            c.problem = fmt::format("its index cannot be locked ({})", std::strerror(errno));
        lock.l_type = F_UNLCK;
        fcntl(fd, F_SETLK, &lock);
        close(fd);
    }
#endif
    return c;
}

void UsePrivateCudaJitCacheIfNeeded() {
#ifdef __linux__
    if (const char *path = std::getenv("CUDA_CACHE_PATH"); path != nullptr && *path != '\0')
        return;
    if (const char *disable = std::getenv("CUDA_CACHE_DISABLE"); disable != nullptr && std::string(disable) == "1")
        return;
    const CudaJitCache c = InspectCudaJitCache();
    if (c.problem.empty())
        return;

    // /var/tmp is shared with every user: the directory is used only if it is a real directory (not a
    // symlink), this user's, and closed to everyone else.
    const std::string dir = fmt::format("/var/tmp/jfjoch-cuda-cache-{}", getuid());
    if (mkdir(dir.c_str(), 0700) != 0 && errno != EEXIST)
        return;
    struct stat st{};
    if (lstat(dir.c_str(), &st) != 0 || !S_ISDIR(st.st_mode) || st.st_uid != getuid() || (st.st_mode & 077) != 0)
        return;
    setenv("CUDA_CACHE_PATH", dir.c_str(), 1);
    private_cache_reason = c.problem;
#endif
}

void LogCudaJitCache(Logger &logger, double plan_ms) {
    if (std::isnan(plan_ms))
        return;
    const CudaJitCache c = InspectCudaJitCache();
    // Measured: 0.4-1.7 ms from a working cache, 60-130 ms with none.
    const bool slow = plan_ms > 10.0;
    const std::string where = c.path.empty() ? std::string("none")
            : c.filesystem.empty() ? c.path : fmt::format("{} ({})", c.path, c.filesystem);
    if (c.problem.empty() && !slow) {
        logger.Info("CUDA JIT cache {}{}: a cuFFT plan takes {:.1f} ms", where,
                    private_cache_reason.empty() ? std::string()
                            : fmt::format(", used because {}", private_cache_reason),
                    plan_ms);
        return;
    }
    logger.Warning("CUDA JIT cache {}: {}a cuFFT plan takes {:.0f} ms{}, which can make molecular "
                   "replacement several times slower. Give the process a cache "
                   "on a local disk that only its user can write: CUDA_CACHE_PATH=<dir>, for instance "
                   "Environment=CUDA_CACHE_PATH=/var/cache/jfjoch/nv in a systemd unit",
                   where, c.problem.empty() ? std::string() : c.problem + "; ", plan_ms,
                   slow ? " (~1 ms from a working cache)" : "");
}