#include "consistent.h"
#include "config_wd.h"

#include <saas/util/types/coverage.h>

#include <library/cpp/balloc/optional/operators.h>
#include <library/cpp/containers/mh_heap/mh_heap.h>
#include <library/cpp/logger/global/global.h>

#include <util/generic/ymath.h>
#include <util/random/easy.h>

#include <tuple>

namespace {
    NFusion::TConsistentClient::TDocument::TVersion& GetVersionId(NFusion::TConsistentClient::TDocumentPtr document, ui16 versionId) {
        for (auto&& version : document->Versions) {
            if (version.VersionId != versionId) {
                continue;
            }

            return version;
        }
        FAIL_LOG("Incorrect routing");
    }

    NFusion::TConsistentClient::TDocument::TVersion& FinishFetch(NFusion::TConsistentClient::TDocument::TVersion& version) {
        CHECK_WITH_LOG(version.FetchInProgress > 0);
        version.FetchInProgress -= 1;
        return version;
    }

    class TFetchRequest: public NRealTime::TDistributorClient::IFetchHandler {
    public:
        TFetchRequest(NFusion::TConsistentClient::TDocumentPtr document, TServiceMetricPtr metric, ui16 versionId)
            : Document(document)
            , Metric(metric)
            , Start(Now())
            , VersionId(versionId)
        {}

        void OnReply(ui64 version, const TString& id, const TString& data) override {
            THolder<TFetchRequest> cleanup(this);
            Metric->OnSuccess(Now() - Start);
            auto metric = Metric;
            auto document = Document;
            auto versionId = VersionId;
            NNamedLock::GuardedRun(document->Id, [version, id, data, document, versionId, metric] {
                auto& documentVersion = FinishFetch(GetVersionId(document, versionId));
                auto fetchFailed = [&] {
                    documentVersion.FetchFailed += 1;
                    metric->OnUserFailure();
                };

                if (document->Id != id) {
                    ERROR_LOG << "Id mismatch for fetched data: " << id << " != " << document->Id << Endl;
                    fetchFailed();
                    return;
                }
                if (documentVersion.GetVersion() != version) {
                    ERROR_LOG << "Version mismatch for " << id << ": " << version << " != " << documentVersion.GetVersion() << Endl;
                    fetchFailed();
                    return;
                }
                if (version && document->FetchedVersion >= version) {
                    DEBUG_LOG << "Fetched older version for " << id << ": " << version << " <= " << document->FetchedVersion << Endl;
                    return;
                }
                if (!data) {
                    ERROR_LOG << "Fetched empty data for " << id << " from replica " << int(documentVersion.GetReplica()) << Endl;
                    fetchFailed();
                    return;
                }

                DEBUG_LOG << "Fetched data for " << id << " from replica " << int(documentVersion.GetReplica()) << Endl;
                document->Data = data;
                document->FetchedVersion = version;
                if (version == document->LastVersion) {
                    document->FetchLastFailed = false;
                }
            });
        }

        void OnError(NBus::EMessageStatus status, const TString& extendedMessage) override {
            THolder<TFetchRequest> cleanup(this);
            Metric->OnFailure();
            auto document = Document;
            auto versionId = VersionId;
            NNamedLock::GuardedRun(document->Id, [document, versionId, status, extendedMessage] {
                auto& version = FinishFetch(GetVersionId(document, versionId));
                ERROR_LOG << "Cannot fetch data for " << document->Id << " from replica " << int(version.GetReplica()) << ": "
                            << NBus::MessageStatusDescription(status) << " " << extendedMessage << Endl;
                version.FetchFailed += 1;
            });
        }
    private:
        NFusion::TConsistentClient::TDocumentPtr Document;
        TServiceMetricPtr Metric;
        const TInstant Start;
        const ui16 VersionId;
    };
}

namespace NFusion {
    class TConsistentClient::TReplicaContext::TIterator {
    public:
        typedef ui64 value_type;

    public:
        TIterator(TReplicaContext& ctx)
            : Ctx(ctx)
            , Limit(0)
            , Count(0)
            , HasDocument(false)
        {
        }

