// 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 <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.
//
// Entries are held weakly, so the tables are released once the last engine using them is gone.

namespace jfjoch_cuda_shared_tables {
    struct Registry {
        std::mutex m;
        std::map<std::pair<int, const void *>, std::weak_ptr<void>> tables;
    };

    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));

    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);
    auto &slot = reg.tables[{device, key}];
    if (auto cached = slot.lock())
        return std::static_pointer_cast<CudaDevicePtr<T>>(cached);

    // 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, count * sizeof(T), cudaMemcpyHostToDevice, stream));
    jfjoch_cuda_shared_tables::check(cudaStreamSynchronize(stream));

    slot = std::shared_ptr<void>(table);
    return table;
}