// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#include <optional>
#include <thread>

#ifdef JFJOCH_USE_CUDA

#include <cstring>
#include <future>
#include <latch>
#include <set>
#include <thread>

#include "../image_analysis/indexing/CUDAMemHelpers.h"
#include "../image_analysis/indexing/IndexerThreadPool.h"

// Whether the current GPU shares the host's memory (GB10, Jetson). A test that fills the card until an
// allocation fails would there fill the host's RAM - every process on the machine runs out with it.
static bool gpu_is_integrated() {
    int device = 0, integrated = 0;
    return cudaGetDevice(&device) == cudaSuccess
           && cudaDeviceGetAttribute(&integrated, cudaDevAttrIntegrated, device) == cudaSuccess && integrated;
}

// The broker starts fresh worker threads for every data collection. A stream per thread ever started
// is a device-memory leak that ends a long-running broker, so threads that follow one another have to
// end up on the same stream.
TEST_CASE("CudaAllocationStream_ReusedByLaterThreads", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    std::set<cudaStream_t> seen;
    for (int i = 0; i < 100; i++)
        std::thread([&seen] {
            CudaDevicePtr<uint8_t> buffer(1 << 16);
            seen.insert(cuda_allocation_stream());
        }).join();

    CHECK(seen.size() == 1);
    CHECK(*seen.begin() != nullptr);
}

TEST_CASE("CudaAllocationStream_DistinctForConcurrentThreads", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    // Neither thread may exit - and hand its stream back - before the other has taken its own.
    std::latch both_have_one(2);
    auto take = [&both_have_one] {
        const cudaStream_t stream = cuda_allocation_stream();
        both_have_one.count_down();
        both_have_one.wait();
        return stream;
    };
    auto a = std::async(std::launch::async, take);
    auto b = std::async(std::launch::async, take);
    const cudaStream_t stream_a = a.get();
    const cudaStream_t stream_b = b.get();

    CHECK(stream_a != nullptr);
    CHECK(stream_b != nullptr);
    CHECK(stream_a != stream_b);
}

// The shadow finder builds its GPU accumulator on a throw-away thread and frees it from another one
// long after. The free is ordered on the allocating thread's stream, so that stream has to outlive
// the thread - including while a later thread has borrowed it.
TEST_CASE("CudaDevicePtr_FreedAfterAllocatingThreadExited", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    auto buffer = std::async(std::launch::async, [] {
        return std::make_unique<CudaDevicePtr<uint8_t>>(1 << 20);
    }).get();
    std::thread([] { CudaDevicePtr<uint8_t> other(1 << 16); }).join();

    REQUIRE(cudaMemset(buffer->get(), 0, 1 << 20) == cudaSuccess);
    buffer.reset();
    CHECK(cudaDeviceSynchronize() == cudaSuccess);
    CHECK(cudaGetLastError() == cudaSuccess);
}

// A pool that cannot serve the request is not an error: the buffer comes from cudaMalloc instead. It
// must then not stay behind as the thread's last error, or the cudaGetLastError() after the next
// kernel launch throws "out of memory" over work that went fine.
TEST_CASE("CudaDevicePtr_HandledPoolFailureLeavesNoError", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    int device = 0;
    REQUIRE(cudaGetDevice(&device) == cudaSuccess);

    cudaMemPoolProps props{};
    props.allocType = cudaMemAllocationTypePinned;
    props.location.type = cudaMemLocationTypeDevice;
    props.location.id = device;
    props.maxSize = 4 << 20;

    cudaMemPool_t previous = nullptr, capped = nullptr;
    REQUIRE(cudaDeviceGetMemPool(&previous, device) == cudaSuccess);
    REQUIRE(cudaMemPoolCreate(&capped, &props) == cudaSuccess);
    REQUIRE(cudaDeviceSetMemPool(device, capped) == cudaSuccess);
    {
        CudaDevicePtr<uint8_t> buffer(64 << 20);
        CHECK(buffer.get() != nullptr);
        CHECK(cudaGetLastError() == cudaSuccess);
    }
    REQUIRE(cudaDeviceSetMemPool(device, previous) == cudaSuccess);
    cudaMemPoolDestroy(capped);
}

// Engines built one after another take the pinned buffers the previous one gave back - also a little
// larger ones, within the size class - while a frame-sized buffer is not kept at all.
TEST_CASE("CudaPinnedPool_ReusedByTheNextEngine", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    cuda_release_pinned_pool();
    void *first = nullptr;
    {
        CudaHostPtr<float> a(1000);
        first = a.get();
    }
    {
        PinnedVector<float> b(1020);
        CHECK(static_cast<void *>(b.data()) == first);
    }
    {
        CudaHostPtr<uint8_t> large(CUDA_PINNED_POOL_MAX_BUFFER + 1);
        large.get()[CUDA_PINNED_POOL_MAX_BUFFER] = 1;
    }
    CHECK(cuda_release_pinned_pool() == 4096);
    CHECK(cuda_release_pinned_pool() == 0);
}

