#pragma once

#include <balancer/kernel/ctl/ctl.h>
#include <balancer/kernel/helpers/default_instance.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/module/module_face.h>

#include <library/cpp/config/sax.h>
#include <library/cpp/regex/pire/regexp.h>

namespace NSrvKernel {
    struct TXYandexRetryFsm : public NRegExp::TFsm, public TWithDefaultInstance<TXYandexRetryFsm> {
        TXYandexRetryFsm() noexcept
            : TFsm("X-Yandex-Balancer-Retry", NRegExp::TFsm::TOptions().SetCaseInsensitive(true))
        {}
    };

    class TAttemptsRateLimiter : public NConfig::IConfig::IFunc {
    public:
        TAttemptsRateLimiter() noexcept {
            Y_VERIFY(AttemptsRate_.is_lock_free(), "atomic double is not lock free");
            AttemptsRate_.store(0);
        }

        TAttemptsRateLimiter(double limit, double coeff, double maxBudget = -1) noexcept
            : Exists_{true}
            , IsNew_{maxBudget > 0}
            , Limit_{limit}
            , Coeff_{coeff}
            , MaxBudget_{maxBudget}
            , AttemptsRate_{MaxBudget_}
        {
            Y_VERIFY(AttemptsRate_.is_lock_free(), "atomic double is not lock free");
        }

        void Verify(TDuration hedgedDelay) const {
            if (Limit_ <= 0) {
                ythrow TConfigParseError{} << "attempts rate limit must be positive";
            }
            if (IsNew()) {
                if (MaxBudget_ < 0) {
                    ythrow TConfigParseError{} << "max budget for rate limiter must be non-negative";
                }
                if (Limit_ > MaxBudget_) {
                    ythrow TConfigParseError{} << "attempts rate limit must not exceed max budget";
                }
            } else {
                if (Exists_ && hedgedDelay != TDuration::Zero()) {
                    ythrow TConfigParseError{} << "trying to use hedged requests with old rate limiter";
                }
            }
        }

        void Init() noexcept {
            if (!IsNew()) {
                Y_VERIFY(Limit_ > 0);
            }
        }

        bool RetryAllowed(size_t retriesPassedToBackend) noexcept {
            if (!Exists_) {
                return true;
            }

            if (IsNew()) {
                bool allowed = false;
                for (;;) {
                        double oldValue = AttemptsRate_.load();
                        if (oldValue >= 1) {
                            allowed = true;
                            if (AttemptsRate_.compare_exchange_weak(oldValue, oldValue - 1)) {
                                break;
                            }
                        } else {
                            allowed = false;
                            break;
                        }
                }
                return allowed;
            } else {
                double newRate = Coeff_ * AttemptsRate_.load() + (1.0 - Coeff_) * retriesPassedToBackend;
                return newRate <= Limit_;
            }
        }

        void RegisterRequest() noexcept {
            if (!IsNew() || !Exists_) {
                return;
            }

            for (;;) {
                double oldValue = AttemptsRate_.load();
                double newValue = Min(MaxBudget_, oldValue + Limit_);
                if (AttemptsRate_.compare_exchange_weak(oldValue, newValue)) {
                    break;
                }
            }
        }

        void RegisterAttempts(size_t count) noexcept {
            if (IsNew() || !Exists_) {
                return;
            }

            for (;;) {
                double oldValue = AttemptsRate_.load();
                double newValue = Coeff_ * oldValue + (1.0 - Coeff_) * count;
                if (AttemptsRate_.compare_exchange_weak(oldValue, newValue)) {
                    break;
                }
            }
        }

        void RegisterFastAttempt() noexcept {
            if (!IsNew() || !Exists_) {
                return;
            }

            for (;;) {
                double oldValue = AttemptsRate_.load();
                double newValue = Min(MaxBudget_, oldValue + 1);
                if (AttemptsRate_.compare_exchange_weak(oldValue, newValue)) {
                    break;
                }
            }
        }

    private:
        bool IsNew() const noexcept {
            return IsNew_.GetOrElse(false);
        }

        void VerifyOption(bool isNew) {
            if (IsNew_.Defined() && *IsNew_ != isNew) {
                ythrow TConfigParseError{} << "inconsistent use of old and new attempts rate limiter options";
            }
            IsNew_ = isNew;
        }

        START_PARSE {
            Exists_ = true;

            ON_KEY("coeff", Coeff_) {
                VerifyOption(false);
                return;
            }

            ON_KEY("limit", Limit_) {
                return;
            }

            ON_KEY("max_budget", MaxBudget_) {
                VerifyOption(true);
                AttemptsRate_.store(MaxBudget_);
                return;
            }

            if (key == "switch_default" || key == "switch_file" || key == "switch_key") {
                return;
            }

            if (key == "switch") {
                return;
            }
        } END_PARSE

    private:
        bool Exists_ = false;
        TMaybe<bool> IsNew_;

        double Limit_ = 1.0;
        double Coeff_ = 0.99;   // Old version
        double MaxBudget_ = -1; // New version

        std::atomic<double> AttemptsRate_;
    };
}
