//
// Created by Alexander Gryanko on 07/09/2017.
//

#include "backends_factory.h"
#include "base_algorithm.h"
#include "rendezvous_hashing.h"
#include "weights_file.h"

#include <balancer/kernel/balancing/per_worker_backend.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/fs/kv_file_consumer.h>
#include <balancer/kernel/fs/shared_files.h>
#include <balancer/kernel/module/iface.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/pinger/pinger.h>

#include <contrib/libs/highwayhash/highwayhash/sip_hash.h>

#include <util/datetime/base.h>
#include <util/digest/murmur.h>
#include <util/generic/set.h>
#include <util/generic/strbuf.h>
#include <util/generic/ymath.h>

namespace NSrvKernel::NModBalancer {

using namespace highwayhash;

static constexpr ui64 FIFTY_THREE_ONES = (0xFFFFFFFFFFFFFFFF >> (64 - 53));
static constexpr double FIFTY_THREE_ZEROS = 1L << 53;

namespace {

class TBackend : public TPerWorkerBaseBackend {
    public:
        TBackend(TBackendDescriptor::TRef descr)
            : TPerWorkerBaseBackend(std::move(descr))
            , Weight_(OriginalWeight())
            , HHKey_{
                MurmurHash<ui64>(Name().data(), Name().size()),
                std::numeric_limits<ui64>::max() - MurmurHash<ui64>(Name().data(), Name().size())
            }
        {}

        void DoOnCompleteRequest(const TDuration&) noexcept override {}

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

        double CalcScore(const TRequestHash& data) const noexcept {
            static constexpr size_t hashSize = sizeof(data);

            char in[hashSize];
            memcpy(&in, &data, hashSize);
            const ui64 hash = SipHash(HHKey_, in, hashSize);

            const double doubleValue = (hash & FIFTY_THREE_ONES) / FIFTY_THREE_ZEROS;

            return Weight_ * (1.0 / -Log2(doubleValue));
        }

        void SetWeight(double weight) noexcept {
            Weight_ = weight;
        }

        void RestoreWeight() noexcept {
            Weight_ = OriginalWeight();
        }

        double Weight() const noexcept {
            return Weight_;
        }

        void PrintInfo(NJson::TJsonWriter& out) const noexcept {
            out.Write("seed", HHKey_[0]);
            out.Write("enabled", Enabled());
            out.Write("weight", Weight_);
            PrintProxyInfo(out);
            PrintSuccFailRate(out);
        }

        bool operator<(const TBackend& other) const noexcept {
            return Weight_ < other.Weight_;
        }

        bool Enabled() const noexcept override {
            return Enabled_;
        }

        void SetEnabled(bool value) noexcept override {
            Enabled_ = value;
        }