        bool Valid() const {
            return HasDocument;
        }
        void Next() {
            if (Count < Limit) {
                HasDocument = Ctx.GetNextVersion(Document);
                Count += HasDocument ? 1 : 0;
            } else {
                HasDocument = false;
            }
        }
        const TDocument::TVersion& GetDoc() const {
            return Document;
        }
        ui64 Current() const {
            return GetDoc().GetTimestamp().GetValue();
        }
        void operator++() {
            Next();
        }
        void Restart() {
        }
        void SetLimit(size_t limit) {
            Limit = limit;
            if (HasDocument) {
                Count = 1;
            } else {
                Count = 0;
                Next();
            }
        }
    private:
        TReplicaContext& Ctx;
        TDocument::TVersion Document;
        size_t Limit;
        size_t Count;
        bool HasDocument;
    };
}

NFusion::TConsistentClientReplicas NFusion::SplitReplicas(const NRealTime::TDistributors& distributors) {
    NFusion::TConsistentClientReplicas replicas;
    if (distributors.empty()) {
        return replicas;
    }

    NUtil::TCoverage<NBus::TBusKey, NRealTime::TDistributor> coverage(NBus::YBUS_KEYMIN, NBus::YBUS_KEYMAX);

    for (auto&& distributor : distributors) {
        NUtil::TInterval<NBus::TBusKey> interval(distributor.KeyMin, distributor.KeyMax);
        coverage.AddHandler(interval, distributor);
    }

    size_t numberOfReplicas = Max<size_t>();
    ForEach(coverage.begin(), coverage.end(), [&](const std::pair<TInterval<NBus::TBusKey>, TVector<NRealTime::TDistributor> >& i) {
        numberOfReplicas = Min<size_t>(numberOfReplicas, i.second.size());
    });
    INFO_LOG << "Distributors splitted into " << numberOfReplicas << " replicas" << Endl;

    replicas.resize(numberOfReplicas);
    for (auto&& interval : coverage) {
        for (size_t i = 0; i < interval.second.size(); ++i) {
            NRealTime::TDistributor d(interval.second[i]);
            d.KeyMin = interval.first.GetMin();
            d.KeyMax = interval.first.GetMax();

            const ui64 id = Min(i, numberOfReplicas - 1);
            replicas[id].push_back(d);
            replicas[id].Id = id;
        }
    }
    for (size_t i = 0; i < numberOfReplicas; ++i) {
        INFO_LOG << "Replica " << i << ":" << Endl;
        for (auto&& d : replicas[i]) {
            INFO_LOG << ToString(d) << Endl;
        }
    }
    return replicas;
}

bool NFusion::TConsistentClient::SetPreferredDistributor(const char* /*hostName*/, int /*port*/) {
    FAIL_LOG("Consistent distributor client does not support preferred distributor");
    return false;
}

void NFusion::TConsistentClient::SetKeyRange(ui64 start, ui64 end) {
    GuardedRun([=](NRealTime::IDistributorClient& client) {
        client.SetKeyRange(start, end);
    });
}

void NFusion::TConsistentClient::SetMinRank(double minRank) {
    MinRank = minRank;
    GuardedRun([=](NRealTime::IDistributorClient& client) {
        client.SetMinRank(minRank);
    });
}

void NFusion::TConsistentClient::SetAttribute(const TString& attr) {
    GuardedRun([=](NRealTime::IDistributorClient& client) {
        client.SetAttribute(attr);
    });
}

void NFusion::TConsistentClient::SetMaxChunks(size_t maxChunks) {
    GuardedRun([=](NRealTime::IDistributorClient& client) {
        client.SetMaxChunks(maxChunks);
    });
}

void NFusion::TConsistentClient::SetIdleTime(time_t idleTime) {
    GuardedRun([=](NRealTime::IDistributorClient& client) {
        client.SetIdleTime(idleTime);
    });
}

void NFusion::TConsistentClient::SetStream(const TString& stream) {
    GuardedRun([=](NRealTime::IDistributorClient& client) {
        client.SetStream(stream);
    });
}

void NFusion::TConsistentClient::SetAge(int secondsAgo /*= 0*/) {

    auto distributorGuard = Guard(DistributorLock);
    auto storageGuard = Guard(StorageLock);
    for (auto&& ctx : Contexts) {
        CHECK_WITH_LOG(ctx.second);
        ctx.second->Run([=](NRealTime::IDistributorClient& client) {
            client.SetAge(secondsAgo);
        });
    }
    Reset();
}

