#include "active.h"
#include "backends_factory.h"
#include "base_algorithm.h"

#include <balancer/kernel/balancing/per_worker_backend.h>
#include <balancer/kernel/coro/coroutine.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/pinger/pinger.h>
#include <balancer/kernel/pinger/pinger2.h>
#include <balancer/kernel/ctl/children_process_common.h>

namespace NSrvKernel::NModBalancer {

TSharedBool GetSharedState(const TBackendsUID& uid, ICtl* ctl) {
    struct TSharedState: public IModuleData {
        THashMap<TBackendsUID::TValue, TSharedBool> Map;
    };

    THolder<IModuleData>& moduleData = ctl->GetModuleData("backends_active");

    if (!moduleData) {
        moduleData.Reset(new TSharedState);
    }

    TSharedState* sharedState = dynamic_cast<TSharedState*>(moduleData.Get());
    Y_VERIFY(sharedState);

    auto& stateMap = sharedState->Map;

    if (TSharedBool* res = stateMap.FindPtr(uid.Value)) {
        return *res;
    } else {
        return stateMap.emplace(uid.Value, ctl->SharedAllocator()).first->second;
    }
}

namespace {
    class TBackend : public TPerWorkerBaseBackend {
    public:
        TBackend(TBackendDescriptor::TRef descr, bool pingerTurnedOn)
            : TPerWorkerBaseBackend(std::move(descr))
            , PingerTurnedOn_(pingerTurnedOn)
        {}

        void DoOnCompleteRequest(const TDuration& /*answerTime*/) noexcept override {}

        void DoOnFailRequest(const TError&, const TDuration& /*answerTime*/) noexcept override {}

        double GetCurrentWeight() const noexcept {
            return CurrentWeight_;
        }

        void SetCurrentWeight(double currentWeight) noexcept {
            CurrentWeight_ = currentWeight;
        }

        void AddCurrentWeight(double value) noexcept {
            CurrentWeight_ += value;
        }

        void PrintInfo(NJson::TJsonWriter& out) const noexcept {
            out.Write("weight", Backend_->Weight());
            out.Write("enabled", Enabled());
            out.Write("currentWeight", CurrentWeight_);
            PrintSuccFailRate(out);
            PrintProxyInfo(out);
        }

        void SetCounter(TSharedCounter* enabledCounter) {
            EnabledBackendsCounter_ = enabledCounter;
        }

        void InitCounter() {
            if (EnabledBackendsCounter_) {
                EnabledBackendsCounter_->Add(1);
            }
        }

        bool Enabled() const noexcept override {
            if (PingerTurnedOn_) {
                return TPerWorkerBaseBackend::Enabled();
            } else {
                return Enabled_;
            }
        }

        void SetEnabled(bool value) noexcept override {
            if (value != Enabled() && EnabledBackendsCounter_) {
                if (value) {
                    EnabledBackendsCounter_->Add(1);
                } else {
                    EnabledBackendsCounter_->Sub(1);
                }
            }

            if (PingerTurnedOn_) {
                TPerWorkerBaseBackend::SetEnabled(value);
            } else {
                Enabled_ = value;
            }
        }

        double GetWeightFromPing(bool disabled) const noexcept {
            return disabled ? 1. : WeightFromPing();
        }

        double WeightFromPing() const noexcept override {
            if (PingerTurnedOn_) {
                return TPerWorkerBaseBackend::WeightFromPing();
            } else {
                return WeightFromPing_;
            }
        }

        void SetWeightFromPing(double value, bool valueFromBackend) noexcept override {
            if (PingerTurnedOn_) {
                TPerWorkerBaseBackend::SetWeightFromPing(value, valueFromBackend);
            } else {
                WeightFromPing_ = value;
            }
        }

