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

#pragma once

#include <map>
#include <memory>
#include <mutex>
#include <tuple>
#include <utility>

#include "CUDAMemHelpers.h"

// Read-only lookup tables that depend only on the detector geometry (pixel -> azimuthal bin, the
// per-pixel correction factors, the pixel mask). One analysis engine is built per worker thread, so
// each of those used to upload its own copy: on a 18 Mpx detector that is ~220 MB per thread, and
// 32 threads spent ~7 GB of device memory on identical data.
//
// Upload once per GPU instead and hand every engine on that GPU a shared pointer to the same table.
// The cache is keyed by (device, key) because a worker thread is pinned round-robin to a device
// (pin_gpu()), so on a multi-GPU node each device keeps its own copy - a kernel may only read memory
// resident on the device it runs on. `key` identifies the table's source data; use the address of the
// host vector that produced it, which lives in the experiment / integration mapping and therefore
// outlives every engine.
//
// A bare address is not enough on its own to say "same table", though: a host buffer can be mutated
// in place, or freed and a new one allocated at the same address, and either would hand the caller a
// device copy of something else - silently, since the data is only ever read. So the byte length and
// a checksum of the bytes actually uploaded are part of the key too. Both are computed once per
// engine construction, against an upload of the same buffer, so they cost nothing measurable.
//
// Entries are held weakly, so the tables are released once the last engine using them is gone.

namespace jfjoch_cuda_shared_tables {
    // (device, source address, byte length, checksum of the bytes)
    using TableKey = std::tuple<int, const void *, size_t, uint64_t>;

    struct Registry {
        std::mutex m;
        std::map<TableKey, std::weak_ptr<void>> tables;
    };

    // FNV-1a. Not a cryptographic hash and does not need to be - it exists to notice that the bytes
    // behind a reused address changed, not to resist anyone.
    inline uint64_t checksum(const void *data, size_t bytes) {
        const auto *p = static_cast<const unsigned char *>(data);
        uint64_t h = 1469598103934665603ULL;
        for (size_t i = 0; i < bytes; i++) {
            h ^= p[i];
            h *= 1099511628211ULL;
        }
        return h;
    }

    inline Registry &registry() {
        static Registry r;
        return r;
    }

    // Not called cuda_err: the .cu files that include this header define their own such helper in an
    // anonymous namespace, and a second one at global scope would make every call ambiguous.
    inline void check(cudaError_t val) {
        if (val != cudaSuccess)
            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
    }
}

// Return the device-resident copy of `host` (`count` elements) for the calling thread's GPU,
// uploading it on `stream` the first time it is asked for.
template <typename T>
std::shared_ptr<CudaDevicePtr<T>> SharedDeviceTable(const void *key, size_t count, const T *host,
                                                    cudaStream_t stream) {
    int device = 0;
    jfjoch_cuda_shared_tables::check(cudaGetDevice(&device));

    const size_t bytes = count * sizeof(T);
    const jfjoch_cuda_shared_tables::TableKey table_key{
        device, key, bytes, jfjoch_cuda_shared_tables::checksum(host, bytes)};

    auto &reg = jfjoch_cuda_shared_tables::registry();
    // The upload happens while the lock is held: another worker must not obtain the pointer before
    // its content is on the device.
    std::lock_guard lock(reg.m);
    if (auto it = reg.tables.find(table_key); it != reg.tables.end()) {
        if (auto cached = it->second.lock())
            return std::static_pointer_cast<CudaDevicePtr<T>>(cached);
    }
    // Drop entries whose table is gone before adding one. Without this a long session that reloads
    // masks or remaps geometry accumulates a dead entry per distinct content, for ever.
    for (auto it = reg.tables.begin(); it != reg.tables.end();)
        it = it->second.expired() ? reg.tables.erase(it) : std::next(it);

    // Free on the device that allocated it - the last engine to drop the table may well be a worker
    // pinned to a different GPU.
    std::shared_ptr<CudaDevicePtr<T>> table(new CudaDevicePtr<T>(count), [device](CudaDevicePtr<T> *p) {
        int current = 0;
        cudaGetDevice(&current);
        cudaSetDevice(device);
        delete p;
        cudaSetDevice(current);
    });
    jfjoch_cuda_shared_tables::check(
        cudaMemcpyAsync(table->get(), host, bytes, cudaMemcpyHostToDevice, stream));
    jfjoch_cuda_shared_tables::check(cudaStreamSynchronize(stream));

    reg.tables[table_key] = std::shared_ptr<void>(table);
    return table;
}