void NFusion::TConsistentClient::SetClientActive(bool value) {
    GuardedRun([=](NRealTime::TDistributorClient& client) {
        client.SetClientActive(value);
    });
}

void NFusion::TConsistentClient::GetClientInfos(IClientInfoHandler* handler) {
    GuardedRun([=](NRealTime::TDistributorClient& client) {
        client.GetClientInfos(handler);
    });
}

bool NFusion::TConsistentClient::IsExhausted() const {
    for (auto&& ctx : Contexts) {
        if (!ctx.second->IsExhausted())
            return false;
    }
    return true;
}

TString NFusion::TConsistentClient::GetStatusMsg() const {
    TString result;
    for (auto&& ctx : Contexts) {
        result += ctx.second->GetStatusMsg();
        result += "\t";
    }
    return result;
}

NFusion::TConsistentClient::TStatus NFusion::TConsistentClient::GetStatus() const {
    NFusion::TConsistentClient::TStatus result;
    for (auto&& ctx : Contexts) {
        result.push_back(ctx.second->GetStatus());
    }
    return result;
}

bool NFusion::TConsistentClient::GetNextDoc(::google::protobuf::Message& doc, NRealTime::TDistributorClient::TSignature& signature, TDocument::TSourceInfo& sourceInfo) {
    if (!GetNextRawRecord(signature, sourceInfo)) {
        return false;
    }

    doc.Clear();
    if (!doc.ParseFromString(signature.Data)) {
        Metric.DataParsingError.Inc();
        return false;
    }

    return true;
}

bool NFusion::TConsistentClient::GetNextRawRecord(NRealTime::TDistributorClient::TSignature& signature, TDocument::TSourceInfo& sourceInfo) {
    if (MinRank.IsPaused())
        return false;

    TDocument::TVersion version;
    bool result = GetVersion(version, sourceInfo);
    if (!result) {
        StartUpdater();
    }

    signature = version;
    return result;
}

bool NFusion::TConsistentClient::GetNextDocAndTime(::google::protobuf::Message &doc,
                                                      TInstant& docTime,
                                                      TInstant& nowTime,
                                                      ui64 *key /*= nullptr*/,
                                                      TVector<TString>* attrs /*= nullptr*/,
                                                      ui64* distTime /*= nullptr*/)
{
    TString data;
    if (!GetNextRawRecordAndTime(data, docTime, nowTime, key, attrs, distTime)) {
        return false;
    }

    doc.Clear();
    if (!doc.ParseFromString(data)) {
        Metric.DataParsingError.Inc();
        return false;
    }
    return true;
}

bool NFusion::TConsistentClient::GetNextRawRecordAndTime(TString& data,
                                                      TInstant& docTime,
                                                      TInstant& nowTime,
                                                      ui64 *key /*= nullptr*/,
                                                      TVector<TString>* attrs /*= nullptr*/,
                                                      ui64* distTime /*= nullptr*/)
{
    auto guard = Guard(DistributorLock);

    NRealTime::TDistributorClient::TSignature version;
    TDocument::TSourceInfo info;
    bool result = GetNextRawRecord(version, info);
    if (result) {
        nowTime = Now();
        docTime = Now() - TDuration::Seconds(version.Age);
        data    = version.Data;
        if (!data) {
            WARNING_LOG << "Empty data for " << version.Id << Endl;
        }
        if (key) {
            *key = version.Key;
        }
        if (attrs) {
            *attrs = std::move(version.Attrs);
        }
        if (distTime) {
            *distTime = version.DistTime;
        }
    }
    return result;
}

NFusion::TConsistentClient::TConsistentClient(const TConsistentClientReplicas& replicas, const TConsistentClientOptions& option)
    : Options(option)
    , Records(option.Capacity)
    , Head(0)
    , Reader(0)
    , Metric(option.MetricsPrefix)
    , StopFlag(true)
{
    ui32 maxPriority = 0;
    for (auto&& replica : replicas) {
        maxPriority = Max(maxPriority, replica.Priority);
        if (Priorities.size() <= replica.Id) {
            Priorities.resize(replica.Id + 1);
            Shuffling.resize(replica.Id + 1);
        }
        Priorities[replica.Id] = replica.Priority;
        Shuffling[replica.Id] = Random();

        Contexts[replica.Id] = MakeAtomicShared<TReplicaContext>(replica, Options);
    }

    if (Options.Metrics) {
        Metric.Register(*Options.Metrics);
    }

    HashedIds.reserve(Options.Capacity);

    ReassignReplicas();
}