// Within a scope a cuFFT plan is kept once made and lent to one user at a time: two users at once get
// two plans, the next one after them gets one of those back, and a transform on a lent plan - another
// stream, a work area of the caller's - gives the same result bit for bit. When the scope ends, every
// kept plan is destroyed; outside a scope nothing is kept.
TEST_CASE("CudaFFTPlan_KeptAndLentToOneUserAtATime", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    REQUIRE(jfjoch_cufft_plans::kept() == 0);
    auto scope = std::make_unique<CudaFFTPlanScope>();

    const int n = 1000, batch = 7, nc = n / 2 + 1;
    const CudaFFTShape shape{.n = {n}, .idist = n, .odist = nc, .type = CUFFT_R2C, .batch = batch};
    std::vector<float> input(n * batch);
    for (size_t i = 0; i < input.size(); i++)
        input[i] = static_cast<float>((i * 7919) % 1013) - 500.0f;
    CudaDevicePtr<float> d_in(input.size());
    CudaDevicePtr<cufftComplex> d_out(nc * batch);
    REQUIRE(cudaMemcpy(d_in, input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice) == cudaSuccess);

    auto transform = [&](const CudaFFTPlan &plan, cudaStream_t stream) {
        REQUIRE(cufftExecR2C(plan, d_in, d_out) == CUFFT_SUCCESS);
        std::vector<cufftComplex> out(nc * batch);
        REQUIRE(cudaMemcpyAsync(out.data(), d_out, out.size() * sizeof(cufftComplex), cudaMemcpyDeviceToHost,
                                stream) == cudaSuccess);
        REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
        return out;
    };

    CudaStream stream_a, stream_c;
    std::vector<cufftComplex> first;
    cufftHandle handle_a = 0;
    {
        CudaFFTPlan a(shape, stream_a);
        CudaFFTPlan b(shape);
        handle_a = a.get();
        CHECK(handle_a != b.get());
        first = transform(a, stream_a);
    }

    const size_t work = CudaFFTPlan::WorkBytes(shape);
    CudaDevicePtr<char> work_area(std::max<size_t>(work, 1));
    CudaFFTPlan c(shape, stream_c, work_area);
    CHECK(c.get() == handle_a);
    CHECK(c.WorkBytes() == work);
    const std::vector<cufftComplex> again = transform(c, stream_c);
    REQUIRE(again.size() == first.size());
    CHECK(memcmp(again.data(), first.data(), first.size() * sizeof(cufftComplex)) == 0);
    CHECK(jfjoch_cufft_plans::kept() == 1);   // b; a is lent to c

    c = CudaFFTPlan();
    CHECK(jfjoch_cufft_plans::kept() == 2);
    scope.reset();
    CHECK(jfjoch_cufft_plans::kept() == 0);

    CHECK(CudaFFTPlan::WorkBytes(shape) == work);
    { CudaFFTPlan d(shape); }
    CHECK(jfjoch_cufft_plans::kept() == 0);
}

