#include "unique_retry.h"

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

namespace NSrvKernel::NUniqueRetry {

IBackend* TUniqueRetryPolicy::Next(IAlgorithm* algo, bool) noexcept {
    IBackend* next = DoNext(algo);
    if (next == nullptr) {
        const bool hasResets = Opts_.ResetsCount ? Opts_.ResetsCount-- : false;
        if (!hasResets) {
            return nullptr;
        }
        algo->Reset();
        return DoNext(algo);
    }
    return next;
}

IBackend* TUniqueRetryPolicy::DoNext(IAlgorithm* algo) noexcept {
    IBackend* next = algo->Next();
    if (next != nullptr) {
        algo->RemoveSelected(next);
    }
    return next;
}


POLICY_FACTORY(unique_retry_policy) {
    TPolicyFactory(const TModuleParams& mp) {
        mp.Config->ForEach(this);

        if (!Opts_.ResetsCount) {
            ythrow TConfigParseError{} << "balancer2/unique_retry_policy: resets_count is zero, use unique_policy!";
        }
    }

    THolder<IPolicy> ConstructPolicy(const TStepParams&) noexcept override {
        return MakeHolder<TUniqueRetryPolicy>(Opts_);
    }

    void FillFeatures(TPolicyFeatures&) const noexcept override {}

    void Init(IWorkerCtl*) override {}

private:
    START_PARSE {
        ON_KEY("resets_count", Opts_.ResetsCount) {
            return;
        }
    } END_PARSE

private:
    TUniqueRetryOptions Opts_;
};

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

}  // namespace NSrvKernel::NUniqueRetry
