#pragma once

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

#include "libutil/ref.hh"
#include "libutil/sync.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 std::function<ref<R>()> Factory;

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

 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:
    Handle(Handle&& h) : pool(h.pool), r(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);
  }
};

}  // namespace nix