#include "hashing.h"
#include "backends_factory.h"

#include <balancer/kernel/balancing/per_worker_backend.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/module/iface.h>
#include <balancer/kernel/module/module.h>

#include <util/digest/murmur.h>

namespace NSrvKernel::NModBalancer {

namespace NHashing {

using namespace Nhashing;

TBackend::TBackend(TBackendDescriptor::TRef descr, TTls& tls)
    : TPerWorkerBaseBackend(std::move(descr))
    , Tls_(tls)
{}

void TBackend::PrintInfo(NJson::TJsonWriter& out) const noexcept {
    if (Tls_.PingersStarted) {
        out.Write("enabled", Enabled());
    }
    PrintSuccFailRate(out);
    PrintProxyInfo(out);
}

THashingAlgorithm::THashingAlgorithm(TTls& tls, const TStepParams& params, const TVector<double>& weights) noexcept
    : TAlgorithmWithRemovals(&params.Descr->Process())
    , Tls_(tls)
    , Weights_(weights)
    , Hash_(*params.Hash) {}

IBackend* THashingAlgorithm::Next() noexcept {
    if (Weights_.empty()) {
        return nullptr;
    }

    IBackend* retval = nullptr;
    for (unsigned attempt = 0; attempt < 20; ++attempt) { // 20 attempts to rehash is all that anybody would ever need
        const double hash = Min(Hash_ / double(std::numeric_limits<TRequestHash>::max()), Weights_.back());
        const size_t ret = LowerBound(Weights_.begin(), Weights_.end(), hash) - Weights_.begin();

        retval = &Tls_.Backends[ret];
        Hash_ += attempt;
        Hash_ = MurmurHash<TRequestHash>(&Hash_, sizeof(Hash_));

        if (!BackendWasSelected(retval)) {
            break;
        }
    }

    if (BackendWasSelected(retval)) {
        return nullptr;
    }

    return retval;
}

void THashingAlgorithm::FillWeights() noexcept {
    TAlgorithmWithRemovals::Reset();
    LastDisabledBackend_ = nullptr;
}

bool THashingAlgorithm::BackendWasSelected(IBackend* backend) noexcept {
    return backend == LastDisabledBackend_ || IsRemoved(backend);
}

void THashingAlgorithm::DoRemove(IBackend* backend) noexcept {
    if (LastDisabledBackend_) {
        TAlgorithmWithRemovals::RemoveSelected(LastDisabledBackend_);
    }
    LastDisabledBackend_ = backend;
}

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

}  // namespace NHashing

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

        // initialize weights and fill names
        double sum = 0.;
        for (auto& i : BackendDescriptors()) {
            sum += i->Weight();
            BackendWeights_.emplace_back(sum);
        }

        for (auto& backEndWeights : BackendWeights_) {
            backEndWeights /= sum;
        }
    }

    THolder<TTls> DoInit(IWorkerCtl* process) noexcept override {
        auto tls = MakeHolder<TTls>(PingerConfig_);
        for (auto& i : BackendDescriptors()) {
            tls->Backends.emplace_back(i, *tls.Get());
        }
        for (auto& i : tls->Backends) {
            tls->NamedBackends[i.Name()] = &i;
        }
        if (PingerConfig_.UnparsedRequest) { // do not start pingers in admin process
            for (auto& backend : tls->Backends) {
                tls->Pingers.emplace_back(backend, tls->LastRequestTime, &process->Executor(),
                tls->PingerConfig, *process, Steady_);
            }
            tls->PingersStarted = true;
        }
        return tls;
    }

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

private:
    START_PARSE {
        ON_KEY("steady", Steady_) {
            return;
        }
        if (PingerConfig_.TryConsume(key, value)) {
            return;
        }

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

        return;
    } END_PARSE

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

    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.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<NHashing::THashingAlgorithm>(GetTls(params.Descr->Process()), params, BackendWeights_);
    }

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

private:
    TVector<double> BackendWeights_;
    bool Steady_ = false;
    TPingerConfigUnparsed PingerConfig_;
    size_t BackendsId_ = 0;
};

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

}  // namespace NSrvKernel::NModBalancer
