#include "backends_factory.h"
#include "base_algorithm.h"
#include "consistent_hashing.h"

#include <balancer/kernel/balancing/per_worker_backend.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/module/iface.h>
#include <balancer/kernel/module/module.h>

#include <library/cpp/consistent_hash_ring/consistent_hash_ring.h>

#include <util/digest/murmur.h>
#include <util/generic/algorithm.h>
#include <util/generic/scope.h>

namespace NSrvKernel::NModBalancer {

namespace {
    using TBackend = TPerWorkerBaseBackend;
}

BACKENDS_TLS(consistent_hashing) {
    THashMap<TString, TBackend*> NamedBackends;
    TDeque<TBackend> Backends;
};

BACKENDS_WITH_TLS(consistent_hashing), public TModuleParams {
private:
    class TConsistentHashingAlgorithm : public TAlgorithmWithRemovals {
    public:
        TConsistentHashingAlgorithm(TBackends* parent, const TStepParams& params) noexcept
            : TAlgorithmWithRemovals(&params.Descr->Process())
            , Parent_(parent)
            , Hash_(*params.Hash)
        {}

        IBackend* Next() noexcept override {
            for (size_t attempt = 0; attempt < Parent_->Size(); ++attempt) {
                Y_DEFER {
                    Hash_ = MurmurHash<TRequestHash>(&Hash_, sizeof(TRequestHash));
                };
                const ui64 ret = Parent_->HashRing_.FindNode(Hash_);
                if (!IsRemoved(&Parent_->GetTls(*Process_).Backends[ret])) {
                    return &Parent_->GetTls(*Process_).Backends[ret];
                }
            }

            return nullptr;
        }

        IBackend* NextByName(TStringBuf name, bool /*allowZeroWeights*/) noexcept override {
            auto it = Parent_->GetTls(*Process_).NamedBackends.find(name);
            if (it != Parent_->GetTls(*Process_).NamedBackends.end()) {
                return it->second;
            }
            return nullptr;
        }
    private:
        TBackends* Parent_ = nullptr;
        TRequestHash Hash_;
    };

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

        const auto& bds = BackendDescriptors();
        if (bds.empty()) {
            return;
        }
        auto minWeight = Max(1., MinElement(bds.begin(), bds.end(), [](auto& lhs, auto& rhs) {
            return lhs->Weight() < rhs->Weight();
        })->Get()->Weight());

        for (ui64 backendNum = 0; backendNum < bds.size(); ++backendNum) {
            ui64 nodesCount = VirtualNodesCount_ * (bds[backendNum]->Weight() / minWeight);
            HashRing_.AddNode(backendNum, nodesCount);
        }
    }

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

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

    START_PARSE {
        ON_KEY("virtual_nodes", VirtualNodesCount_) {
            if (VirtualNodesCount_ == 0) {
                ythrow TConfigParseError{} << "virtual_nodes must be greater then 0";
            }
            return;
        }

        auto descr = MakeHolder<TBackendDescriptor>(Copy(value->AsSubConfig()), key);

        if (descr->Weight() < 1.0 || descr->Weight() >= Max<int>()) {
            ythrow TConfigParseError{} << "Backend weight must be in [1, INT_MAX)";
        }

        Add(std::move(descr));

        return;
    } END_PARSE

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


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


// Functionality
// --------------------------------------------------------------------------------
    THolder<IAlgorithm> ConstructAlgorithm(const TStepParams& params) noexcept override {
        return MakeHolder<TConsistentHashingAlgorithm>(this, params);
    }

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


// State
// --------------------------------------------------------------------------------
private:
    TConsistentHashRing<ui64> HashRing_;
    static constexpr ui64 DEF_VIRTUAL_NODES_COUNT = 1000;
    ui64 VirtualNodesCount_ = DEF_VIRTUAL_NODES_COUNT;
    size_t BackendsId_ = 0;
// --------------------------------------------------------------------------------
};

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

}  // namespace NSrvKernel::NModBalancer
