#include "watermark.h"

#include <balancer/kernel/balancer/algorithm.h>
#include <balancer/kernel/balancing/backend.h>
#include <balancer/kernel/fs/kv_file_consumer.h>
#include <balancer/kernel/log/errorlog.h>

#include <balancer/modules/balancer/policies/policy_factory.h>

namespace NSrvKernel::NWatermark {

double TSharedSuccessTracker::RecalcNewValueSign(double newValue, double oldValue) noexcept {
    /* newValue is always non-negative */
    if (oldValue < 0) {
        /* old state = disabled */
        if (!(newValue < Lo_)) {
            /* keep same sign if no transition occurs */
            newValue = -newValue;
        }
    } else {
        /* old state = enabled */
        if (newValue > Hi_) {
            /* negate value if transition occurs */
            newValue = -newValue;
        }
    }

    return newValue;
}

void TState::DoUpdateParams(const TStringBuf contents, const TConnDescr& descr) noexcept {
    NSrvKernel::TExternalWeightsFileConsumer consumer;
    try {
        ProcessKvData(consumer, contents);
        const auto& kv = consumer.Storage();
        double lo = Params_.Lo;
        double hi = Params_.Hi;

        if (const double* fileLo = MapFindPtr(kv, "lo")) {
            lo = *fileLo;
        }

        if (const double* fileHi = MapFindPtr(kv, "hi")) {
            hi = *fileHi;
        }

        if (const double* fileShared = MapFindPtr(kv, "shared")) {
            Params_.Shared = *fileShared != 0;
        }

        if (lo <= hi) {
            SharedTracker_->SetParams(lo, hi);
            PerWorkerTracker_->SetParams(lo, hi);
        }
    } catch (const yexception& e) {
        constexpr char NAME[] = "watermark_policy";
        LOG_ERROR(TLOG_ERR, descr, "watermark_policy Error parsing params_file: " << e.what());
    }
}

void TState::DoUpdateSwitchedState(TStringBuf contents, const TConnDescr& descr) noexcept {
    NSrvKernel::TExternalWeightsFileConsumer consumer;
    try {
        ProcessKvData(consumer, contents);
        const auto& kv = consumer.Storage();

        if (const double* enabled = MapFindPtr(kv, Params_.SwitchKey)) {
            SwitchedOn_ = *enabled == 1;
        } else {
            SwitchedOn_ = Params_.SwitchDefault;
        }
    } catch (const yexception& e) {
        constexpr char NAME[] = "watermark_policy";
        LOG_ERROR(TLOG_ERR, descr, "watermark_policy Error parsing switch_file: " << e.what());
    }
}

void TState::UpdateParams(const TConnDescr& descr) noexcept {
    {
        const auto& data = ParamsFileReader_.Data();
        if (data.Id() != ParamsFileData_.Id()) {
            ParamsFileData_ = data;
            DoUpdateParams(ParamsFileData_.Data(), descr);
        }
    }

    {
        const auto& data = SwitchFileReader_.Data();
        if (data.Id() != SwitchFileData_.Id()) {
            SwitchFileData_ = data;
            DoUpdateSwitchedState(SwitchFileData_.Data(), descr);
        }
    }
}


POLICY_FACTORY_WITH_TLS(watermark_policy, TState), public TModuleParams {
public:
    TPolicyFactory(const TModuleParams& mp)
        : TPolicyBase(mp.Control->GetCountOfChildren())
        , TModuleParams(mp)
    {
        Y_VERIFY(FailRate_.is_lock_free(), "std::atomic<double> is not lock free");
        FailRate_.store(0);

        Config->ForEach(this);
        CheckConfiguration();

        if (Params_.Lo > Params_.Hi) {
            ythrow TConfigParseError{} << "lo must be less or equal to hi in watermark_policy";
        }

        if (Params_.Coeff <= 0.0 || Params_.Coeff >= 1.0) {
            ythrow TConfigParseError{} << "\"coeff\" must be in (0.0, 1.0) in watermark_policy";
        }
    }


    THolder<IPolicy> ConstructPolicy(const TStepParams& params) noexcept override {
        return MakeHolder<TWatermarkPolicy>(TPolicyBase::ConstructPolicy(params), GetTls(&params.Descr->Process()), *params.Descr);
    }

    THolder<TState> InitTls(IWorkerCtl* process) override {
        auto tls = MakeHolder<TState>(Params_, FailRate_);
        if (ParamsFilename_) {
            tls->ParamsFileReader_ = process->SharedFiles()->FileReReader(ParamsFilename_, TDuration::Seconds(1));
        }

        if (SwitchFilename_ && Params_.SwitchKey) {
            tls->SwitchFileReader_ = process->SharedFiles()->FileReReader(SwitchFilename_, TDuration::Seconds(1));
        }

        return tls;
    }

private:
    START_PARSE {
        ON_KEY("lo", Params_.Lo) {
            return;
        }

        ON_KEY("hi", Params_.Hi) {
            return;
        }

        ON_KEY("coeff", Params_.Coeff) {
            return;
        }

        ON_KEY("params_file", ParamsFilename_) {
            return;
        }

        ON_KEY("switch_file", SwitchFilename_) {
            return;
        }

        ON_KEY("switch_key", Params_.SwitchKey) {
            return;
        }

        ON_KEY("switch_default", Params_.SwitchDefault) {
            return;
        }

        ON_KEY("is_shared", Params_.Shared) {
            return;
        }

        ON_KEY("shared", Params_.Shared) {
            Cerr << "'shared' option in watermark_policy is deprecated, use 'is_shared' instead" << Endl;
            return;
        }

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

private:
    std::atomic<double> FailRate_;

    TString ParamsFilename_;
    TString SwitchFilename_;
    TParams Params_;
};

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

}  // namespace NSrvKernel::NWatermark
