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

#ifndef JUNGFRAUJOCH_THREADSAFEFIFO_H
#define JUNGFRAUJOCH_THREADSAFEFIFO_H

#include <queue>
#include <mutex>
#include <condition_variable>
#include <set>

template <class T> class ThreadSafeFIFO {
    std::queue<T> queue;
    std::condition_variable c_empty, c_full;
    mutable std::mutex m;
    const size_t max_size;
    size_t max_utilization;
    size_t utilization;
public:
    explicit ThreadSafeFIFO(size_t in_max_size = UINT32_MAX) : max_size(in_max_size), max_utilization(0), utilization(0) {}

    void Clear() {
        std::unique_lock ul(m);
        queue = {};
        utilization = 0;
        max_utilization = 0;
    }

    bool Put(T val) {
        std::unique_lock ul(m);
        if (queue.size() < max_size) {
            queue.push(val);
            c_empty.notify_one();
            utilization++;
            if (utilization > max_utilization)
                max_utilization = utilization;
            return true;
        } else
            return false;
    };

    void PutBlocking(T val) {
        std::unique_lock ul(m);
        c_full.wait(ul, [&]{return queue.size() < max_size;});
        queue.push(val);
        utilization++;
        if (utilization > max_utilization)
            max_utilization = utilization;
        c_empty.notify_one();
    };

    int Get(T &val) {
        std::unique_lock ul(m);

        if (queue.empty())
            return 0;
        else {
            val = queue.front();
            queue.pop();
            c_full.notify_one();
            utilization--;
            return 1;
        }
    }

    T GetBlocking() {
        std::unique_lock ul(m);
        c_empty.wait(ul, [&]{return !queue.empty();});
        T tmp = queue.front();
        queue.pop();
        c_full.notify_one();
        utilization--;
        return tmp;
    };

    int GetTimeout(T &val, std::chrono::microseconds timeout) {
        std::unique_lock ul(m);
        if (queue.empty())
            c_empty.wait_for(ul, timeout, [&]{return !queue.empty();});

        if (queue.empty())
            return 0;
        else {
            val = queue.front();
            queue.pop();
            c_full.notify_one();
            utilization--;
            return 1;
        }
    }

    [[nodiscard]] size_t Size() const {
        return queue.size();
    }

    void ClearMaxUtilization() {
        std::unique_lock ul(m);
        max_utilization = utilization;
    }

    [[nodiscard]] size_t GetMaxUtilization() const {
        std::unique_lock ul(m);
        return max_utilization;
    }

    [[nodiscard]] size_t GetCurrentUtilization() const {
        std::unique_lock ul(m);
        return utilization;
    }
};

template <class T> class ThreadSafeSet {
    std::set<T> set;
    std::condition_variable c;
    std::mutex m;
public:
    void Clear() {
        std::unique_lock ul(m);
        set = {};
    }
    void Put(T val) {
        std::unique_lock ul(m);
        set.insert(val);
        c.notify_one();
    };

    int Get(T &val) {
        std::unique_lock ul(m);

        if (set.empty()) return 0;
        else {
            auto iter = set.begin();
            val = *iter;
            set.erase(iter);
            return 1;
        }
    }

    T GetBlocking() {
        std::unique_lock ul(m);
        c.wait(ul, [&]{return !set.empty();});
        auto iter = set.begin();
        T tmp = *iter;
        set.erase(iter);
        return tmp;
    };

    size_t Size() const {
        return set.size();
    }
};

#endif //JUNGFRAUJOCH_THREADSAFEFIFO_H