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

#include <catch2/catch_all.hpp>

#include <algorithm>
#include <atomic>
#include <chrono>
#include <thread>
#include <tuple>
#include <vector>

#include "../common/HostMemory.h"

TEST_CASE("HostMemory_BudgetFromTheTotal", "[HostMemory]") {
    const size_t total = get_host_total_memory();
    REQUIRE(total > 0);
    CHECK(DefaultHostMemoryBudget() == static_cast<size_t>(HOST_MEMORY_BUDGET_FRACTION * static_cast<double>(total)));
}

// As many reservations at once as the pool holds, and one at a time where a single one is larger than it.
TEST_CASE("HostWorkerPool_BoundsWorkersByBytes", "[HostMemory]") {
    HostWorkerPool &pool = HostWorkerPool::Instance();
    for (const auto &[size, bytes, at_most] : std::vector<std::tuple<size_t, size_t, int>>{
             {100, 40, 2}, {100, 100, 1}, {100, 1000, 1}, {1000, 10, 8}}) {
        CAPTURE(size, bytes);
        pool.SetSize(size);
        std::atomic<int> running{0}, most{0}, done{0};
        std::vector<std::thread> workers;
        for (int i = 0; i < 8; i++)
            workers.emplace_back([&] {
                const HostWorkerPool::Reservation reserved(pool, bytes);
                const int now = ++running;
                int seen = most.load();
                while (now > seen && !most.compare_exchange_weak(seen, now)) {}
                std::this_thread::sleep_for(std::chrono::milliseconds(20));
                --running;
                ++done;
            });
        for (auto &w : workers)
            w.join();
        CHECK(done == 8);
        CHECK(most >= 1);
        CHECK(most <= at_most);
    }
    pool.SetSize(0);
    CHECK(pool.Size() == static_cast<size_t>(HOST_WORKER_POOL_FRACTION * static_cast<double>(DefaultHostMemoryBudget())));
}