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

#include <saas/api/action.h>
#include <saas/rtyserver/common/common_messages.h>
#include <saas/rtyserver/common/stream_messages.h>
#include <saas/rtyserver/docfetcher/library/stream_basics.h>
#include <saas/rtyserver/docfetcher/library/types.h>

#include <saas/protos/reply.pb.h>

#include <library/cpp/watchdog/lib/factory.h>
#include <library/cpp/linear_regression/welford.h>
#include <library/cpp/logger/global/global.h>

#include <util/random/random.h>
#include <util/system/fs.h>
#include <library/cpp/deprecated/atomic/atomic.h>

namespace {
    class TSignatureProvider final: public NFusion::ISignatureProvider {
    public:
        TSignature Extract(const TString& data) const override {
            TSignature result;
            NRTYServer::TMessage message;
            if (message.ParseFromString(data)) {
                result.Id = NSaas::GetActionDescription(message);
                result.Version = NSaas::GetActionVersion(message);
            }
            return result;
        }
    };

    template <class TClient>
    class TBaseDistributorIterator: public NFusion::TDistributorStream::IIterator {
    protected:
        void SetCommonOptions(const NFusion::TStreamConfig& sc) {
            CHECK_WITH_LOG(Client);
            Client->SetStream(sc.DistributorStream);
            Client->SetIdleTime(1);
            Client->SetKeyRange(sc.ShardMin, sc.ShardMax);

            for (auto&& attribute: SplitString(sc.DistributorAttributes, ",")) {
                Client->SetAttribute(attribute);
            }
        }

        void SetActive(bool value) override {
            Client->SetClientActive(value);
        }

        void SetAge(ui32 value) override {
            Client->SetAge(value);
        }

        void SetMinRank(const NFusion::TMinRankCondition& minRank) override {
            Client->SetMinRank(minRank.GetValue());
        }

        NJson::TJsonValue GetStatus() const final {
            return Serialize(Client->GetStatus());
        }

        bool IsExhausted() const override {
            return Client->IsExhausted();
        }

        ui32 GetStreamTimestamp(const NRTYServer::TTimestampSnapshot& snapshot, NRTYServer::TStreamId stream) const {
            const auto element = snapshot.find(stream);
            if (element == snapshot.end()) {
                return 0;
            }

            const ui64 raw = SafeIntegerCast<ui64>(element->second.MaxValue);
            if (raw > Max<ui32>()) {
                ERROR_LOG << "Incorrect timestamp for stream " << stream << ": " << raw << Endl;
                return 0;
            }

            return static_cast<ui32>(raw);
        }

        NJson::TJsonValue Serialize(const NRealTime::TDistributorClient::TStatus& status) const {
            NJson::TJsonValue result;
            for (auto&& shard : status) {
                const TString& range = ToString(shard.ShardMin) + "-" + ToString(shard.ShardMax);
                NJson::TJsonValue& v = result[range];
                if (shard.Host) { v["Host"] = shard.Host; }
                if (shard.RequestedAge) { v["RequestedAge"] = shard.RequestedAge; }
                if (shard.RequestedCookie) { v["RequestedCookie"] = shard.RequestedCookie; }
                if (shard.ResponseMessage) { v["ResponseMessage"] = shard.ResponseMessage; }
                v["ResponseAge"] = shard.ResponseAge;
                v["ResponseCookie"] = shard.ResponseCookie;
                v["ResponseStatus"] = NRealTime::TIndexedDocResponse_EStatus_Name(shard.ResponseStatus);
                v["DocsCount"] = shard.DocsCount;
                v["Exhausted"] = shard.Exhausted;
            }
            return result;
        }

        NJson::TJsonValue Serialize(const NFusion::TConsistentClient::TStatus& status) const {
            NJson::TJsonValue result;
            for (auto&& replica : status) {
                NJson::TJsonValue& v = result[ToString(replica.Id)];
                v["Status"] = Serialize(replica.ClientStatus);
                v["Age"] = replica.Age;
                v["Primary"] = replica.Primary;
            }
            return result;
        }

    protected:
        THolder<TClient> Client;
        NRTYServer::TMessage Message;
    };

    class TLegacyDistributorClientFilter : public NRealTime::TDistributorClient {
    public:
        using NRealTime::TDistributorClient::TDistributorClient;