NFusion::TConsistentClient::~TConsistentClient() {
    if (Options.Metrics) {
        Metric.Deregister(*Options.Metrics);
    }

    StopUpdater();
}

void NFusion::TConsistentClient::StopUpdater() {
    if (!StopFlag) {
        StopFlag = true;

        CHECK_WITH_LOG(Updater);
        Updater->Join();
        Updater.Destroy();
        INFO_LOG << "Stopped updater thread for " << Options.MetricsNamePrefix << Endl;
    }
}

void NFusion::TConsistentClient::StartUpdater() {
    if (StopFlag) {
        StopFlag = false;
        TThread::TParams updaterThreadParams(&UpdaterProc, this);
        updaterThreadParams.SetName("consistent-client-updater");

        CHECK_WITH_LOG(!Updater);
        Updater = MakeHolder<TThread>(updaterThreadParams);
        Updater->Start();
        INFO_LOG << "Started updater thread for " << Options.MetricsNamePrefix << Endl;
    }
}

void* NFusion::TConsistentClient::UpdaterProc(void* object) {
    ThreadDisableBalloc();

    auto this_ = reinterpret_cast<TConsistentClient*>(object);
    TInstant switchUpdateTime = Now();
    while (!this_->StopFlag) {
        const ui32 preload = Min<ui32>(this_->Vacant() * this_->Options.PreloadFactor, this_->Options.MaxPreloadCount);
        const ui32 prefetch = Min<ui32>(this_->Ahead() * this_->Options.PrefetchFactor, this_->Options.MaxPrefetchCount);
        this_->Preload(preload);
        this_->Prefetch(prefetch);

        Sleep(this_->Options.UpdateInterval);
        const TInstant now = Now();
        if (this_->Options.DynamicSwitch && (now > switchUpdateTime)) {
            this_->ReassignReplicas();
            switchUpdateTime = now + this_->Options.DynamicSwitchInterval;
        }
    }

    return nullptr;
}

void NFusion::TConsistentClient::Preload(size_t count) {
    if (!count) {
        return;
    }

    TVector<TReplicaContext::TIterator*> iterators;
    for (auto&& ctx : Contexts) {
        iterators.push_back(ctx.second->GetIterator(count));
    }

    size_t c = 0;
    MultHitHeap<TReplicaContext::TIterator> merger(&iterators[0], iterators.size());
    for (merger.Restart(); merger.Valid() && c <= count; ++merger) {
        const TDocument::TVersion& doc = merger.TopIter()->GetDoc();

        if (!MinRank.IsAllowedRank(doc.Rank)) {
            //This happens only if MinRank handle is changed when the client have already preloaded a document
            DEBUG_LOG << "Discarded preloaded " << doc.Id << " by rank" << Endl;
            continue;
        }

        c += AddVersion(doc);
    }
}

void NFusion::TConsistentClient::Prefetch(size_t count) {
    if (!count) {
        return;
    }

    size_t reader = Reader;
    for (size_t i = 0; i < count; ++i) {
        TGuardedDocument d;
        with_lock (StorageLock) {
            d = Guarded(GetDocument(reader + i));
        }
        if (d && NeedFetch(d) && CanFetch(d) && CheckFilterRank(d)) {
            Metric.PrefetchMiss.Inc();
            Fetch(d);
        }
    }
}

