#pragma once

#include <algorithm>

#include <iosfwd>

#include <memory>

#include <queue>

#include <utility>

#include <vector>

#include "envoy/common/random_generator.h"

#include "envoy/upstream/scheduler.h"

#include "source/common/common/assert.h"

#include "source/common/common/logger.h"

#include "absl/container/flat_hash_map.h"

namespace Envoy {

namespace Upstream {

// Weighted Random Selection Queue (WRSQ) Scheduler

// ------------------------------------------------

// This scheduler keeps a queue for each unique weight among all objects inserted and adds the

// objects to their respective queue based on weight. When performing a pick operation, a queue is

// selected and an object is pulled. Each queue gets its own selection probability which is weighted

// as the sum of all weights of objects contained within. Once a queue is picked, you can simply

// pull from the top and honor the expected selection probability of each object.

//

// Adding an object will cause the scheduler to rebuild internal structures on the first pick that

// follows. This first pick operation will be linear on the number of unique weights among objects

// inserted. Subsequent picks will be logarithmic with the number of unique weights. Adding objects

// is always constant time.

//

// For the case where all object weights are the same, WRSQ behaves identical to vanilla

// round-robin. If all object weights are different, it behaves identical to weighted random

// selection.

//

// NOTE: While the base scheduler interface allows for mutation of object weights with each pick,

// this implementation is not meant for circumstances where the object weights change with each pick

// (like in the least request LB). This scheduler implementation will perform quite poorly if the

// object weights change often.

template <class C>

class WRSQScheduler : public Scheduler<C>, protected Logger::Loggable<Logger::Id::upstream> {

public:

    // Burn through the pre-pick queue.

private:

  using ObjQueue = std::queue<std::weak_ptr<C>>;

  // TODO(tonya11en): We can reduce memory utilization by using an absl::flat_hash_map of QueueInfo

  // with heterogeneous lookup on the weight. This would allow us to save 8 bytes per unique weight.

  using QueueMap = absl::flat_hash_map<double, ObjQueue>;

  // Used to store a queue's weight info necessary to perform the weighted random selection.

  struct QueueInfo {

    double cumulative_weight;

    double weight;

    ObjQueue& q;

  };

  // If needed, such as after object expiry or addition, rebuild the cumulative weights vector.

  // Performs a weighted random selection on the queues containing objects of the same weight.

  // Popping off the top of the queue to pick an object will honor the selection probability based

  // on the weight provided when the object was added.

  // Remove objects from the queue until it's empty or there is an unexpired object at the front. If

  // the queue is purged to empty, it's removed from the queue map and we return true.

        // The chosen queue was purged to empty and removed from the queue map. Try again.

        // The object expired after the purge.

        // The weight has changed for this object, so we must re-add it to the scheduler.

  Random::RandomGenerator& random_;

  // Objects already picked via peekAgain().

  ObjQueue prepick_queue_;

  // A mapping from an object weight to the associated queue.

  QueueMap queue_map_;

  // Stores the necessary information to perform a weighted random selection of the different

  // queues. A cumulative sum is also kept of the total object weights for a queue, which allows for

  // a single random number generation and a binary search to pick a queue.

  std::vector<QueueInfo> cumulative_weights_;

  // Keeps state that determines whether the cumulative weights need to be rebuilt. If any objects

  // contained in a queue change from addition or expiry, it throws off the cumulative weight

  // values. Therefore, they must be recalculated.

  bool rebuild_cumulative_weights_{true};

};

} // namespace Upstream

} // namespace Envoy