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

#pragma once

#include <algorithm>
#include <atomic>
#include <condition_variable>
#include <exception>
#include <functional>
#include <mutex>
#include <thread>
#include <vector>

// Two shapes of "run this over a range on several threads", used by the analysis code. Both take the
// worker count from the caller rather than asking the hardware, so a run that was told how many
// threads to use keeps to it.

// How many workers a pass over `n` cheap items should use: enough that each gets at least
// `min_per_thread` of them, and never more than the caller was given. Entering a parallel region below
// costs a condition-variable wakeup rather than a thread, but a pass over a few thousand items still
// finishes before the last worker has woken up, so the floor stays. That is not a large machine's
// problem: it is what makes the same code right on an 8-core laptop, a 16-core desktop and a
// two-socket node, none of which should be handed 48 chunks of a few thousand items.
inline size_t ThreadsForWork(size_t n, size_t nthreads, size_t min_per_thread = 32768) {
    if (nthreads <= 1 || n == 0) return 1;
    return std::clamp<size_t>(n / min_per_thread, 1, nthreads);
}

namespace parallel_detail {

// One parallel region: `ntasks` invocations of one body, each claimed and run by exactly one thread.
struct Region {
    const std::function<void(int)> *body;
    int ntasks;
    int next = 0;                  // next unclaimed task, guarded by the pool's mutex
    std::atomic<int> done{0};
    std::exception_ptr error;      // the first task to throw, guarded by the pool's mutex
};

// The threads both helpers below run on: started once, parked on a condition variable between regions.
// The analysis code enters tens of thousands of parallel regions in a run - several from inside
// iterative fits - and a thread per chunk per call was tens of microseconds of clone and first-touch
// each time, paid again on every iteration.
class Pool {
public:
    static Pool &Instance() {
        static Pool pool;
        return pool;
    }

    // Run body(0) ... body(ntasks-1), each exactly once, on the pool's threads and on this one, and
    // return when they have all finished. An exception from a task is rethrown here, as it was when
    // these were futures.
    void Run(int ntasks, const std::function<void(int)> &body) {
        Region region{&body, ntasks};
        {
            std::lock_guard lock(mutex_);
            regions_.push_back(&region);
        }
        // One wakeup per task the caller is not going to run itself. notify_all instead would wake every
        // parked thread for a region of two chunks - on a node with 192 of them that is 190 threads woken
        // to find nothing and go back to sleep, per region, tens of thousands of times in a run.
        for (int i = 1; i < ntasks; i++)
            work_.notify_one();
        // The calling thread is one of the hands and keeps claiming until nothing of this region is
        // left unclaimed. That is what makes a region entered from inside a region safe: a region never
        // waits for a pool thread to come free, only for tasks that are already running.
        for (int i = Claim(region); i >= 0; i = Claim(region))
            Execute(region, i);
        std::unique_lock lock(mutex_);
        done_.wait(lock, [&] { return region.done.load() == ntasks; });
        if (region.error)
            std::rethrow_exception(region.error);
    }

private:
    Pool() {
        const unsigned n = std::max(1u, std::thread::hardware_concurrency());
        threads_.reserve(n);
        for (unsigned i = 0; i < n; i++)
            threads_.emplace_back([this] { Work(); });
    }

    ~Pool() {
        {
            std::lock_guard lock(mutex_);
            stop_ = true;
        }
        work_.notify_all();
        for (auto &t : threads_)
            t.join();
    }

    // Claim one of `region`'s tasks, or -1 if it has none left. A region leaves the list the moment its
    // last task is claimed, so no pool thread can look at it after Run returns and it goes out of scope.
    int Claim(Region &region) {
        std::lock_guard lock(mutex_);
        if (region.next == region.ntasks)
            return -1;
        const int i = region.next++;
        if (region.next == region.ntasks)
            regions_.erase(std::find(regions_.begin(), regions_.end(), &region));
        return i;
    }

    // Claim one task of the newest region that still has one, called with the mutex held. Newest first,
    // so an inner region is finished before more hands go to the outer one waiting on it.
    Region *ClaimAny(int &i) {
        for (size_t at = regions_.size(); at-- > 0;) {
            Region *region = regions_[at];
            if (region->next < region->ntasks) {
                i = region->next++;
                if (region->next == region->ntasks)
                    regions_.erase(regions_.begin() + at);
                return region;
            }
        }
        return nullptr;
    }

    // Run one task and count it. The last one to finish wakes whoever is waiting in Run; it reads ntasks
    // before the count, because after the count the region may already be gone.
    void Execute(Region &region, int i) {
        try {
            (*region.body)(i);
        } catch (...) {
            std::lock_guard lock(mutex_);
            if (!region.error)
                region.error = std::current_exception();
        }
        const int ntasks = region.ntasks;
        if (region.done.fetch_add(1) + 1 == ntasks) {
            std::lock_guard lock(mutex_);
            done_.notify_all();
        }
    }

    void Work() {
        std::unique_lock lock(mutex_);
        for (;;) {
            int i = 0;
            Region *region = ClaimAny(i);
            if (!region) {
                if (stop_)
                    return;
                work_.wait(lock);
                continue;
            }
            lock.unlock();
            Execute(*region, i);
            lock.lock();
        }
    }

    std::mutex mutex_;
    std::condition_variable work_, done_;
    std::vector<Region *> regions_;
    std::vector<std::thread> threads_;
    bool stop_ = false;
};

}  // namespace parallel_detail

// Chunked: each worker gets one contiguous [lo, hi) range and there is no per-item synchronisation.
// Right for millions of cheap uniform items - the CPU stand-in for a flat CUDA grid-stride kernel.
// The split is fixed and deterministic, so a pass whose per-element work is independent gives the
// same answer as the serial loop, bit for bit.
template <typename Fn>
void ParallelChunks(int n, size_t nthreads, Fn fn) {
    if (n <= 0) return;
    const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
    if (nt == 1) { fn(0, n); return; }
    const int chunk = (n + nt - 1) / nt;
    const std::function<void(int)> body = [&fn, n, chunk](int t) {
        const int lo = t * chunk, hi = std::min(n, lo + chunk);
        if (lo < hi) fn(lo, hi);
    };
    parallel_detail::Pool::Instance().Run(nt, body);
}

// Work-stealing per item, off a shared atomic counter: one atomic per item, so use it only where the
// per-item work is heavy and uneven (per-frame fits, per-ring selections) and the atomic amortises.
// For millions of tiny uniform items a per-item atomic is pure contention - use ParallelChunks.
template <typename Fn>
void ParallelFor(int n, size_t nthreads, Fn fn) {
    if (n <= 0) return;
    if (nthreads <= 1 || n == 1) {
        for (int i = 0; i < n; i++) fn(i);
        return;
    }
    const int nt = static_cast<int>(std::min(nthreads, static_cast<size_t>(n)));
    std::atomic<int> next = 0;
    const std::function<void(int)> body = [&fn, &next, n](int) {
        for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
            fn(i);
    };
    parallel_detail::Pool::Instance().Run(nt, body);
}