        void SetMinRank(double minRank) override {
            RankCondition = minRank;
        }
    protected:
        NRealTime::TRecord::TPtr GetNextRecord() override {
            NRealTime::TRecord::TPtr record;

            if (RankCondition.IsPaused())
                return record;

            do {
                record = NRealTime::TDistributorClient::GetNextRecord();
                if (!record) {
                    return nullptr;
                }
            } while (!RankCondition.IsAllowedRank(record->ProtoRecord.GetRank()));
            return record;
        }
    protected:
        NFusion::TMinRankCondition RankCondition;
    };

    class TLegacyDistributorIterator: public TBaseDistributorIterator<TLegacyDistributorClientFilter> {
    public:
        TLegacyDistributorIterator(const NFusion::TStreamConfig& sc, const NRTYServer::TTimestampSnapshot& snapshot)
            : StreamId(sc.StreamId)
        {
            Client = MakeHolder<TLegacyDistributorClientFilter>(
                NRealTime::SerializeDistributorString(sc.DistributorServers),
                /*bus=*/ nullptr,
                /*blocking=*/ true,
                /*proto=*/ nullptr,
                /*requireIpVersion=*/ NBus::EIP_VERSION_ANY,
                /*preferIpVersion=*/ NBus::EIP_VERSION_6,
                sc.UseCompression
            );
            Client->SetSwitchOverlap(sc.OverlapAge);
            const ui32 mainStream = NFusion::GetSubStream(sc.StreamId, NFusion::SubStreamDocumentAux);
            Client->SetAge(Min<ui32>(Seconds() + sc.OverlapAge - GetStreamTimestamp(snapshot, mainStream), sc.MaxAgeToGetSec));

            for (auto&& d : sc.PreferredDistributor) {
                const bool result = Client->SetPreferredDistributor(d.Host.data(), d.Port);
                NOTICE_LOG << sc.Name << ": " << (result ? "" : "unable to ") << "set distributor " << d << " as preferred" << Endl;
            }

            SetCommonOptions(sc);
        }

        NFusion::TActionPtr GetDoc() final {
            TInstant now;
            TInstant diststamp;
            if (!Client->GetNextDocAndTime(Message, diststamp, now)) {
                return nullptr;
            }
            //note: Message.Document.FilterRank is not filled (because legacy)

            auto result = MakeAtomicShared<NSaas::TAction>(Message);
            result->GetDocument().SetStreamId(NFusion::GetSubStream(StreamId, NFusion::SubStreamDocumentAux));
            result->SetReceiveTimestamp(diststamp);
            return result;
        }

    private:
        const ui32 StreamId;
    };

    class TConsistentDistributorIterator: public TBaseDistributorIterator<NFusion::TConsistentClient> {
    public:
        TConsistentDistributorIterator(
            const NFusion::TStreamConfig& sc,
            const NFusion::TConsistentClientReplicas& replicas,
            const NFusion::TConsistentClientOptions& options,
            const TInstant& minTimestamp,
            const NRTYServer::TTimestampSnapshot& snapshot
        )
            : Replicas(replicas)
            , Options(options)
            , StreamId(sc.StreamId)
        {
            const ui32 mainStream = NFusion::GetSubStream(sc.StreamId, NFusion::SubStreamDocumentAux);
            const ui32 mainTimestamp = GetStreamTimestamp(snapshot, mainStream);
            for (auto&& replica: Replicas) {
                const ui8  replicaId = replica.Id;
                const ui32 substream = NFusion::GetSubStream(sc.StreamId, replicaId);
                const ui32 substreamTimestamp = GetStreamTimestamp(snapshot, substream);
                const ui32 timestamp = Max<ui32>(substreamTimestamp != 0 ? substreamTimestamp : mainTimestamp, minTimestamp.Seconds());

                ui32 overlapTimestamp = Seconds() + sc.OverlapAge;
                CHECK_WITH_LOG(overlapTimestamp > timestamp);
                replica.StartAge = Min<ui32>(overlapTimestamp - timestamp, sc.MaxAgeToGetSec);
                INFO_LOG << sc.Name << ": initializing replica " << int(replica.Id) << " with age " << replica.StartAge << Endl;
            }

            Options.RequireIpVersion = NBus::EIP_VERSION_ANY;
            Options.PreferIpVersion = NBus::EIP_VERSION_6;
            Options.UseCompression = sc.UseCompression;
            Options.ClientId = sc.ClientId;
            Options.Metrics = &GetGlobalMetrics();
            Options.MetricsPrefix = GetMetricsPrefix() + JoinMetricName(sc.Name, "ConsistentClient");
            Options.SignatureProvider = Singleton<TSignatureProvider>();

            Client = MakeHolder<NFusion::TConsistentClient>(Replicas, Options);
            SetCommonOptions(sc);
        }

