#pragma once

#include <optional>
#include <algorithm>
#include <atomic>
#include <condition_variable>
#include <deque>
#include <memory>
#include <mutex>
#include <tuple>
#include <utility>
#include <vector>

#include "optionalVersion.h"

struct MultiQueueWaiter;

struct BaseThreadQueue
{
    virtual ~BaseThreadQueue() = default;

    virtual bool IsEmpty() = 0;

    std::shared_ptr<MultiQueueWaiter> waiter;
};

// std::lock accepts two or more arguments. We define an overload for one
// argument.
namespace std
{
template<typename Lockable>
void lock(Lockable& l)
{
    l.lock();
}
} // namespace std

template<typename... Queue>
struct MultiQueueLock
{
    MultiQueueLock(Queue... lockable) : tuple_ {lockable...}
    {
        lock();
    }
    ~MultiQueueLock()
    {
        unlock();
    }
    void lock()
    {
        lock_impl(typename std::index_sequence_for<Queue...> {});
    }
    void unlock()
    {
        unlock_impl(typename std::index_sequence_for<Queue...> {});
    }

private:
    template<size_t... Is>
    void lock_impl(std::index_sequence<Is...>)
    {
        std::lock(std::get<Is>(tuple_)->mutex...);
    }

    template<size_t... Is>
    void unlock_impl(std::index_sequence<Is...>)
    {
        (void)std::initializer_list<int> {(std::get<Is>(tuple_)->mutex.unlock(), 0)...};
    }

    std::tuple<Queue...> tuple_;
};

struct MultiQueueWaiter
{
    static bool HasState(std::initializer_list<BaseThreadQueue*> queues);

    bool ValidateWaiter(std::initializer_list<BaseThreadQueue*> queues);

    template<typename... BaseThreadQueue>
    bool Wait(std::atomic<bool>& quit, BaseThreadQueue... queues)
    {
        MultiQueueLock<BaseThreadQueue...> l(queues...);
        while (!quit.load(std::memory_order_relaxed))
        {
            if (HasState({queues...}))
            {
                return false;
            }
            cv.wait(l);
        }
        return true;
    }
    template<typename... BaseThreadQueue>
    void WaitUntil(std::chrono::steady_clock::time_point t, BaseThreadQueue... queues)
    {
        MultiQueueLock<BaseThreadQueue...> l(queues...);
        if (!HasState({queues...}))
        {
            cv.wait_until(l, t);
        }
    }
    template<typename... BaseThreadQueue>
    void Wait(BaseThreadQueue... queues)
    {
        assert(ValidateWaiter({queues...}));

        MultiQueueLock<BaseThreadQueue...> l(queues...);
        while (!HasState({queues...}))
        {
            cv.wait(l);
        }
    }

    std::condition_variable_any cv;
};

// A threadsafe-queue. http://stackoverflow.com/a/16075550
template<class T>
struct ThreadedQueue : public BaseThreadQueue
{
public:
    ThreadedQueue() : ThreadedQueue(std::make_shared<MultiQueueWaiter>())
    {
    }

    explicit ThreadedQueue(std::shared_ptr<MultiQueueWaiter> waiter) : total_count_(0)
    {
        this->waiter = waiter;
    }

    // Returns the number of elements in the queue. This is lock-free.
    size_t Size() const
    {
        return total_count_;
    }

    // Add an element to the queue.
    void Enqueue(T&& t, bool priority)
    {
        {
            std::lock_guard<std::mutex> lock(mutex);
            if (priority)
            {
                priority_.push_back(std::move(t));
            }
            else
            {
                queue_.push_back(std::move(t));
            }
            ++total_count_;
        }
        waiter->cv.notify_one();
    }

    // Add a set of elements to the queue.
    void EnqueueAll(std::vector<T>&& elements, bool priority)
    {
        if (elements.empty())
        {
            return;
        }

        {
            std::lock_guard<std::mutex> lock(mutex);
            total_count_ += elements.size();
            for (T& element : elements)
            {
                if (priority)
                {
                    priority_.push_back(std::move(element));
                }
                else
                {
                    queue_.push_back(std::move(element));
                }
            }
            elements.clear();
        }

        waiter->cv.notify_all();
    }

    // Returns true if the queue is empty. This is lock-free.
    bool IsEmpty()
    {
        return total_count_ == 0;
    }

    // Get the first element from the queue. Blocks until one is available.
    T Dequeue()
    {
        std::unique_lock<std::mutex> lock(mutex);
        waiter->cv.wait(lock, [&]() { return !priority_.empty() || !queue_.empty(); });

        auto execute = [&](std::deque<T>* q)
        {
            auto val = std::move(q->front());
            q->pop_front();
            --total_count_;
            return std::move(val);
        };
        if (!priority_.empty())
        {
            return execute(&priority_);
        }
        return execute(&queue_);
    }

    // Get the first element from the queue without blocking. Returns a null
    // value if the queue is empty.
    optional<T> TryDequeue(bool priority)
    {
        std::lock_guard<std::mutex> lock(mutex);

        auto pop = [&](std::deque<T>* q)
        {
            auto val = std::move(q->front());
            q->pop_front();
            --total_count_;
            return std::move(val);
        };

        auto get_result = [&](std::deque<T>* first, std::deque<T>* second) -> optional<T>
        {
            if (!first->empty())
            {
                return pop(first);
            }
            if (!second->empty())
            {
                return pop(second);
            }
            return {};
        };

        if (priority)
        {
            return get_result(&priority_, &queue_);
        }
        return get_result(&queue_, &priority_);
    }
    // Return all elements in the queue.
    std::vector<T> DequeueAll()
    {
        std::lock_guard<std::mutex> lock(mutex);

        total_count_ = 0;

        std::vector<T> result;
        result.reserve(priority_.size() + queue_.size());
        while (!priority_.empty())
        {
            result.emplace_back(std::move(priority_.front()));
            priority_.pop_front();
        }
        while (!queue_.empty())
        {
            result.emplace_back(std::move(queue_.front()));
            queue_.pop_front();
        }

        return result;
    }
    std::vector<T> TryDequeueSome(size_t num)
    {
        std::lock_guard<std::mutex> lock(mutex);

        std::vector<T> result;
        num = std::min(num, priority_.size() + queue_.size());
        total_count_ -= num;
        result.reserve(num);
        while (num)
        {
            if (!priority_.empty())
            {
                result.emplace_back(std::move(priority_.front()));
                priority_.pop_front();
            }
            else
            {
                break;
            }
            num -= 1;
        }
        while (num)
        {
            if (!queue_.empty())
            {
                result.emplace_back(std::move(queue_.front()));
                queue_.pop_front();
            }
            else
            {
                break;
            }
            num -= 1;
        }
        return result;
    }
    template<typename Fn>
    void Iterate(Fn fn)
    {
        std::lock_guard<std::mutex> lock(mutex);
        for (auto& entry : priority_)
        {
            fn(entry);
        }
        for (auto& entry : queue_)
        {
            fn(entry);
        }
    }

    mutable std::mutex mutex;

private:
    std::atomic<int> total_count_;
    std::deque<T> priority_;
    std::deque<T> queue_;
};