#include "back_off.h"

#include <balancer/kernel/balancer/algorithm.h>
#include <balancer/kernel/balancing/backend.h>
#include <balancer/modules/balancer/policies/policy_factory.h>

#include <balancer/kernel/helpers/common_parsers.h>

namespace NSrvKernel::NBackOff {

IBackend* TBackOffPolicy::BackoffDoNext(IAlgorithm* algo, double probability, bool fastAttempt) noexcept {
    const double current = RandomNumber<double>();
    if (current < probability) {
        return Slave_->Next(algo, fastAttempt);
    }

    return nullptr;
}

IBackend* TBackOffPolicy::Next(IAlgorithm* algo, bool fastAttempt) noexcept {
    if (Attempt_ == 0 || Probabilities_.empty()) {
        ++Attempt_;
        return Slave_->Next(algo, fastAttempt);
    } else if ((Attempt_ - 1) < Probabilities_.size()) {
        const size_t idx = Attempt_ - 1;
        ++Attempt_;
        return BackoffDoNext(algo, Probabilities_[idx], fastAttempt);
    } else {
        return BackoffDoNext(algo, Probabilities_.back(), fastAttempt);
    }
}


POLICY_FACTORY_BASE(backoff_policy, TPolicyFactoryWithSlave), public TModuleParams {
    TPolicyFactory(const TModuleParams& mp)
        : TModuleParams(mp)
    {
        Config->ForEach(this);
        CheckConfiguration();
    }

    THolder<IPolicy> ConstructPolicy(const TStepParams& params) noexcept override {
        return MakeHolder<TBackOffPolicy>(Probabilities_, TPolicyBase::ConstructPolicy(params));
    }

private:
    START_PARSE {
        if (key == "probabilities") {
            ParseMap(value->AsSubConfig(), [&m = Probabilities_](const auto&, auto* value) {
                const double probability = FromString<double>(value->AsString());
                if (probability < 0.0 || probability > 1.0) {
                    ythrow TConfigParseError{} << "bad probability " << value->AsString();
                }
                m.push_back(probability);
            });
            return;
        }

        if (Configure(key, Copy(value->AsSubConfig()))) {
            return;
        }
    } END_PARSE

private:
    TVector<double> Probabilities_;
};

INodeHandle<IPolicyFactory>* Handle() {
    return TPolicyFactory::Handle();
}

}  // namespace NSrvKernel::NBackOff