        NFusion::TActionPtr GetDoc() final {
            NRealTime::TDistributorClient::TSignature signature;
            NFusion::TConsistentClient::TDocument::TSourceInfo info;
            if (!Client->GetNextDoc(Message, signature, info)) {
                return nullptr;
            }

            if (Message.HasDocument())
                Message.MutableDocument()->SetFilterRank(signature.Rank);

            NFusion::TActionPtr result = MakeAtomicShared<NSaas::TAction>(Message);
            NSaas::TDocument& document = result->GetDocument();

            const ui32 timestamp = document.GetTimestamp();
            result->SetReceiveTimestamp(Now() - TDuration::Seconds(signature.Age));
            if (const size_t replicasCount = info.Replicas.size()) {
                document.SetStreamId(
                    NFusion::GetSubStream(StreamId, info.Replicas[RandomNumber(replicasCount)])
                );
            }
            for (auto&& replica : info.Replicas) {
                document.AddExtraTimestamp(
                    NFusion::GetSubStream(StreamId, replica),
                    timestamp
                );
            }

            return result;
        }

    private:
        NFusion::TConsistentClientReplicas Replicas;
        NFusion::TConsistentClientOptions Options;
        const ui32 StreamId;
    };

    NFusion::TConsistentClientReplicas GetDistributorReplicas(const NFusion::TStreamConfig& config) {
        if (!config.DistributorReplicas.empty()) {
            return config.DistributorReplicas;
        } else {
            return NFusion::SplitReplicas(config.DistributorServers);
        }
    }
}

NFusion::TDistributorStream::TDistributorStream(const TDocFetcherModule& owner, const TStreamConfig& config, const TRTYServerConfig& globalConfig, TLog& log)
    : TBase(owner, config, globalConfig, log)
    , SignalsHolder<>(config)
    , Config(config)
    , Metrics(config.Name)
{
}

NFusion::TDistributorStream::~TDistributorStream() {
}

void NFusion::TDistributorStream::OnDocStreamStart() {
    CHECK_WITH_LOG(!Iterator);
    const auto& snapshot = GetTimestampSnapshot();
    if (Config.ConsistentClient) {
        Iterator = MakeHolder<TConsistentDistributorIterator>(
            Config,
            GetDistributorReplicas(Config),
            Config.ConsistentClientOptions,
            ItsOptions.MinTimestamp,
            snapshot
        );
    } else {
        Iterator = MakeHolder<TLegacyDistributorIterator>(Config, snapshot);
    }

    Iterator->SetMinRank(ItsOptions.MinRank);
}

void NFusion::TDistributorStream::OnDocStreamStop() {
    CHECK_WITH_LOG(Iterator);
    Iterator.Destroy();
}

void NFusion::TDistributorStream::OnReply(const NRTYServer::TReply& reply, NFusion::TActionPtr action) {
    if (static_cast<NRTYServer::TReply::TRTYStatus>(reply.GetStatus()) == NRTYServer::TReply::OK) {
        const ui32 ackAge = Seconds() - action->GetReceiveTimestamp().Seconds();
        Metrics.AckAge.Set(ackAge);
    }

    TBase::OnReply(reply, action);
}

bool NFusion::TDistributorStream::IsExhausted() const {
    return TBase::IsExhausted() && Iterator->IsExhausted();
}

void NFusion::TDistributorStream::OnExhausted(TTimeToSleep& timeToSleep) {
    Metrics.DistAge.Set(0);
    Metrics.DistributorTimestamp.Set(Seconds());
    TBase::OnExhausted(timeToSleep);
}

NFusion::TActionPtr NFusion::TDistributorStream::GetDoc(TMap<TString, TString>& /*propsToLog*/, bool& /*skipSleepOnError*/) {
    CHECK_WITH_LOG(Iterator);
    NFusion::TActionPtr action = Iterator->GetDoc();
    if (action) {
        const i64 diststamp = action->GetReceiveTimestamp().Seconds();
        const i64 distributorAge = Seconds() - diststamp;
        const i64 version = action->GetDocument().GetVersion();

        Metrics.DistAge.Set(distributorAge);
        Metrics.DistributorTimestamp.Set(diststamp);
        if (version) {
            Metrics.IndexedDocTimestamp.Set(version); // compliance with legacy Fusion monitoring
        }
    }
    return action;
}

