#pragma once

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

#include <balancer/kernel/module/module_face.h>

namespace NSrvKernel::NWatermark {

class IAttemptTracker {
public:
    virtual ~IAttemptTracker() {};
    virtual void RegisterSuccess() noexcept = 0;
    virtual void RegisterFail() noexcept = 0;
    virtual bool Enabled() noexcept = 0;
    virtual void SetParams(double lo, double hi) noexcept = 0;
};

class TSuccessTracker : public IAttemptTracker {
public:
    TSuccessTracker(double lo, double hi, double coeff) noexcept
        : Coeff_(coeff)
        , Lo_(lo)
        , Hi_(hi)
    {}

    bool Enabled() noexcept override {
        return Enabled_;
    }

    void RegisterSuccess() noexcept override {
        UpdateStateAfterAttempt(false);
    }

    void RegisterFail() noexcept override {
        UpdateStateAfterAttempt(true);
    }

    void SetParams(double lo, double hi) noexcept override {
        Lo_ = lo;
        Hi_ = hi;
        Enabled_ = FailRate_ < Lo_;
    }

    double FailRate() const noexcept {
        return FailRate_;
    }

private:
    void UpdateStateAfterAttempt(bool failedAttempt) noexcept {
        FailRate_ = Coeff_ * FailRate_ + (1. - Coeff_) * (failedAttempt ? 1 : 0);

        if (!Enabled_ && (FailRate_ < Lo_)) {
            Enabled_ = true;
        }

        if (Enabled_ && (FailRate_ > Hi_)) {
            Enabled_ = false;
        }
    }

    double Coeff_ = 0;
    double Lo_ = 0;
    double Hi_ = 0;
    double FailRate_ = 0;
    bool Enabled_ = true;
};

class TSharedSuccessTracker : public IAttemptTracker {
private:
    /* sign used to represent current state: < 0 for disabled, >= 0 for enabled */

public:
    TSharedSuccessTracker(std::atomic<double>& failRate, double lo, double hi, double coeff) noexcept
        : FailRate_(failRate)
        , Lo_(lo)
        , Hi_(hi)
        , Coeff_(coeff)
    {}

    void SetParams(double lo, double hi) noexcept override {
        Lo_ = lo;
        Hi_ = hi;
        UpdateState();
    }

    bool Enabled() noexcept override {
        return FailRate_.load() >= 0;
    }

    void RegisterSuccess() noexcept override {
        UpdateStateAfterAttempt(false);
    }

    void RegisterFail() noexcept override {
        UpdateStateAfterAttempt(true);
    }

    double FailRate() const noexcept {
        return FailRate_.load();
    }

private:
    double RecalcNewValueSign(double newValue, double oldValue) noexcept;

    void UpdateState() noexcept {
        for (;;) {
            double oldValue = FailRate_.load();
            double newValue = RecalcNewValueSign(Abs(oldValue), oldValue);
            if (FailRate_.compare_exchange_weak(oldValue, newValue)) {
                break;
            }
        }
    }

    void UpdateStateAfterAttempt(bool failedAttempt) noexcept {
        for (;;) {
            double oldValue = FailRate_.load();
            double newValue = Coeff_ * Abs(oldValue) + (1 - Coeff_) * (failedAttempt ? 1 : 0);
            newValue = RecalcNewValueSign(newValue, oldValue);
            if (FailRate_.compare_exchange_weak(oldValue, newValue)) {
                break;
            }
        }
    }

    std::atomic<double>& FailRate_;
    double Lo_ = 0;
    double Hi_ = 0;
    double Coeff_ = 0;
};

struct TParams {
    TString SwitchKey;
    double Lo = 0;
    double Hi = 0;
    double Coeff = 0.99;
    bool SwitchDefault = true;
    bool Shared = false;
};

class TState {
public:
    TState(const TParams& params, std::atomic<double>& failRate)
        : Params_(params)
        , SharedTracker_(MakeHolder<TSharedSuccessTracker>(failRate, params.Lo, params.Hi, params.Coeff))
        , PerWorkerTracker_(MakeHolder<TSuccessTracker>(params.Lo, params.Hi, params.Coeff))
    {}

    void RegisterSuccess(const TConnDescr& descr) noexcept {
        UpdateParams(descr);
        SharedTracker_->RegisterSuccess();
        PerWorkerTracker_->RegisterSuccess();
    }

    void RegisterFail(const TConnDescr& descr) noexcept {
        UpdateParams(descr);
        SharedTracker_->RegisterFail();
        PerWorkerTracker_->RegisterFail();
    }

    bool Enabled(const TConnDescr& descr) noexcept {
        UpdateParams(descr);

        if (!SwitchedOn_) {
            return true;
        } else if (Params_.Shared) {
            return SharedTracker_->Enabled();
        } else {
            return PerWorkerTracker_->Enabled();
        }
    }

private:
    void DoUpdateParams(const TStringBuf contents, const TConnDescr& descr) noexcept;

    void DoUpdateSwitchedState(TStringBuf contents, const TConnDescr& descr) noexcept;

    void UpdateParams(const TConnDescr& descr) noexcept;

public:
    TParams Params_;
    TSharedFileReReader ParamsFileReader_;
    TSharedFileReReader SwitchFileReader_;
    TSharedFileReReader::TData ParamsFileData_;
    TSharedFileReReader::TData SwitchFileData_;

    THolder<TSharedSuccessTracker> SharedTracker_;
    THolder<TSuccessTracker> PerWorkerTracker_;

    bool SwitchedOn_ = true;
};

class TWatermarkPolicy : public IPolicy {
public:
    TWatermarkPolicy(THolder<IPolicy> slave, TState& state, const TConnDescr& descr) noexcept
        : Slave_(std::move(slave))
        , State_(state)
        , Descr_(descr)
    {}

    IBackend* Next(IAlgorithm* algo, bool fastAttempt) noexcept override {
        if (First_ || fastAttempt || State_.Enabled(Descr_)) {
            First_ = false;
            return Slave_->Next(algo, fastAttempt);
        } else {
            return nullptr;
        }
    }

    void MarkAsRetry() noexcept override {
        First_ = false;
    }

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

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

private:
    THolder<IPolicy> Slave_;
    TState& State_;
    const TConnDescr& Descr_;
    bool First_ = true;
};

INodeHandle<IPolicyFactory>* Handle();

}  // namespace NSrvKernel::NWatermark