    private:
        bool Enabled_ = true;
        double Weight_ = 0.;
        const HH_U64 HHKey_[2]; // seed from hash name
    };

}

BACKENDS_TLS(rendezvous_hashing) {
    TTls(const TPingerConfigUnparsed& config)
        : PingerConfig(config)
    {}

    TDeque<TBackend> Backends;
    THashMap<TString, TBackend*> NamedBackends;

    TInstant LastRequestTime;
    TPingerConfig PingerConfig;
    TDeque<TPinger> Pingers;

    TExternalWeightsFileReReader WeightsFileChecker;
};

BACKENDS_WITH_TLS(rendezvous_hashing), public TModuleParams {
private:

    class TRendezvousHashingAlgorithm : public TAlgorithmWithRemovals {
    public:
        TRendezvousHashingAlgorithm(const TStepParams &params, TTls& tls) noexcept
            : TAlgorithmWithRemovals(&params.Descr->Process())
            , Tls_(tls)
            , Hash_(*params.Hash)
        {}

        IBackend* Next() noexcept override {
            IBackend* winner = nullptr;

            double maxScore = -1.;

            // More details about rendezvous hashing here:
            // http://www.snia.org/sites/default/files/SDC15_presentations/dist_sys/Jason_Resch_New_Consistent_Hashings_Rev.pdf
            for (auto& backend: Tls_.Backends) {
                if (!IsRemoved(&backend) && backend.Weight() > 0. && backend.Enabled()) {
                    const double score = backend.CalcScore(Hash_);
                    if (score > maxScore) {
                        maxScore = score;
                        winner = &backend;
                    }
                }
            }

            return winner;
        }

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

    private:
        TTls& Tls_;
        TRequestHash Hash_;
    };

// Initialization
// --------------------------------------------------------------------------------
public:
    TBackends(const TModuleParams& mp, const TBackendsUID& uid)
        : TBackendsWithTLS(mp)
        , TModuleParams(mp)
        , BackendsId_(uid.Value)
    {
        Config->ForEach(this);
    }

private:
    THolder<TTls> DoInit(IWorkerCtl* process) noexcept override {
        auto tls = MakeHolder<TTls>(PingerConfig_);
        for (auto& i : BackendDescriptors()) {
            tls->Backends.emplace_back(i);
        }
        for (auto& i : tls->Backends) {
            tls->NamedBackends[i.Name()] = &i;
        }

        if (WeightsFilename_) {
            tls->WeightsFileChecker = TExternalWeightsFileReReader(*process, WeightsFilename_, UpdateDuration_, NAME);
        }

        if (PingerConfig_.UnparsedRequest) {
            for (auto& backend : tls->Backends) {
                tls->Pingers.emplace_back(backend, tls->LastRequestTime, &process->Executor(),
                                      tls->PingerConfig, *process, Steady_);
            }
        }
        return tls;
    }

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

    START_PARSE {
        ON_KEY("steady", Steady_) {
            return;
        }
        ON_KEY("weights_file", WeightsFilename_) {
            return;
        }

        ON_KEY("reload_duration", UpdateDuration_) {
            return;
        }

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

        Add(MakeHolder<TBackendDescriptor>(Copy(value->AsSubConfig()), key));

        return;
    } END_PARSE
// --------------------------------------------------------------------------------


// Statistics
// --------------------------------------------------------------------------------
void DumpBackends(NJson::TJsonWriter& out, const TTls& tls) const noexcept override {
        out.OpenMap();
        out.Write("id", BackendsId_);
        out.Write("update_weights_id", tls.WeightsFileChecker.LatestData().Id());
        out.OpenArray("backends");
        for (const auto& backend : tls.Backends) {
            out.OpenMap();
            backend.PrintInfo(out);
            out.CloseMap();
        }
        out.CloseArray();
        out.CloseMap();
    }
// --------------------------------------------------------------------------------

void DumpWeightsFileTags(NJson::TJsonWriter& out, const TTls& tls) const noexcept override {
    tls.WeightsFileChecker.WriteWeightsFileTag(out);
}

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

    bool IsHashing() const noexcept override {
        return true;
    }

    void UpdateWeights(const TConnDescr& descr) noexcept {
        auto& tls = GetTls(descr.Process());
        if (tls.WeightsFileChecker.UpdateWeights(descr)) {
            const auto& backendsWeights = tls.WeightsFileChecker.Entries();
            for (const auto& backend : tls.NamedBackends) {
                const auto newData = backendsWeights.find(backend.first);
                if (newData != backendsWeights.end()) {
                    backend.second->SetWeight(newData->second);
                } else {
                    backend.second->RestoreWeight();
                }
            }
        }
    }
// --------------------------------------------------------------------------------


// State
// --------------------------------------------------------------------------------
private:
    TDuration UpdateDuration_ = TDuration::Seconds(1);
    TString WeightsFilename_;
    bool Steady_ = true;
    TPingerConfigUnparsed PingerConfig_;
    size_t BackendsId_ = 0;
// --------------------------------------------------------------------------------
};

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

}  // namespace NSrvKernel::NModBalancer