void NFusion::TConsistentClient::ReassignReplicas() {
    struct TReplicaInfo {
        size_t Index;
        float FailRate;
        ui32 Age;
        ui16 Id;
        bool Banned;
    };

    const size_t count = Contexts.size();

    ui32 maxPriority = 0;
    ui32 maxPriorityCount = 0;
    TVector<TReplicaInfo> infos;
    for (auto&& ctx: Contexts) {
        auto context = ctx.second;
        CHECK_WITH_LOG(context);

        const TReplicaInfo info = {
            ctx.first,
            context->GetFailRate(),
            context->GetAge(),
            context->GetId(),
            false
        };
        infos.push_back(info);

        maxPriority = Max(maxPriority, Priorities[context->GetId()]);
    }

    float averageFailRate = 0;
    averageFailRate = Accumulate(infos.begin(), infos.end(), averageFailRate, [](float value, const TReplicaInfo& e) {
        return value + e.FailRate;
    }) / count;
    float averageAge = 0;
    averageAge = Accumulate(infos.begin(), infos.end(), averageAge, [](float value, const TReplicaInfo& e) {
        return value + e.Age;
    }) / count;

    for (size_t i = 0; i < count; ++i) {
        if ((infos[i].FailRate > averageFailRate * (1 + Options.FailRateSwitchFactor)) ||
            (infos[i].Age > averageAge * (1 + Options.AgeSwitchFactor)))
        {
            infos[i].Banned = true;
        }

        if (Priorities[infos[i].Id] == maxPriority) {
            maxPriorityCount += 1;
        }
    }

    Sort(infos.begin(), infos.end(), [this](const TReplicaInfo& l, const TReplicaInfo& r) {
        return std::make_tuple(Priorities[l.Id], Shuffling[l.Id]) > std::make_tuple(Priorities[r.Id], Shuffling[r.Id]);
    });

    ui32 primaryReplicasCount = 0;
    for (size_t i = 0; i < count; ++i) {
        auto replica = Contexts[infos[i].Index];
        CHECK_WITH_LOG(replica);

        const bool previous = replica->IsPrimary();
        if (infos[i].Banned || primaryReplicasCount >= Min(Options.MaxPrimaryReplicas, maxPriorityCount)) {
            replica->SetSecondary();
        } else {
            replica->SetPrimary();
            primaryReplicasCount++;
        }
        const bool current = replica->IsPrimary();

        if (current != previous) {
            INFO_LOG << "Replica " << infos[i].Index << " role: " << previous << " -> " << current << Endl;
        }
    }
}

void NFusion::TConsistentClient::Reset() {
    HashedIds.clear(Options.Capacity);
    Records.clear();
    Records.resize(Options.Capacity);
    Head = 0;
    Reader = 0;
}

template <class F>
void NFusion::TConsistentClient::GuardedRun(F f) {
    auto guard = Guard(DistributorLock);
    for (auto&& ctx : Contexts) {
        CHECK_WITH_LOG(ctx.second);
        ctx.second->Run(f);
    }
}

bool NFusion::TConsistentClient::GetVersion(TDocument::TVersion& doc, TDocument::TSourceInfo& info) {
    Metric.Requested.Inc();

    while (true) {

    auto guard = Guard(StorageLock);
    if (Reader >= Head) {
        return false;
    }

    CHECK_WITH_LOG(GetDocument(Reader));
    auto d = Guarded(GetDocument(Reader));
    if (!d) {
        return false;
    }

    if (!CheckFilterRank(d)) {
        DiscardFiltered(d);
        AdvanceReader();
        continue;
    }

    if (NeedFetch(d)) {
        DEBUG_LOG << "Need fetch for " << d->Id << " " << d->FetchedVersion << " -> " << d->LastVersion << Endl;
        if (!CanFetch(d)) {
            RescheduleOrDiscard(d);
            AdvanceReader();
            continue;
        }

        Metric.FetchMiss.Inc();
        Fetch(d);
        return false;
    }

    doc = GetFetched(d);
    info = GetSourceInfo(d);
    ClearFetched(d);
    AdvanceReader();

    Metric.Acquired.Inc();
    return true;

    } // while (true)

    // unreachable
    return false;
}

