#pragma once

#include <balancer/kernel/balancer/policy.h>

#include <util/datetime/base.h>

namespace NSrvKernel::NTimeout {

class TTimeoutPolicy : public IPolicy {
public:
    TTimeoutPolicy(THolder<IPolicy> slave, TDuration timeout) noexcept
        : Slave_(std::move(slave))
        , Timeout_(timeout)
    {}

    IBackend* Next(IAlgorithm* algo, bool fastAttempt) noexcept override {
        if (HasTimeout()) {
            if (Deadline_ == TInstant()) {
                Deadline_ = Timeout_.ToDeadLine();
                return Slave_->Next(algo, fastAttempt);
            } else if (Now() < Deadline_) {
                return Slave_->Next(algo, fastAttempt);
            } else {
                return nullptr;
            }
        } else {
            return Slave_->Next(algo, fastAttempt);
        }
    }

    void RegisterSuccess() noexcept override {
        Slave_->RegisterSuccess();
    }

    void RegisterFail() noexcept override {
        Slave_->RegisterFail();
    }

    void MarkAsRetry() noexcept override {}

private:
    bool HasTimeout() const noexcept {
        return Timeout_ != TDuration();
    }

private:
    THolder<IPolicy> Slave_;
    TDuration Timeout_;
    TInstant Deadline_;
};

INodeHandle<IPolicyFactory>* Handle();

}  // namespace NSrvKernel::NTimeout
