#ifndef SLED_FUTURES_WHEN_ALL_H
#define SLED_FUTURES_WHEN_ALL_H

#pragma once
#include "future.h"
#include "sled/meta/type_traits.h"
#include "sled/synchronization/mutex.h"

namespace sled {

namespace futures {

struct WhenAllState {
    sled::Mutex mutex;
    sled::ConditionVariable cv;
    std::atomic<int> count{0};
    std::atomic<bool> has_failed{false};
};

template<typename T, typename FailureT>
void
WhenAllImpl(std::weak_ptr<WhenAllState> &weak_state, Future<T, FailureT> &future)
{
    {
        auto state = weak_state.lock();
        if (!state) { return; }
        state->count++;
    }
    future.OnSuccess([weak_state](const T &) {
        auto state = weak_state.lock();
        if (!state) { return; }
        if (state->count.fetch_sub(1) == 0) {
            sled::MutexLock lock(&state->mutex);
            state->cv.NotifyAll();
        }
    });
    future.OnFailure([weak_state](const FailureT &) {
        auto state = weak_state.lock();
        if (!state) { return; }
        state->has_failed = true;
        sled::MutexLock lock(&state->mutex);
        state->cv.NotifyAll();
    });
}

template<typename T, typename FailureT, typename... Args>
void
WhenAllImpl(std::weak_ptr<WhenAllState> &weak_state, Future<T, FailureT> &future, Args &&...futures)
{
    {
        auto state = weak_state.lock();
        if (!state) { return; }
        state->count++;
    }

    future.OnSuccess([weak_state](const T &) {
        auto state = weak_state.lock();
        if (!state) { return; }
        if (state->count.fetch_sub(1) == 1) {
            sled::MutexLock lock(&state->mutex);
            state->cv.NotifyAll();
        }
    });

    future.OnFailure([weak_state](const FailureT &) {
        auto state = weak_state.lock();
        if (!state) { return; }
        state->has_failed = true;
        sled::MutexLock lock(&state->mutex);
        state->cv.NotifyAll();
    });

    WhenAllImpl(weak_state, std::forward<Args>(futures)...);
}

}// namespace futures

template<typename T, typename FailureT, typename... Args>
bool
WhenAll(Future<T, FailureT> &future, Args &&...futures)
{
    auto state                                      = std::make_shared<futures::WhenAllState>();
    std::weak_ptr<futures::WhenAllState> weak_state = state;
    futures::WhenAllImpl(weak_state, future, std::forward<Args>(futures)...);

    {
        sled::MutexLock lock(&state->mutex);
        state->cv.Wait(lock, [state] { return state->count.load() == 0 || state->has_failed.load(); });
    }

    return !state->has_failed.load();
}

template<typename ContainerT>
bool
WhenAll(ContainerT container)
{
    auto state                                      = std::make_shared<futures::WhenAllState>();
    std::weak_ptr<futures::WhenAllState> weak_state = state;
    for (auto &f : container) { futures::WhenAllImpl(weak_state, f); }
    {
        sled::MutexLock lock(&state->mutex);
        state->cv.Wait(lock, [state] { return state->count.load() == 0 || state->has_failed.load(); });
    }
    return !state->has_failed.load();
}

template<typename IteratorT>
bool
WhenAll(IteratorT begin, IteratorT end)
{
    auto state                                      = std::make_shared<futures::WhenAllState>();
    std::weak_ptr<futures::WhenAllState> weak_state = state;
    for (auto it = begin; it != end; ++it) { futures::WhenAllImpl(weak_state, *it); }
    {
        sled::MutexLock lock(&state->mutex);
        state->cv.Wait(lock, [state] { return state->count.load() == 0 || state->has_failed.load(); });
    }
    return !state->has_failed.load();
}

}// namespace sled

#endif// SLED_FUTURES_WHEN_ALL_H
