// 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 <future>
#include <latch>
#include <set>
#include <thread>

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

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

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

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

#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();
}