// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license.

#pragma once
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <mutex>
#include <queue>
#include <thread>
#include <type_traits>
#include <unordered_set>

namespace diskann
{

template <typename T> class ConcurrentQueue
{
    typedef std::chrono::microseconds chrono_us_t;
    typedef std::unique_lock<std::mutex> mutex_locker;

    std::queue<T> q;
    std::mutex mut;
    std::mutex push_mut;
    std::mutex pop_mut;
    std::condition_variable push_cv;
    std::condition_variable pop_cv;
    T null_T;

  public:
    ConcurrentQueue()
    {
    }

    ConcurrentQueue(T nullT)
    {
        this->null_T = nullT;
    }

    ~ConcurrentQueue()
    {
        this->push_cv.notify_all();
        this->pop_cv.notify_all();
    }

    // queue stats
    uint64_t size()
    {
        mutex_locker lk(this->mut);
        uint64_t ret = q.size();
        lk.unlock();
        return ret;
    }

    bool empty()
    {
        return (this->size() == 0);
    }

    // PUSH BACK
    void push(T &new_val)
    {
        mutex_locker lk(this->mut);
        this->q.push(new_val);
        lk.unlock();
    }

    template <class Iterator> void insert(Iterator iter_begin, Iterator iter_end)
    {
        mutex_locker lk(this->mut);
        for (Iterator it = iter_begin; it != iter_end; it++)
        {
            this->q.push(*it);
        }
        lk.unlock();
    }

    // POP FRONT
    T pop()
    {
        mutex_locker lk(this->mut);
        if (this->q.empty())
        {
            lk.unlock();
            return this->null_T;
        }
        else
        {
            T ret = this->q.front();
            this->q.pop();
            // diskann::cout << "thread_id: " << std::this_thread::get_id() <<
            // ", ctx: "
            // << ret.ctx << "\n";
            lk.unlock();
            return ret;
        }
    }

    // register for notifications
    void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10})
    {
        mutex_locker lk(this->push_mut);
        this->push_cv.wait_for(lk, wait_time);
        lk.unlock();
    }

    void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10})
    {
        mutex_locker lk(this->pop_mut);
        this->pop_cv.wait_for(lk, wait_time);
        lk.unlock();
    }

    // just notify functions
    void push_notify_one()
    {
        this->push_cv.notify_one();
    }
    void push_notify_all()
    {
        this->push_cv.notify_all();
    }
    void pop_notify_one()
    {
        this->pop_cv.notify_one();
    }
    void pop_notify_all()
    {
        this->pop_cv.notify_all();
    }
};
} // namespace diskann