ui32 NFusion::TConsistentClient::AddVersion(const TDocument::TVersion& doc) {
    const TString& id = doc.GetId();
    if (!id) {
        Metric.EmptySignature.Inc();
        DEBUG_LOG << "Skipping document due to empty signature" << Endl;
        return 0;
    }

    auto guard = Guard(StorageLock);
    Metric.Added.Inc();
    DEBUG_LOG << "Adding document " << id << ":" << doc.GetVersion() << ":" << doc.HasData() << "@"
                << int(doc.GetReplica()) << ":" << doc.GetKey() << ":" << doc.GetTimestamp().GetValue() << Endl;

    if (ShouldInsert(doc)) {
        return AddDocument(MakeAtomicShared<TDocument>(doc));
    }

    DEBUG_LOG << "Appending document " << id << " to position " << GetPos(id) << Endl;
    auto record = GetDocument(id);
    NNamedLock::GuardedRun(id, [record, doc, id, this] {
        CHECK_WITH_LOG(record);
        CHECK_WITH_LOG(!record->Versions.empty());
        Y_ASSERT(IsSorted(record));

        TDocument::TVersions& versions = record->Versions;
        if (versions.size() >= Max<ui16>()) {
            ERROR_LOG << "Max versions per document limit reached for " << id << Endl;
            return;
        }

        const ui64 before = versions.front().GetVersion();
        versions.push_back(doc);
        versions.back().VersionId = versions.size() - 1;
        Sort(versions.begin(), versions.end(), [&](const TDocument::TVersion& l, const TDocument::TVersion& r) {
            auto left = std::make_tuple(l.GetVersion(), Priorities[l.GetReplica()], Shuffling[l.GetReplica()]);
            auto right = std::make_tuple(r.GetVersion(), Priorities[r.GetReplica()], Shuffling[r.GetReplica()]);
            return left > right;
        });
        const ui64 after = versions.front().GetVersion();

        DEBUG_LOG << "Version of document " << id << " " << before << " -> " << after << Endl;
        CHECK_WITH_LOG(after >= before);

        if (after != before) {
            Y_ASSERT(record->LastVersion == before);
            record->LastVersion = after;
        }
        if (doc.HasData() && (record->Data.empty() || doc.GetVersion() > record->FetchedVersion)) {
            record->Data = doc.GetData();
            record->FetchedVersion = doc.GetVersion();
        }

        record->FetchAllFailed = false;
        if (doc.GetVersion() >= record->LastVersion) {
            record->FetchLastFailed = false;
        }
    });

    return 0;
}

ui32 NFusion::TConsistentClient::AddDocument(TDocumentPtr doc) {
    const size_t pos = AdvanceHead();
    DEBUG_LOG << "Inserting document " << doc->Id << " to position " << pos << Endl;

    auto old = GetDocument(pos);
    if (old && old->Versions.size()) {
        TDocument::TVersions& versions = old->Versions;
        if (IsRead(GetPos(old->Id))) {
            HashedIds.erase(old->Id);
        }
        if (versions.size() == 1) {
            Metric.Unique.Inc();
        }
    }

    Records[GetIndex(pos)] = doc;
    HashedIds[doc->Id] = pos;

    return 1;
}

NFusion::TConsistentClient::TDocumentPtr NFusion::TConsistentClient::GetDocument(const TString& id) const {
    auto p = HashedIds.find(id);
    if (p == HashedIds.end()) {
        return nullptr;
    }

    return Records[GetIndex(p->second)];
}

NFusion::TConsistentClient::TDocumentPtr NFusion::TConsistentClient::GetDocument(size_t pos) const {
    return Records[GetIndex(pos)];
}

size_t NFusion::TConsistentClient::GetPos(const TString& id) const {
    auto p = HashedIds.find(id);
    return p != HashedIds.end() ? p->second : size_t(-1);
}

NFusion::TConsistentClient::TGuardedDocument NFusion::TConsistentClient::Guarded(TDocumentPtr doc) {
    if (!doc) {
        return TGuardedDocument();
    }

    auto lock = NNamedLock::TryAcquireLock(doc->Id);
    if (!lock) {
        Metric.LockMiss.Inc();
        return TGuardedDocument();
    }

    return TGuardedDocument(doc, lock);
}

bool NFusion::TConsistentClient::ShouldInsert(const TDocument::TVersion& doc) const {
    auto record = GetDocument(doc.GetId());
    if (!record) {
        return true;
    }

    if (record->LastVersion >= doc.GetVersion()) {
        return false;
    }

    if (!IsRead(GetPos(doc.GetId()))) {
        return false;
    }

    return true;
}