    private:
        const bool PingerTurnedOn_ = false;
        bool Enabled_ = true;
        double CurrentWeight_ = 0.;
        double WeightFromPing_ = 1.;
        TSharedCounter* EnabledBackendsCounter_ = nullptr;
    };
}

BACKENDS_TLS(active) {
    TTls(const TSharedHistogram& hist, const TSharedCounter& disabled, const TMaybe<TSharedCounter>& enabled,
        const TPingerConfigUnparsed& pingerConfig, size_t workerId)
        : StatsDelayRatioHistogram(hist, workerId)
        , DisabledEntireGroupCounter(disabled, workerId)
        , EnabledBackendsCounter(enabled.Defined() ? TSharedCounter(*enabled, workerId) : TMaybe<TSharedCounter>())
        , PingerConfig(pingerConfig)
    {}

    bool IsBackendWeightDisable() const noexcept {
        return BackendWeightDisableFileChecker.Exists();
    }

    bool EntireGroupDisabled = false;
    TDeque<TBackend> Backends;
    TSharedHistogram StatsDelayRatioHistogram;
    TSharedCounter DisabledEntireGroupCounter;
    TMaybe<TSharedCounter> EnabledBackendsCounter;

    TInstant LastRequestTime;
    TPingerConfig PingerConfig;
    THolder<IPingerBackendGroupManager> PingerBackendGroupManager;
    TDeque<TPinger> Pingers;
    TSharedFileExistsChecker BackendWeightDisableFileChecker;
    TCoroutine StateUpdater;
};

BACKENDS_WITH_TLS(active), public TModuleParams {
// Initialization
// --------------------------------------------------------------------------------
private:
    class TActiveAlgorithm : public TAlgorithmWithRemovals {
    private:
        struct TBackendState {
            TBackend* Backend;
            double CurrentWeight;
        };

    public:
        TActiveAlgorithm(TBackends* parent, const TStepParams& params) noexcept
            : TAlgorithmWithRemovals(&params.Descr->Process())
            , Parent_(parent)
        {}

        void RemoveSelected(IBackend* backend) noexcept override {
            TAlgorithmWithRemovals::RemoveSelected(backend);
            DoRemove(backend);
        }

        IBackend* Next() noexcept override {
            Parent_->UpdateQuorum(*Process_);
            if (!Parent_->IgnoreStateOnFailedQuorum_ && !Parent_->Available_) { // no valid backends
                return nullptr;
            }

            TBackend* ret = nullptr;

            double max = -std::numeric_limits<double>::max();
            double total = 0.;
            const bool ignoreActiveChecks = Parent_->IgnoreStateOnFailedQuorum_ && !Parent_->Available_;
            const bool backendWeightDisable = Parent_->GetTls(*Process_).IsBackendWeightDisable();
            for (auto& backend : Parent_->GetTls(*Process_).Backends) {
                if (!backend.Enabled() && !ignoreActiveChecks) {
                    continue;
                }

                if (IsRemoved(&backend)) {
                    continue;
                }

                const double weight = backend.OriginalWeight() * backend.GetWeightFromPing(backendWeightDisable);
                total += weight;
                backend.AddCurrentWeight(weight);
                const double cur = backend.GetCurrentWeight();
                if (cur > max) {
                    max = cur;
                    ret = &backend;
                }
            }

            if (ret != nullptr) {
                CurrentWeights_[ret] = TBackendState{ ret, total };
            }

            return ret;
        }

        IBackend* NextByName(TStringBuf name, bool /*allowZeroWeights*/) noexcept override {
            auto it = Parent_->NamedBackends_.find(name);
            if (it == Parent_->NamedBackends_.end()) {
                return nullptr;
            }

            auto ret = &Parent_->GetTls(*Process_).Backends[it->second];

            if (ret->Enabled() && !IsRemoved(ret)) {
                const double weight = ret->OriginalWeight() * ret->GetWeightFromPing(Parent_->GetTls(*Process_).IsBackendWeightDisable());
                ret->AddCurrentWeight(weight);
                // TODO: check whether it is ok to have such a small currentWeight
                CurrentWeights_[ret] = TBackendState{ ret, weight };
                return ret;
            }

            return nullptr;
        }

    private:
        void DoRemove(IBackend* backend) noexcept {
            auto it = CurrentWeights_.find(backend);
            if (it != CurrentWeights_.end()) {
                auto& state = it->second;
                if (state.Backend) {
                    state.Backend->AddCurrentWeight(-state.CurrentWeight);
                }
            }
        }

    private:
        TBackends* Parent_ = nullptr;
        TMap<IBackend*, TBackendState> CurrentWeights_;
    };

public:
    TBackends(const TModuleParams& mp, const TBackendsUID& uid)
        : TBackendsWithTLS(mp)
        , TModuleParams(mp)
        , Available_(GetSharedState(uid, mp.Control))
        , BackendsUIDValue_(uid.Value)
        , PingerTurnedOn_(Control->IsPingerEnabled())
    {
        Config->ForEach(this);

        if (!PingerConfig_.UnparsedRequest && !PingerConfig_.TcpCheck) {
            ythrow TConfigParseError{} << "You must enable tcp check or http check";
        }

        if (PingerConfig_.UnparsedRequest && PingerConfig_.TcpCheck) {
            ythrow TConfigParseError{} << "Can't use tcp check and http check simultaneously";
        }

        if (UseBackendWeight_ && PingerConfig_.TcpCheck) {
            ythrow TConfigParseError{} << "Can't use backend dynamic weight with tcp check";
        }

        if (PingerConfig_.Delay == TDuration::Zero()) {
            ythrow TConfigParseError{} << "No ping delay configured";
        }

        if (UseBackendWeight_ && !PingerTurnedOn_) {
            ythrow TConfigParseError{} << "BackendWeight work only with pinger";
        }

        const TVector<ui64> intervals({0, 50, 100, 200, 400, 800, 1600, 10000});

        TString delayHistName = "max_delay_of_pings";
        if (!StatsUuid_.Empty()) {
            delayHistName = "active-" + StatsUuid_ + "-" + delayHistName;
        }
        StatsDelayRatioHistogram_ = Control->SharedStatsManager().MakeHistogram(delayHistName, intervals).AllowDuplicate().Build();

        DisabledEntireGroupCounter_ = Control->SharedStatsManager()
                .MakeGauge("active-entire_group_disabled").AllowDuplicate().Build();

        if (!StatsUuid_.Empty()) {
            EnabledBackendsCounter_ = Control->SharedStatsManager()
                .MakeGauge("active-" + StatsUuid_ + "-all_alive_flag").AllowDuplicate().Build();
        }
    }

private:
    THolder<TTls> DoInit(IWorkerCtl* process) noexcept override {
        auto type = process->WorkerType();
        auto tls = MakeHolder<TTls>(*StatsDelayRatioHistogram_, *DisabledEntireGroupCounter_, EnabledBackendsCounter_,
            PingerConfig_, process->WorkerId());
        for (auto& i : BackendDescriptors()) {
            tls->Backends.emplace_back(i, PingerTurnedOn_);
            if (!StatsUuid_.Empty()) {
                tls->Backends.back().SetCounter(tls->EnabledBackendsCounter.Get());
            }
        }

        if (!PingerTurnedOn_) {
            for (auto& backend : tls->Backends) {
                tls->Pingers.emplace_back(backend, tls->LastRequestTime, &process->Executor(),
                        tls->PingerConfig, *process, Steady_);
            }
        } else if (type == NProcessCore::TChildProcessType::Pinger) {
            tls->PingerBackendGroupManager = process->SharedPingerManager().GeneratePingerBackendGroupManager();
            for (auto& backend : tls->Backends) {
                THolder<IBackendPinger> pinger = MakeHolder<TPingerV2>(backend, tls->PingerConfig, *process, UseBackendWeight_);
                tls->PingerBackendGroupManager->AddPinger(backend.Name(), std::move(pinger));
                backend.InitCounter();
            }

            tls->StateUpdater = TCoroutine{
                ECoroType::Service,
                "pingtime_updater", &process->Executor(), [this, process](TContExecutor* const exec) {
                    auto* const cont = exec->Running();
                    const TDuration sleepTime = Min(TDuration::Seconds(1), PingerConfig_.Delay);
                    while (!cont->Cancelled()) {
                        UpdateLastPingCounter(*process, GetTheOldestPingTime(*process));
                        cont->SleepT(sleepTime);
                    }
                }, &process->Executor()
            };
        }

        if (process->WorkerType() == NSrvKernel::NProcessCore::TChildProcessType::Default) {
            if (!PingerTurnedOn_) {
                for (auto& backend : tls->Backends) {
                    backend.InitCounter();
                }
            }

            tls->StateUpdater = TCoroutine{
                ECoroType::Service,
                "state_updater", &process->Executor(), [this, process](TContExecutor* const exec) {
                    auto* const cont = exec->Running();
                    cont->SleepT(TDuration::Seconds(RandomNumber<double>() * PingerConfig_.Delay.SecondsFloat()));

                    TInstant lastQuorumUpdate = Now();
                    const TDuration sleepTime = Min(TDuration::Seconds(1), PingerConfig_.Delay);
                    while (!cont->Cancelled()) {
                        if (Steady_ &&  lastQuorumUpdate + PingerConfig_.Delay < Now()) {
                            UpdateQuorum(*process);
                            lastQuorumUpdate = Now();
                        }
                        UpdateEntireGroupState(*process);
                        if (!PingerTurnedOn_) {
                            UpdateLastPingCounter(*process, GetTheOldestPingTime(*process));
                        }
                        cont->SleepT(sleepTime);
                    }
                }, &process->Executor()
            };
        }

        if (UseBackendWeight_ && BackendWeightDisableFile_) {
            tls->BackendWeightDisableFileChecker = process->SharedFiles()->FileChecker(
                    BackendWeightDisableFile_, TDuration::Seconds(1)
            );
        }
        return tls;
    }

    void ProcessPolicyFeatures(const TPolicyFeatures& features) override {
        if (features.WantsHash) {
            PrintOnce("WARNING in balancer2/active: policies with hash "
                      "are not supported and will be ignored");
        }
    }

    START_PARSE {
        if (PingerConfig_.TryConsume(key, value)) {
            return;
        }

        ON_KEY("steady", Steady_) {
            return;
        }

        ON_KEY("quorum", Quorum_) {
            if (Quorum_ < 0) {
                ythrow TConfigParseError{} << "balancer2: quorum can't be negative";
            }
            return;
        }

        ON_KEY("hysteresis", Hysteresis_) {
            if (Hysteresis_ < 0) {
                ythrow TConfigParseError{} << "balancer2: hysteresis can't be negative";
            }
            return;
        }

        // mode that ignores active check results during losing of quorum
        ON_KEY("ignore_state_on_failed_quorum", IgnoreStateOnFailedQuorum_) {
            return;
        }

        ON_KEY("use_backend_weight", UseBackendWeight_) {
            return;
        }

        ON_KEY("backend_weight_disable_file", BackendWeightDisableFile_) {
            return;
        }

        if (key == "uuid") {
            StatsUuid_ = value->AsString();
            return;
        }

        auto descr = MakeHolder<TBackendDescriptor>(Copy(value->AsSubConfig()), key);
        if (descr->Weight() < 0) {
            ythrow TConfigParseError{} << "Backend weight must be positive";
        }

        NamedBackends_[descr->Name()] = Size();
        Add(std::move(descr));

        return;
    } END_PARSE

// --------------------------------------------------------------------------------


    void DumpBackends(NJson::TJsonWriter& out, const TTls& tls) const noexcept override {
        out.OpenMap();
        out.Write("id", BackendsUIDValue_);
        out.OpenArray("backends");
        for (const auto& backend : tls.Backends) {
            out.OpenMap();
            backend.PrintInfo(out);
            out.CloseMap();
        }
        out.CloseArray();
        out.CloseMap();
    }

// Functionality
// --------------------------------------------------------------------------------
    THolder<IAlgorithm> ConstructAlgorithm(const TStepParams& params) noexcept override {
        GetTls(params.Descr->Process()).LastRequestTime = Now();
        return MakeHolder<TActiveAlgorithm>(this, params);
    }

    void UpdateQuorum(const IWorkerCtl& process) {
        double total = 0.;
        for (const auto& backend : GetTls(process).Backends) {
            if (backend.Enabled()) {
                total += backend.OriginalWeight();
            }
        }

        if (Available_
            && total < Max<double>(
            std::numeric_limits<double>::epsilon(), Quorum_ - Hysteresis_))
        {
            Available_ = false;
        } else if (!Available_ && total >= Quorum_ + Hysteresis_) {
            Available_ = true;
        }
    }


    void UpdateEntireGroupState(const IWorkerCtl& process) {
        size_t totalEnabled = 0;
        auto& tls = GetTls(process);

        for (const auto& backend :tls.Backends) {
            if (backend.Enabled()) {
                totalEnabled += 1;
            }
        }

        if (totalEnabled == 0) {
            if (tls.EntireGroupDisabled) {
                return;
            }

            tls.EntireGroupDisabled = true;
            tls.DisabledEntireGroupCounter.Inc();
        } else {
            if (!tls.EntireGroupDisabled) {
                return;
            }

            tls.EntireGroupDisabled = false;
            tls.DisabledEntireGroupCounter.Dec();
        }
    }

    void UpdateLastPingCounter(const IWorkerCtl& process, ui64 theOldestPingTime) {
        const ui16 ratio = 100.0 * (Now().Seconds() - theOldestPingTime) / PingerConfig_.Delay.SecondsFloat();
        GetTls(process).StatsDelayRatioHistogram.AddValue(ratio);
    }

    ui64 GetTheOldestPingTime(const IWorkerCtl& process) const {
        auto& tls = GetTls(process);
        if (PingerTurnedOn_) {
            return tls.PingerBackendGroupManager->GetTheOldestPingTime();
        } else {
            if (tls.Pingers.empty()) {
                return 0;
            }
            const ui64 theOldestPingTime = (*MinElement(tls.Pingers.begin(), tls.Pingers.end(), [](const TPinger& left, const TPinger& right){
                return left.GetLastPingTime() < right.GetLastPingTime();
            })).GetLastPingTime();
            return theOldestPingTime;
        }
    }
// --------------------------------------------------------------------------------


// State
// --------------------------------------------------------------------------------
private:
    TSharedBool Available_;
    size_t BackendsUIDValue_ = -1;
    double Quorum_ = std::numeric_limits<double>::epsilon();
    double Hysteresis_ = 0;

    THashMap<TString, size_t> NamedBackends_;

    bool IgnoreStateOnFailedQuorum_ = false;
    bool Steady_ = true;
    TPingerConfigUnparsed PingerConfig_;
    bool PingerTurnedOn_ = false;
    bool UseBackendWeight_ = false;
    TString BackendWeightDisableFile_;
    TMaybe<TSharedHistogram> StatsDelayRatioHistogram_;
    TString StatsUuid_;
    TMaybe<TSharedCounter> EnabledBackendsCounter_;
    TMaybe<TSharedCounter> DisabledEntireGroupCounter_;
// --------------------------------------------------------------------------------
};

INodeHandle<IBackends>* NActive::Handle() {
    return TBackends::Handle();
}

}  // namespace NSrvKernel::NModBalancer
