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

#include "CUDAWrapper.h"

#include <condition_variable>
#include <mutex>

// Build-independent: the CUDA build gets get_gpu_names() from CUDAWrapper.cu, the CPU-only build from
// the stub below, and this collapses whichever list came back. Four identical cards read better as
// "4x <name>" than as the same name four times, and a mixed machine keeps one group per model.
std::string get_gpu_description() {
    const auto names = get_gpu_names();

    std::string out;
    for (size_t i = 0; i < names.size();) {
        size_t n = 1;
        while (i + n < names.size() && names[i + n] == names[i])
            n++;
        if (!out.empty())
            out += ", ";
        if (n > 1)
            out += std::to_string(n) + "x ";
        out += names[i];
        i += n;
    }
    return out;
}

namespace {
    // How many probe passes (GPUWorkBeside) and main-line image passes (GPUImagePass) are running, for a
    // thread to wait until none of one kind is.
    std::mutex gpu_work_m;
    std::condition_variable gpu_work_changed;
    int gpu_work_beside_n = 0, gpu_image_passes_n = 0;

    // A card too small for a probe pass beside the main line's image pass, or beside another probe pass:
    // measured on a 16 Mpx detector, two of them together held 8.6 GB. On such a card a probe pass waits,
    // before it starts, for every other one and for the image passes to end, and an image pass for the
    // probe passes - so they take turns, and only when each runs depends on the card, never what it
    // computes. None waits once it is running, so nothing waits in a circle.
    constexpr size_t TAKE_TURNS_BELOW_BYTES = 12'000'000'000ull;

    bool take_turns() {
        static const bool small = get_gpu_count() > 0 && get_gpu_total_memory(0) < TAKE_TURNS_BELOW_BYTES;
        return small;
    }

    void remove(int &n) {
        {
            std::lock_guard lock(gpu_work_m);
            --n;
        }
        gpu_work_changed.notify_all();
    }
}

GPUWorkBeside::GPUWorkBeside() {
    std::unique_lock lock(gpu_work_m);
    if (take_turns())
        gpu_work_changed.wait(lock, [] { return gpu_image_passes_n == 0 && gpu_work_beside_n == 0; });
    ++gpu_work_beside_n;
}

GPUWorkBeside::~GPUWorkBeside() {
    remove(gpu_work_beside_n);
}

bool wait_for_gpu_work_beside(std::chrono::seconds timeout) {
    std::unique_lock lock(gpu_work_m);
    return gpu_work_changed.wait_for(lock, timeout, [] { return gpu_work_beside_n == 0; });
}

GPUImagePass::GPUImagePass() {
    std::unique_lock lock(gpu_work_m);
    if (take_turns())
        gpu_work_changed.wait(lock, [] { return gpu_work_beside_n == 0; });
    ++gpu_image_passes_n;
}

GPUImagePass::~GPUImagePass() {
    remove(gpu_image_passes_n);
}

void wait_for_gpu_image_passes() {
    std::unique_lock lock(gpu_work_m);
    gpu_work_changed.wait(lock, [] { return gpu_image_passes_n == 0; });
}

#ifndef JFJOCH_USE_CUDA

int32_t get_gpu_count() {
    return 0;
}

void hide_gpus() {}

std::vector<std::string> get_gpu_names() {
    return {};
}

size_t get_gpu_total_memory(int32_t) {
    return 0;
}

bool get_gpu_free_memory(size_t &, size_t &) {
    return false;
}

void set_gpu(int32_t dev_id) {}

bool gpu_direct_host_read() {
    return false;
}

void pin_gpu() {}

void pin_gpu(int32_t dev_id) {}

ScopedGpuPin::ScopedGpuPin(int32_t) {}

ScopedGpuPin::~ScopedGpuPin() {}

void enable_gpu_numa_binding() {}

std::vector<int> get_gpu_numa_nodes() { return {}; }

void set_gpu_blocking_sync() {}

void create_gpu_contexts(int32_t) {}

void release_gpu_contexts() {}

void cuda_clear_error() {}

void cuda_throw_if_context_lost() {}

#endif