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

#include "CUDAWrapper.h"

#include <atomic>
#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, gpu_merges_n = 0;
    // What the merges hold of the card, as they said (GPUMerge::Holds), and how many have not said yet.
    size_t gpu_merge_bytes = 0;
    int gpu_merges_unsized_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;
    }

    // Whether work holding `bytes` fits beside the merges: within 80% of the first card's total, as the
    // GPU budgets of the merge and the model validation count it. With no card there is nothing to share.
    bool fits_beside_merges(size_t bytes) {
        static const size_t card = get_gpu_count() > 0 ? get_gpu_total_memory(0) : 0;
        return card == 0 || (gpu_merges_unsized_n == 0 && bytes + gpu_merge_bytes <= card / 5 * 4);
    }

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

bool gpu_takes_turns() {
    return take_turns();
}

GPUWorkBeside::GPUWorkBeside(bool after_merges, size_t bytes) {
    std::unique_lock lock(gpu_work_m);
    gpu_work_changed.wait(lock, [after_merges, bytes] {
        return (!take_turns() || (gpu_image_passes_n == 0 && gpu_work_beside_n == 0))
               && (!after_merges || gpu_merges_n == 0 || fits_beside_merges(bytes));
    });
    ++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);
}

GPUMerge::GPUMerge() {
    std::lock_guard lock(gpu_work_m);
    ++gpu_merges_n;
    ++gpu_merges_unsized_n;
}

GPUMerge::~GPUMerge() {
    {
        std::lock_guard lock(gpu_work_m);
        if (bytes_)
            gpu_merge_bytes -= *bytes_;
        else
            --gpu_merges_unsized_n;
    }
    remove(gpu_merges_n);
}

void GPUMerge::Holds(size_t bytes) {
    {
        std::lock_guard lock(gpu_work_m);
        if (bytes_)
            gpu_merge_bytes -= *bytes_;
        else
            --gpu_merges_unsized_n;
        gpu_merge_bytes += bytes;
        bytes_ = bytes;
    }
    gpu_work_changed.notify_all();
}

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

namespace {
    std::atomic<bool> gpus_hidden{false};
}

void hide_gpus() {
    gpus_hidden = true;
}

bool gpus_are_hidden() {
    return gpus_hidden;
}

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