bool NFusion::TDistributorStream::Process(IMessage* message_) {
    if (auto message = message_->As<TMessageSearchServerControl>()) {
        if (message->ShouldStart()) {
            Iterator->SetActive(true);
        }
        if (message->ShouldStop()) {
            Iterator->SetActive(false);
        }
        return true;
    }
    if (auto message = message_->As<TMessageSetDocfetcherTimestamp>()) {
        SetTimestamp(TInstant::Seconds(SafeIntegerCast<ui64>(message->Timestamp)), message->IncrementOnly);
        return true;
    }
    return TBase::Process(message_);
}

void NFusion::TDistributorStream::ReportStatus(NJson::TJsonValue& status) {
    status[Config.Name] = Iterator->GetStatus();
}

void NFusion::TDistributorStream::SetTimestamp(TInstant ts, bool incrementOnly) {
    CHECK_WITH_LOG(Iterator);

    const ui32 now = Seconds();
    const ui32 vl = ts.Seconds();
    const ui32 age = (now > vl) ? now - vl : 0;

    if (incrementOnly) {
        auto curAge = Metrics.DistAge.Get();
        if (curAge <= age && LastIncoming != TInstant::Zero()) {
            WARNING_LOG << "Stream " << Config.Name << " age is not set due to " << "increment_only" << Endl;
            return;
        }
    }

    Iterator->SetAge(age);
    NOTICE_LOG << "Stream " << Config.Name << " age is set to " << age << Endl;
}

TDuration NFusion::TDistributorStream::GetDelay(const NRTYServer::TTimestampSnapshot& snapshot) const {
    if (Config.ConsistentClient) {
        return GetDelayConsistent(snapshot);
    } else {
        return GetDelayLegacy(snapshot);
    }
}

TDuration NFusion::TDistributorStream::GetDelayConsistent(const NRTYServer::TTimestampSnapshot& snapshot) const {
    TMultiSet<TDuration> delays;
    for (auto&& replica : GetDistributorReplicas(Config)) {
        const ui32 substream = NFusion::GetSubStream(Config.StreamId, replica.Id);
        if (!snapshot.contains(substream)) {
            continue;
        }
        delays.insert(GetDelayFromSubStream(snapshot, substream));
    }

    if (delays.empty()) {
        // probably converting from legacy mode here
        return GetDelayLegacy(snapshot);
    } else {
        TMeanCalculator calculator;
        for (auto&& i : delays) {
            if (calculator.GetSumWeights() < 2 || i.Seconds() < 2 * calculator.GetMean()) {
                calculator.Add(i.Seconds());
            } else {
                WARNING_LOG << "Dropping delay " << i << Endl;
            }
        }
        return TDuration::Seconds(calculator.GetMean());
    }
}

TDuration NFusion::TDistributorStream::GetDelayLegacy(const NRTYServer::TTimestampSnapshot& snapshot) const {
    return GetDelayFromSubStream(
        snapshot,
        NFusion::GetSubStream(Config.StreamId, NFusion::SubStreamDocumentAux)
    );
}

TDuration NFusion::TDistributorStream::GetDelay() const {
    return GetDelay(GetTimestampSnapshot());
}

void NFusion::TDistributorStream::SubscribeToWatchdog(IWatchdogOptions& w) {
    // w.Subscribe will cause OnWatchdogOption() calls with the initial values.
    // That is why we do not use the "Messenger" here - we need _consistency_.
    // ivanmorozov@ seems not to understand stuff like this
    Watchdog = w.Subscribe(this);
}

void NFusion::TDistributorStream::OnWatchdogOption(const TString& key, const TString& value) {
    if (key == "refresh.df.rank_threshold") {
        NFusion::TMinRankCondition cond(value);
        if (ItsOptions.MinRank != cond) {
            ItsOptions.MinRank = cond;
            if (!!Iterator)
                Iterator->SetMinRank(ItsOptions.MinRank);
        }
    } else if (key == "refresh.df.min_timestamp") {
        TInstant ts = TInstant::Seconds(FromString<ui32>(value));
        if (ItsOptions.MinTimestamp != ts) {
            if (ts > Now())
                return;

            ItsOptions.MinTimestamp = ts;
            if (!!Iterator)
                SetTimestamp(ItsOptions.MinTimestamp, /*incrementOnly=*/ true);
        }
    }
}