// An FFT indexer that finds the card full waits for the GPU work beside it to give its memory back, and
// is built then: the main line's for the probes, a probe's for the main line's image pass - never for a
// probe, itself included. With nothing to wait for it fails at once, as one that does not fit always has.
TEST_CASE("IndexerThreadPool_FFTIndexerWaitsForGPUWorkBeside", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    if (gpu_is_integrated())
        SKIP("Fills the card, which on an integrated GPU is the host's memory");

    IndexingSettings settings;
    settings.Algorithm(IndexingAlgorithmEnum::FFT).IndexingThreads(1);
    DiffractionExperiment x;
    x.ImportIndexingSettings(settings);
    const std::vector<Coord> recip{{0.1, 0, 0}, {0, 0.1, 0}, {0, 0, 0.1}};

    // Every card, held by the work beside: an indexer worker takes its card round-robin (pin_gpu), so on
    // a machine with several the one it lands on is not the test thread's. Topped up right before each
    // Run, so memory something else gave back in between does not let the indexer fit.
    const int gpus = get_gpu_count();
    std::vector<std::vector<CudaDevicePtr<uint8_t>>> held(gpus);
    const auto fill = [&] {
        for (int d = 0; d < gpus; d++) {
            ScopedGpuPin pin(d);
            for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2)
                while (true) {
                    try {
                        held[d].emplace_back(chunk, CudaAlloc::Synchronous);
                    } catch (const JFJochException &) {
                        break;
                    }
                }
        }
    };
    const auto release = [&] {
        for (int d = 0; d < gpus; d++) {
            ScopedGpuPin pin(d);
            held[d].clear();
        }
    };
    // The pools' indexers go back to the allocation pool when they are destroyed; trim it after them, so
    // the next test finds the cards free.
    {
        IndexerThreadPool probe(settings, IndexerConstruction::OnFirstUse, /*beside=*/true);
        IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse);

        {
            GPUWorkBeside work;
            fill();
            CHECK(probe.Run(x, recip).error != nullptr);
        }
        {
            std::optional<GPUImagePass> image_pass(std::in_place);
            GPUWorkBeside work;
            fill();
            std::thread ends([&] {
                std::this_thread::sleep_for(std::chrono::milliseconds(500));
                release();
                image_pass.reset();
            });
            CHECK(probe.Run(x, recip).error == nullptr);
            ends.join();
        }
        {
            std::optional<GPUWorkBeside> work(std::in_place);
            fill();
            std::thread ends([&] {
                std::this_thread::sleep_for(std::chrono::milliseconds(500));
                release();
                work.reset();
            });
            CHECK(pool.Run(x, recip).error == nullptr);
            ends.join();
        }
    }
    for (int d = 0; d < gpus; d++) {
        ScopedGpuPin pin(d);
        cuda_trim_allocation_pool();
    }
}

// A merge helper's buffer that finds the card full while a probe pass holds the rest of it waits for
// that work to end and its memory to come back (CudaAlloc::SynchronousAfterWorkBeside); a plain
// synchronous one fails at once.
TEST_CASE("CudaDevicePtr_AfterWorkBesideWaitsForItsMemory", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");
    if (gpu_is_integrated())
        SKIP("Fills the card, which on an integrated GPU is the host's memory");
    std::vector<CudaDevicePtr<uint8_t>> held;
    const auto fill = [&] {
        for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2)
            while (true) {
                try {
                    held.emplace_back(chunk, CudaAlloc::Synchronous);
                } catch (const JFJochException &) {
                    break;
                }
            }
    };
    constexpr size_t BYTES = size_t(64) << 20;
    {
        std::optional<GPUWorkBeside> work(std::in_place);
        fill();
        CHECK_THROWS_AS(CudaDevicePtr<uint8_t>(BYTES, CudaAlloc::Synchronous), JFJochException);
        std::thread ends([&] {
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
            held.clear();
            work.reset();
        });
        CHECK_NOTHROW(CudaDevicePtr<uint8_t>(BYTES, CudaAlloc::SynchronousAfterWorkBeside));
        ends.join();
    }
    cuda_trim_allocation_pool();
}

// The free memory waited on is cudaMemGetInfo's on a discrete card. On one that shares the host's memory
// (GB10, Jetson) it also counts the page cache the host hands back, less a tenth of the card for the host:
// never less than cudaMemGetInfo's. Allocations between the readings move them, hence the 1 GB slack.
TEST_CASE("GetGpuFreeMemory_IntegratedCountsThePageCache", "[CUDAMemHelpers]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    constexpr double SLACK = 1e9;
    size_t free_bytes = 0, total_bytes = 0;
    REQUIRE(get_gpu_free_memory(free_bytes, total_bytes));
    size_t cuda_free = 0, cuda_total = 0;
    REQUIRE(cudaMemGetInfo(&cuda_free, &cuda_total) == cudaSuccess);
    CHECK(total_bytes == cuda_total);
    CHECK(free_bytes <= total_bytes);
    CHECK(double(free_bytes) >= double(cuda_free) - SLACK);
    if (!gpu_is_integrated())
        CHECK(double(free_bytes) == Catch::Approx(double(cuda_free)).margin(SLACK));
}

#endif

// A merge that finds the card full waits for the GPU work running beside it: until the last one
// ends, and not forever.
TEST_CASE("GPUWorkBeside_WaitEndsWithTheWork", "[CUDAMemHelpers]") {
    CHECK(wait_for_gpu_work_beside(std::chrono::seconds(0)));

    std::optional<GPUWorkBeside> work;
    work.emplace();
    CHECK_FALSE(wait_for_gpu_work_beside(std::chrono::seconds(0)));

    std::thread ends([&work] {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
        work.reset();
    });
    CHECK(wait_for_gpu_work_beside(std::chrono::seconds(60)));
    ends.join();
}