bool NFusion::TConsistentClient::IsSorted(TDocumentPtr doc) const {
    TSet<ui16> versionIds;
    for (size_t i = 0; i < doc->Versions.size() - 1; ++i) {
        if (doc->Versions[i].GetVersion() < doc->Versions[i + 1].GetVersion()) {
            ERROR_LOG << "Unsorted versions vector" << Endl;
            return false;
        }
        if (!versionIds.insert(doc->Versions[i].VersionId).second) {
            ERROR_LOG << "Duplicated VersionId " << int(doc->Versions[i].VersionId) << Endl;
            return false;
        }
    }
    return true;
}

void NFusion::TConsistentClient::Fetch(TDocumentPtr doc) {
    NNamedLock::GuardedRun(doc->Id, [doc, this] {
        CHECK_WITH_LOG(doc->Versions.size());
        ui32 fetchFailedCount = 0;
        for (auto&& version : doc->Versions) {
            if (version.FetchInProgress) {
                continue;
            }
            if (version.FetchFailed > doc->FetchFailedThreshold) {
                fetchFailedCount += 1;
                continue;
            }

            if (version.GetVersion() < doc->LastVersion) {
                doc->FetchLastFailed = true;
            }

            Contexts[version.GetReplica()]->Fetch(doc, version.VersionId);
            return;
        }

        if (fetchFailedCount == doc->Versions.size()) {
            doc->FetchAllFailed = true;
            doc->FetchLastFailed = true;
        }
    });
}

void NFusion::TConsistentClient::RescheduleOrDiscard(TGuardedDocument doc) {
    if (doc->FetchFailedThreshold < Options.MaxFetchFails) {
        doc->FetchFailedThreshold += 1;
        doc->FetchAllFailed = false;
        doc->FetchLastFailed = false;

        DEBUG_LOG << "Rescheduling " << doc->Id << Endl;
        AddDocument(doc);
        Metric.Rescheduled.Inc();
    } else {
        ERROR_LOG << "Cannot fetch " << doc->Id << Endl;
        Metric.FetchError.Inc();
    }
}

bool NFusion::TConsistentClient::CanFetch(TGuardedDocument doc) const {
    return !doc->FetchAllFailed;
}

bool NFusion::TConsistentClient::NeedFetch(TGuardedDocument doc) const {
    return doc->Data.empty() || ((doc->LastVersion > doc->FetchedVersion) && !doc->FetchLastFailed);
}

void NFusion::TConsistentClient::ClearFetched(TGuardedDocument doc) const {
    doc->Data.clear();
    for (auto&& v : doc->Versions) {
        v.Data.clear();
    }
}

bool NFusion::TConsistentClient::CheckFilterRank(TGuardedDocument doc) const {
    return doc->Versions.empty() || MinRank.IsAllowedRank(doc->Versions.front().Rank);
}

void NFusion::TConsistentClient::DiscardFiltered(TGuardedDocument doc) const {
    //This happens only if MinRank handle is changed when the client have already preloaded or prefetched a document
    //(normally the distributor returns only the matching records)
    DEBUG_LOG << "Discarded preloaded " << doc->Id << " by rank" << Endl;
    ClearFetched(doc);
}

NFusion::TConsistentClient::TDocument::TVersion NFusion::TConsistentClient::GetFetched(TGuardedDocument doc) const {
    TDocument::TVersion result;
    CHECK_WITH_LOG(doc->Versions.size());
    CHECK_WITH_LOG(!doc->Data.empty());
    result = doc->Versions.front();
    result.Data = doc->Data;
    return result;
}

NFusion::TConsistentClient::TDocument::TSourceInfo NFusion::TConsistentClient::GetSourceInfo(TGuardedDocument doc) const {
    TDocument::TSourceInfo result;
    for (auto&& v : doc->Versions) {
        if (v.GetVersion() != doc->FetchedVersion) {
            break;
        }

        result.Replicas.push_back(v.GetReplica());
    }
    return result;
}

bool NFusion::TConsistentClient::TReplicaContext::GetNextVersion(TDocument::TVersion& version) {
    if (!Client->GetNextDocSignature(version)) {
        SignatureMetric->Failures.Inc();
        return false;
    } else {
        SignatureMetric->Acquires.Inc();
    }

    if (version.Id.empty() && version.HasData() && Options.SignatureProvider) {
        auto extracted = Options.SignatureProvider->Extract(version.GetData());
        version.Id = extracted.Id;
        version.Version = extracted.Version;
    }

    Age = version.Age;
    version.Replica = Replica.Id;
    version.Timestamp = Now() - TDuration::Seconds(version.Age);
    return true;
}

