#pragma once

#include <balancer/kernel/balancer/algorithm.h>
#include <util/generic/hash_set.h>
#include <util/random/fast.h>

#include <functional>

namespace NSrvKernel {

class TAlgorithmWithRemovals : public IAlgorithm {
public:
    explicit TAlgorithmWithRemovals(IWorkerCtl* process);
    void RemoveSelected(IBackend* backend) noexcept override;
    void Reset() noexcept override;
    bool IsRemoved(IBackend* backend) const noexcept;
protected:
    THashSet<IBackend*> Excluded_;
};

using TWrrRng = TReallyFastRng32;

class TBaseWrrAlgorithm : public TAlgorithmWithRemovals {
public:
    TBaseWrrAlgorithm(IWorkerCtl* process, TWrrRng* rng);
    IBackend* SelectBackend(const TBackendsGroupWeights& group, std::function<IBackend*(size_t, double)> resolve = {}) const noexcept;
    size_t Round(const TBackendsGroupWeights& group) const noexcept;
    size_t Round(const TBackendsGroupWeights& group, double diceRoll) const noexcept;
private:
    virtual IBackend* Resolve(size_t index, double diceRoll) const noexcept = 0;
    size_t FindInGroup(const TBackendsGroupWeights& group, double diceRoll) const noexcept;
    IBackend* Fallback(const TBackendsGroupWeights& group, std::function<IBackend*(size_t, double)> resolve, TWrrRng* rng) const noexcept;
    TWrrRng* Rng_;
};

}  // namespace NSrvKernel

