#pragma once
///@file

#include <functional>
#include <limits>
#include <list>
#include <memory>
#include <cassert>

#include "nix/util/fun.hh"
#include "nix/util/sync.hh"
#include "nix/util/ref.hh"

namespace nix {

/**
 * This template class implements a simple pool manager of resources
 * of some type R, such as database connections. It is used as
 * follows:
 *
 *   class Connection { ... };
 *
 *   Pool<Connection> pool;
 *
 *   {
 *     auto conn(pool.get());
 *     conn->exec("select ...");
 *   }
 *
 * Here, the Connection object referenced by ‘conn’ is automatically
 * returned to the pool when ‘conn’ goes out of scope.
 */
template<class R>
class Pool
{
public:

    /**
     * A function that produces new instances of R on demand.
     */
    typedef fun<ref<R>()> Factory;

    /**
     * A function that checks whether an instance of R is still
     * usable. Unusable instances are removed from the pool.
     */
    typedef fun<bool(const ref<R> &)> Validator;

private:

    Factory factory;
    Validator validator;

    struct State
    {
        size_t inUse = 0;
        size_t max;
        std::vector<ref<R>> idle;
    };

    Sync<State> state;

    std::condition_variable wakeup;

public:

    Pool(
        size_t max = std::numeric_limits<size_t>::max(),
        const Factory & factory = []() { return make_ref<R>(); },
        const Validator & validator = [](ref<R> r) { return true; })
        : factory(factory)
        , validator(validator)
    {
        auto state_(state.lock());
        state_->max = max;
    }

    void incCapacity()
    {
        auto state_(state.lock());
        state_->max++;
        /* we could wakeup here, but this is only used when we're
         * about to nest Pool usages, and we want to save the slot for
         * the nested use if we can
         */
    }

    void decCapacity()
    {
        auto state_(state.lock());
        state_->max--;
    }

    ~Pool()
    {
        auto state_(state.lock());
        assert(!state_->inUse);
        state_->max = 0;
        state_->idle.clear();
    }

    class Handle
    {
    private:
        Pool & pool;
        std::shared_ptr<R> r;
        bool bad = false;

        friend Pool;

        Handle(Pool & pool, std::shared_ptr<R> r)
            : pool(pool)
            , r(r)
        {
        }

    public:
        // NOTE: Copying std::shared_ptr and calling a .reset() on it is always noexcept.
        Handle(Handle && h) noexcept
            : pool(h.pool)
            , r(h.r)
        {
            static_assert(noexcept(h.r.reset()));
            static_assert(noexcept(std::shared_ptr(h.r)));
            h.r.reset();
        }

        Handle(const Handle & l) = delete;

        ~Handle()
        {
            if (!r)
                return;
            {
                auto state_(pool.state.lock());
                if (!bad)
                    state_->idle.push_back(ref<R>(r));
                assert(state_->inUse);
                state_->inUse--;
            }
            pool.wakeup.notify_one();
        }

        R * operator->()
        {
            return &*r;
        }

        R & operator*()
        {
            return *r;
        }

        void markBad()
        {
            bad = true;
        }
    };

    Handle get()
    {
        {
            auto state_(state.lock());

            /* If we're over the maximum number of instance, we need
               to wait until a slot becomes available. */
            while (state_->idle.empty() && state_->inUse >= state_->max)
                state_.wait(wakeup);

            while (!state_->idle.empty()) {
                auto p = state_->idle.back();
                state_->idle.pop_back();
                if (validator(p)) {
                    state_->inUse++;
                    return Handle(*this, p);
                }
            }

            state_->inUse++;
        }

        /* We need to create a new instance. Because that might take a
           while, we don't hold the lock in the meantime. */
        try {
            Handle h(*this, factory());
            return h;
        } catch (...) {
            auto state_(state.lock());
            state_->inUse--;
            wakeup.notify_one();
            throw;
        }
    }

    size_t count()
    {
        auto state_(state.lock());
        return state_->idle.size() + state_->inUse;
    }

    size_t capacity()
    {
        return state.lock()->max;
    }

    void flushBad()
    {
        auto state_(state.lock());
        std::vector<ref<R>> left;
        for (auto & p : state_->idle)
            if (validator(p))
                left.push_back(p);
        std::swap(state_->idle, left);
    }

    std::vector<ref<R>> clear()
    {
        auto state_(state.lock());
        return std::exchange(state_->idle, {});
    }
};

} // namespace nix