NFusion::TConsistentClient::TReplicaContext::TIterator* NFusion::TConsistentClient::TReplicaContext::GetIterator(size_t limit) {
    Y_ASSERT(Iterator);
    Iterator->SetLimit(limit);
    return Iterator.Get();
}

void NFusion::TConsistentClient::TReplicaContext::Fetch(TDocumentPtr doc, ui16 versionId) {
    auto& version = GetVersionId(doc, versionId);
    VERIFY_WITH_LOG(version.GetReplica() == Replica.Id, "Incorrect replica routing");
    auto req = MakeHolder<TFetchRequest>(doc, FetchMetric, versionId);
    FetchMetric->OnIncoming();
    if (!Client->AsyncFetch(version, req.Get())) {
        ERROR_LOG << "Cannot send fetch request for " << doc->Id << " to replica " << int(Replica.Id) << Endl;
        FetchMetric->OnFailure();
        version.FetchFailed += 1;
    } else {
        DEBUG_LOG << "Fetching " << doc->Id << " from replica " << int(Replica.Id) << Endl;
        Y_UNUSED(req.Release());
        version.FetchInProgress += 1;
    }
}

void NFusion::TConsistentClient::TReplicaContext::SetPrimary() {
    Client->SetSignaturesOnly(false);
    Primary = true;
}

void NFusion::TConsistentClient::TReplicaContext::SetSecondary() {
    Client->SetSignaturesOnly(true);
    Primary = false;
}

float NFusion::TConsistentClient::TReplicaContext::GetFailRate() const {
    if (!Options.Metrics) {
        return 0;
    }

    auto acquires = GetMetricResult(SignatureMetric->Acquires, *Options.Metrics);
    auto failures = GetMetricResult(SignatureMetric->Failures, *Options.Metrics);
    if (!acquires.Defined() || !failures.Defined()) {
        return 0;
    }
    return (Abs(acquires->AvgRate + failures->AvgRate) > 0.01f) ? (failures->AvgRate / (acquires->AvgRate + failures->AvgRate)) : 0;
}

NFusion::TConsistentClient::TReplicaStatus NFusion::TConsistentClient::TReplicaContext::GetStatus() const {
    NFusion::TConsistentClient::TReplicaStatus result;
    result.ClientStatus = Client->GetStatus();
    result.Age = GetAge();
    result.Id = GetId();
    result.Primary = IsPrimary();
    return result;
}

template <class F>
void NFusion::TConsistentClient::TReplicaContext::Run(F f) {
    CHECK_WITH_LOG(Client);
    f(*Client);
}

NFusion::TConsistentClient::TReplicaContext::TReplicaContext(const TConsistentClientReplica& replica, const TConsistentClientOptions& option)
    : Replica(replica)
    , Options(option)
    , Age(0)
    , Primary(true)
{
    NRealTime::TDistributorClient::TOptions opts(Options);
    opts.MetricsNamePrefix = Options.MetricsPrefix + "_" + ToString(Replica.Id) + "_";

    FetchMetric = MakeAtomicShared<TServiceMetric>("Fetch", opts.MetricsNamePrefix);
    FetchMetric->Register(Options.Metrics);

    SignatureMetric = MakeHolder<TSignatureMetric>(opts.MetricsNamePrefix);
    SignatureMetric->Register(Options.Metrics);

    Client = MakeHolder<NRealTime::TDistributorClient>(replica, opts);
    Client->SetClientId(Options.ClientId);
    Client->SetMaxChunks(Options.MaxResponseChunks);
    Client->SetMaxRecords(Options.MaxPreloadCount);
    if (replica.StartAge) {
        Client->SetAge(replica.StartAge);
    }

    Iterator = MakeHolder<TReplicaContext::TIterator>(*this);
}

NFusion::TConsistentClient::TReplicaContext::~TReplicaContext() {
    SignatureMetric->Deregister(Options.Metrics);
    FetchMetric->Deregister(Options.Metrics);
}
