#include "persqueue.h"

#include <saas/api/action.h>
#include <saas/library/persqueue/common/common.h>
#include <saas/library/persqueue/common/installations.h>
#include <saas/library/persqueue/configuration/api/api.h>
#include <saas/library/persqueue/logger/logger.h>
#include <saas/library/persqueue/lock_consumer/lock.h>
#include <saas/rtyserver/common/common_messages.h>
#include <saas/rtyserver/common/stream_messages.h>

#include <saas/rtyserver/docfetcher/dc_checker/dc_checker.h>
#include <saas/rtyserver/docfetcher/library/stream_basics.h>
#include <saas/rtyserver/docfetcher/library/types.h>

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

#include <kikimr/persqueue/sdk/deprecated/cpp/v2/persqueue.h>
#include <kikimr/persqueue/sdk/deprecated/cpp/v1/consumer.h>

#include <library/cpp/logger/global/global.h>
#include <library/cpp/retry/retry.h>

#include <util/generic/cast.h>
#include <util/generic/maybe.h>
#include <util/generic/queue.h>
#include <util/generic/variant.h>
#include <util/string/join.h>

namespace NFusion {

namespace {
    class IPersQueue2Logger : public NPersQueue2::ILogger {
    public:
        virtual bool IsEnabled(int /* level*/) const { return false; }
    };
}

class IPersQueueReader : public TDatacenterChecker::ISubscriber {
public:
    struct TPersQueueDoc {
        TString Partition = "";
        ui64 Offset = 0;
        TString SourceId = "";
        TInstant CreateTime;
        TInstant WriteTime;
        TActionPtr Action;
        TPersQueueDoc() = default;
        TPersQueueDoc(
            const TString& partition,
            ui64 offset,
            const TString& sourceId,
            TInstant createTime,
            TInstant writeTime,
            TActionPtr action
        )
            : Partition(partition)
            , Offset(offset)
            , SourceId(sourceId)
            , CreateTime(createTime)
            , WriteTime(writeTime)
            , Action(action)
        {
        }
        bool operator< (const TPersQueueDoc& doc) const {
            return WriteTime < doc.WriteTime;
        }
    };

public:
    IPersQueueReader(const TPersQueueStreamConfig& config,
        TAtomicSharedPtr<TDatacenterChecker> datacenterChecker,
        TPersQueueStreamSignals* signals);
    virtual ~IPersQueueReader() = default;

    virtual void Start(const TInstant& ts) = 0;
    virtual void Start(const NRTYServer::TPositionsMap& positions) = 0;
    virtual void Stop() = 0;
    virtual void Restart() = 0;
    void AsyncRestart();

    virtual bool GetDoc(TPersQueueDoc& doc) = 0;
    virtual bool IsExhausted() const = 0;

    virtual NJson::TJsonValue GetStatus() = 0;

    void SetMinRank(const TMinRankCondition& rankCondition);
    void SetMinTimestamp(TInstant minTimestamp);

protected:
    virtual void BeforeStop();

    virtual void InitConsumerSettings() = 0;
    virtual TString GetTopicName() const = 0;

    virtual TString GetCurrentServer() const = 0;
    virtual TString GetCurrentConsumer() const = 0;

    TVector<TString> GetShardsList() const;

    TString ChoseConsumer();
    TString GetServerName(TStringBuf datacenter) const;
    TString ChoseServer() const;

    static TActionPtr ParseMessage(const TString& blob);

private:
    // TDatacenterChecker::ISubscriber
    void DatacenterChangedStatus(const TString& datacenter, bool isFailed) override;

protected:
    const TPersQueueStreamConfig& Config;
    TAtomicSharedPtr<TDatacenterChecker> DatacenterChecker;
    TVector<TString> Datacenters;
    THolder<NSaasLB::TConsumerLocker> ConsumerLocker;
    TPersQueueStreamSignals* Signals;

    TDocfetcherItsOptions ItsOptions;

    NAtomic::TBool NeedStop;
};

class TPersQueueClient : public IPersQueueReader {
private:
    class TMergedQueue {
    private:
        const TDuration ReceiveDelay;
        const TDuration BlockTime;

        TMultiSet<TPersQueueDoc> Docs;
        THashMap<TString, size_t> SizePerTopic;
        TInstant DontPopUntil;
        TSpinLock Lock;
    public:
        TMergedQueue(const TDuration& receiveDelay, const TDuration& blockTime)
            : ReceiveDelay(receiveDelay)
            , BlockTime(blockTime)
        {
        }
        void Init() {
            TGuard<TSpinLock> guard(Lock);
            DontPopUntil = TInstant::Now() + BlockTime;
        }
        void Clear() {
            TGuard<TSpinLock> guard(Lock);
            DontPopUntil = TInstant::Zero();
            Docs.clear();
            SizePerTopic.clear();
        }
        void Add(const TPersQueueDoc& doc) {
            TGuard<TSpinLock> guard(Lock);
            Docs.insert(doc);
            SizePerTopic[doc.Partition]++;
        }
        bool Pop(TPersQueueDoc& doc) {
            TGuard<TSpinLock> guard(Lock);
            if (Docs.empty() || TInstant::Now() < DontPopUntil) {
                return false;
            }
            if (TInstant::Now() - Docs.begin()->WriteTime < ReceiveDelay) {
                return false;
            }
            doc = *Docs.begin();
            Docs.erase(Docs.begin());
            SizePerTopic[doc.Partition]--;
            return true;
        }
        size_t Size(const TString& topic = "") const {
            TGuard<TSpinLock> guard(Lock);
            if (topic == "") {
                return Docs.size();
            }
            const auto it = SizePerTopic.find(topic);
            return it != SizePerTopic.end() ? it->second : 0;
        }
    };

    class TTopicInfo {
    public:
        TTopicInfo(NPersQueue2::TConsumerRecId readOffset = -1, NPersQueue2::TConsumerRecId commitOffset = -1)
            : ReadOffset(readOffset)
            , CommitOffset(commitOffset)
            , LastReceiveTimestamp(TInstant::Zero())
            , LastReceivedBatchTimestamp(TInstant::Zero())
            , Exhausted(true)
        {
        }
        NJson::TJsonValue GetStatus() const {
            NJson::TJsonValue topicValue;
            topicValue.InsertValue("ReadOffset", ReadOffset);
            topicValue.InsertValue("CommitOffset", CommitOffset);
            topicValue.InsertValue("Exhausted", Exhausted);
            topicValue.InsertValue("LastReceiveDuration", GetLeftTime(LastReceiveTimestamp));
            topicValue.InsertValue("PassedAfterReceiving", GetLeftTime(LastReceivedBatchTimestamp));
            return topicValue;
        }
        void Update(NPersQueue2::TConsumerRecId offset) {
            LastReceivedBatchTimestamp = TInstant::Now();
            Exhausted = false;
            ReadOffset = offset;
        }
        void Update(const TPersQueueDoc& doc) {
            Update(doc.Offset);
            if (doc.Action) {
                LastReceiveTimestamp = doc.Action->GetReceiveTimestamp();
            }
        }
        void UpdateCommitOffset(NPersQueue2::TConsumerRecId offset) {
            CommitOffset = ReadOffset = offset;
        }

    private:
        NJson::TJsonValue GetLeftTime(const TInstant ts) const {
            if (ts != TInstant::Zero()) {
                return ToString((TInstant::Now() - ts));
            }
            return "null";
        }

    public:
        NPersQueue2::TConsumerRecId ReadOffset;
        NPersQueue2::TConsumerRecId CommitOffset;
        TInstant LastReceiveTimestamp;
        TInstant LastReceivedBatchTimestamp;
        bool Exhausted;
    };

    using TOffsetSnapshot = THashMap<TString, TTopicInfo>;

    class ICallBackManager : public NPersQueue2::ICallBack {
    protected:
        class ICommonCallBack : public NPersQueue2::ICallBack {
        private:
            TSpinLock Lock;

        protected:
            ICallBackManager* Parent;
            const TString Topic;
            NAtomic::TBool NeedStop;

        public:
            ICommonCallBack(ICallBackManager* parent, const TString& topic, bool needStop)
                : Parent(parent)
                , Topic(topic)
                , NeedStop(needStop)
            {
                INFO_LOG << "Callback created for topic: " << Topic << Endl;
                CHECK_WITH_LOG(Parent);
                Parent->AddChild(this);
            }
            virtual ~ICommonCallBack() {
                Parent->RemoveChild(this);
                INFO_LOG << "Callback removed for topic: " << Topic << Endl;
            }

            TAtomicSharedPtr<ICallBack> Clone(const TString&, NPersQueue2::TConsumerRecId, NPersQueue2::TConsumerRecId, TFuture<void>&&) override {
                VERIFY_WITH_LOG(false, "Incorrect usage of ICommonCallBack::Clone");
                return nullptr;
            }
            void Stop() {
                NeedStop = true;
            }
        };
    private:
        TSet<ICommonCallBack*> Childs;
        TSpinLock Lock;

    protected:
        TPersQueueClient* Client;
        NAtomic::TBool NeedStop;

    protected:
        void AddChild(ICommonCallBack* child) {
            TGuard<TSpinLock> guard(Lock);
            Childs.insert(child);
        }
        void RemoveChild(ICommonCallBack* child) {
            TGuard<TSpinLock> guard(Lock);
            Childs.erase(child);
        }

    public:
        ICallBackManager(TPersQueueClient* client)
            : Client(client)
            , NeedStop(false)
        {
            CHECK_WITH_LOG(Client);
        }
        virtual ~ICallBackManager() = default;

        virtual void Stop() {
            NeedStop = true;
            TGuard<TSpinLock> guard(Lock);
            for (auto& child : Childs) {
                child->Stop();
            }
        }
    };

    class TReadCallBackManager : public ICallBackManager {
    protected:
        class TReadCallBack : public ICommonCallBack {
        private:
            const TDuration WaitInterval = TDuration::MilliSeconds(100);
            TPersQueueClient* Client;
            bool Start = true;

        public:
            TReadCallBack(TPersQueueClient* client, ICallBackManager* parent, const TString& topic, bool needStop)
                : ICommonCallBack(parent, topic, needStop)
                , Client(client)
            {
                CHECK_WITH_LOG(Client);
            }

            NPersQueue2::TConsumerRecId BeforeRead(bool* softCommit) override {
                while (!NeedStop && Client->QueueIsFull(Topic)) {
                    Sleep(WaitInterval);
                }

                *softCommit = !Start;
                Start = false;
                if (Client->GetReadOffset(Topic) < Client->GetCommitOffset(Topic)) {
                    ERROR_LOG << "Read offset (" << Client->GetReadOffset(Topic) << ") less then commit offset for topic="
                        << Topic
                        << ", read offset was set to " << Client->GetCommitOffset(Topic) << Endl;
                    Client->SetReadOffset(Topic, Client->GetCommitOffset(Topic));
                }
                return Client->GetReadOffset(Topic);
            }

            NPersQueue2::TConsumerRecId OnData(NPersQueue2::TBlob* blob, const NPersQueue2::TProducerInfo& info, size_t size, NPersQueue2::TConsumerRecId, bool*) override {
                Client->Apply(Topic, blob, info, size);
                return NPersQueue2::NO_POSITION_CHANGE;
            }
        };

    public:
        TReadCallBackManager(TPersQueueClient* client)
            : ICallBackManager(client)
        {
        }
        TAtomicSharedPtr<ICallBack> Clone(const TString& topic, NPersQueue2::TConsumerRecId, NPersQueue2::TConsumerRecId, TFuture<void>&&) override {
            return new TReadCallBack(Client, this, topic, NeedStop);
        }
    };

    class TGetOffsetsCallBackManager : public ICallBackManager {
    protected:
        class TGetOffsetsCallBack : public ICommonCallBack {
        private:
            const ui32 Timestamp;
            NPersQueue2::TConsumerRecId Left;
            NPersQueue2::TConsumerRecId Right;
        private:
            NPersQueue2::TConsumerRecId GetMiddle() const {
                return Left + (Right - Left) / 2;
            }
        public:
            TGetOffsetsCallBack(
                TGetOffsetsCallBackManager* parent,
                const TString& topic,
                NPersQueue2::TConsumerRecId lastOffset,
                ui32 timestamp,
                bool needStop
            )
                : ICommonCallBack(parent, topic, needStop)
                , Timestamp(timestamp)
                , Left(-1)
                , Right(lastOffset - 1)
            {
            }

            NPersQueue2::TConsumerRecId BeforeRead(bool* softCommit) override {
                DEBUG_LOG << "Find offset in process for topic=" << Topic << " range [" << Left << ":" << Right << "]" << Endl;
                *softCommit = false;
                if (!NeedStop && Left >= Right) {
                    NeedStop = true;
                    dynamic_cast<TGetOffsetsCallBackManager*>(Parent)->SetResult(Topic, Left);
                }
                return GetMiddle();
            }
            NPersQueue2::TConsumerRecId OnData(NPersQueue2::TBlob* blob, const NPersQueue2::TProducerInfo& info, size_t size, NPersQueue2::TConsumerRecId, bool*) override {
                if (size != 0) {
                    if (info[0].Offset != GetMiddle() + 1) {
                        Left = info[0].Offset;
                    } else {
                        auto action = ParseMessage({blob[0].data(), blob[0].size()});
                        if (!action || action->GetReceiveTimestamp().Seconds() < Timestamp) {
                            Left = GetMiddle() + 1;
                        } else {
                            Right = GetMiddle();
                        }
                    }
                } else {
                    Left = Right;
                }
                return NPersQueue2::NO_POSITION_CHANGE;
            }
        };
    private:
        const TDuration WaitingTime = TDuration::Minutes(1);
        const ui32 Timestamp;
        ui32 WaitPartitions;
        TManualEvent SnapshotReady;
        TOffsetSnapshot Snapshot;
        THashMap<TString, bool> TopicFinished;
        TSpinLock Lock;
    protected:
        void SetResult(const TString& topic, NPersQueue2::TConsumerRecId offset) {
            TGuard<TSpinLock> guard(Lock);
            if (TopicFinished[topic]) {
                return;
            }
            Snapshot[topic] = TTopicInfo(offset, offset);
            TopicFinished[topic] = true;
            WaitPartitions--;
            INFO_LOG << "Offset found for topic=" << topic << " offset=" << offset << ". Waiting " << WaitPartitions << " partitions" << Endl;
            if (!WaitPartitions) {
                SnapshotReady.Signal();
            }
        }
    public:
        TGetOffsetsCallBackManager(TPersQueueClient* client, ui32 timestamp)
            : ICallBackManager(client)
            , Timestamp(timestamp)
            , WaitPartitions(0)
        {
        }
        TMaybe<TOffsetSnapshot> GetSnapshot() {
            Sleep(WaitingTime);
            SnapshotReady.Wait();
            if (NeedStop) {
                return Nothing();
            }
            return Snapshot;
        }
        TAtomicSharedPtr<ICallBack> Clone(const TString &topic, NPersQueue2::TConsumerRecId offset, NPersQueue2::TConsumerRecId lag, TFuture<void>&&) override {
            TGuard<TSpinLock> guard(Lock);
            bool finished = NeedStop;
            NPersQueue2::TConsumerRecId recordsCount = offset + lag;
            if (!finished) {
                auto it = TopicFinished.find(topic);
                if (it == TopicFinished.end()) {
                    WaitPartitions++;
                    SnapshotReady.Reset();
                    TopicFinished[topic] = false;
                    INFO_LOG << "Starting search offset for topic=" << topic << " records=" << recordsCount << Endl;
                } else {
                    finished = it->second;
                }
            }

            return new TGetOffsetsCallBack(this, topic, recordsCount, Timestamp, finished);
        }
        void Stop() override {
            ICallBackManager::Stop();
            SnapshotReady.Signal();
        }
    };

private:
    THolder<IPersQueue2Logger> Logger;
    NPersQueue2::TSettings Settings;
    TVector<TString> PQShards;
    THolder<ICallBackManager> CallbackManager;
    TVector<THolder<NPersQueue2::TConsumer>> Consumers;
    using TPQLibPtr = TAtomicSharedPtr<NPersQueue2::TPQLib>;
    TPQLibPtr PQLib;
    static TPQLibPtr GlobalPQLib;
    TAtomic CurrentDatacenterIndex;
    TMutex Lock;
    TSpinLock CallbackManagerLock;
    TSpinLock LockTopicInfo;

    TMaybe<std::variant<NRTYServer::TPositionsMap, TInstant>> StartValue;

    NAtomic::TBool Running;

    TOffsetSnapshot Snapshot;
    TMergedQueue Queue;
private:
    void BeforeStop() override;

    template<class TStartValue>
    void StartNoLock(const TStartValue& startValue);
    void StartNoLock();
    template <class... TArgs>
    void RestartNoLock(TArgs&... args);
    void StopNoLock();

    static void* RestartProc(void* object);

    void StartConsumers(bool streamRead, size_t batchSize, ICallBackManager* callbackManager);
    void StopConsumers();
    void ClearQueue();

    TMaybe<TOffsetSnapshot> GetSnapshot(const TInstant& ts);
    TMaybe<TOffsetSnapshot> GetSnapshot(const NRTYServer::TPositionsMap& positions);
    void SetSnapshot(const TOffsetSnapshot& snapshot);
    void InitPQLib();

protected:
    void InitConsumerSettings() override;
    TString GetTopicName() const override;

    TString GetCurrentServer() const override;
    TString GetCurrentConsumer() const override;

public:
    TPersQueueClient(const TPersQueueStreamConfig& config, TAtomicSharedPtr<TDatacenterChecker> datacenterChecker, TPersQueueStreamSignals* signals, TLog& log);
    virtual ~TPersQueueClient() = default;

    void Start(const TInstant& ts) override;
    void Start(const NRTYServer::TPositionsMap& positions) override;
    void Stop() override;
    void Restart() override;
    template<class TStartValue>
    void Restart(const TStartValue& startValue);

    bool GetDoc(TPersQueueDoc& doc) override;
    bool IsExhausted() const override;

    NJson::TJsonValue GetStatus() override;

    bool QueueIsFull(const TString& topic = "") const;
    NPersQueue2::TConsumerRecId GetReadOffset(const TString& topic);
    NPersQueue2::TConsumerRecId GetCommitOffset(const TString& topic);
    std::optional<TPersQueueDoc> CreateDoc(const TString& topic, const NPersQueue2::TOneProducerInfo& info, const NPersQueue2::TBlob& blob);
    void Apply(const TString& topic, NPersQueue2::TBlob* blob, const NPersQueue2::TProducerInfo& info, size_t size);

    void ResetReadOffsets();
    void SetReadOffset(const TString& topic, NPersQueue2::TConsumerRecId offset);
    void SetCommitOffset(const TString& topic, NPersQueue2::TConsumerRecId offset);
};

class TPersQueueReader : public IPersQueueReader {
private:
    using TPQLibPtr = TAtomicSharedPtr<NPersQueue::TPQLib>;
    using TLbTopicKey = std::tuple<TString, TString>;

    struct TLbTopicInfo {
        TString cluster;
        TString qualifiedTopicName;

        TLbTopicInfo(TString cluster, TString qualifiedTopicName)
            : cluster(cluster)
            , qualifiedTopicName(qualifiedTopicName) {}
    };

    class TLockQueue {
    public:
        void Push(const TPersQueueDoc& doc) {
            TGuard<TSpinLock> guard(Lock);
            Queue.push(doc);
        }
        bool Pop(TPersQueueDoc& doc) {
            TGuard<TSpinLock> guard(Lock);
            if (Queue.empty()) {
                return false;
            }
            doc = Queue.front();
            Queue.pop();
            return true;
        }
        void Clear() {
            TGuard<TSpinLock> guard(Lock);
            Queue.clear();
        }
        ui64 Size() const {
            TGuard<TSpinLock> guard(Lock);
            return Queue.size();
        }
    private:
        TQueue<TPersQueueDoc> Queue;
        mutable TSpinLock Lock;
    };

    class TOffsets {
    private:
        const static ui64 UnsetOffset = -1;
        struct TPartition {
            ui64 Offset = UnsetOffset;
            TInstant ReceiveTimestamp = TInstant::Zero();
            TInstant AccessTimestamp = TInstant::Zero();
        };
    public:
        ui64 GetOffset(const TString& partition) const {
            TGuard<TSpinLock> guard(Lock);
            auto it = Partitions.find(partition);
            return it == Partitions.end() ? -1 : it->second.Offset;
        }
        void SetOffset(
            const TString& partition,
            ui64 offset,
            TInstant receiveTimestamp = TInstant::Zero(),
            TInstant accessTimestamp = TInstant::Zero()
        ) {
            TGuard<TSpinLock> guard(Lock);
            auto& info = Partitions[partition];
            info.Offset = offset;
            info.ReceiveTimestamp = receiveTimestamp;
            info.AccessTimestamp = accessTimestamp;

            if (offset != UnsetOffset && PartitionToMirrorPartitions.contains(partition)) {
                for (auto& mirrorPartition : PartitionToMirrorPartitions[partition]) {
                    auto& mirrorInfo = Partitions[mirrorPartition];
                    if (mirrorInfo.Offset != UnsetOffset) {
                        mirrorInfo.Offset = std::max(mirrorInfo.Offset, offset);
                    } else {
                        mirrorInfo.Offset = offset;
                    }
                }
            }
        }
        void DeclareMirrorPartitions(const TSet<TString>& partitions) {
            for (const auto& partition : partitions) {
                TSet<TString> curr(partitions);
                curr.erase(partition);
                PartitionToMirrorPartitions[partition] = curr;
            }
        }
        NJson::TJsonValue GetStatus() const {
            TGuard<TSpinLock> guard(Lock);
            NJson::TJsonValue status;
            for (auto partition : Partitions) {
                NJson::TJsonValue info;
                info.InsertValue("Offset", partition.second.Offset);
                info.InsertValue("LastReceiveDuration", GetPassedDuration(partition.second.ReceiveTimestamp));
                info.InsertValue("PassedAfterAccess", GetPassedDuration(partition.second.AccessTimestamp));
                status.InsertValue(partition.first, info);
            }
            return status;
        }
        TVector<std::pair<TString, ui64>> GetSnapshot() const {
            TGuard<TSpinLock> guard(Lock);
            TVector<std::pair<TString, ui64>> snapshot;
            for (auto partition : Partitions) {
                snapshot.push_back({ partition.first, partition.second.Offset });
            }
            return snapshot;
        }
        void Clear() {
            TGuard<TSpinLock> guard(Lock);
            Partitions.clear();
        }
    private:
        NJson::TJsonValue GetPassedDuration(TInstant ts) const {
            if (ts != TInstant::Zero()) {
                return (TInstant::Now() - ts).ToString();
            }
            return "null";
        }
    private:
        THashMap<TString, TPartition> Partitions;
        THashMap<TString, TSet<TString>> PartitionToMirrorPartitions;
        mutable TSpinLock Lock;
    };
public:
    TPersQueueReader(const TPersQueueStreamConfig& config, TAtomicSharedPtr<TDatacenterChecker> datacenterChecker, TPersQueueStreamSignals* signals, TLog& log);
    virtual ~TPersQueueReader() = default;

    void ResetState();
    void Start(const TInstant& ts) override;
    void Start(const NRTYServer::TPositionsMap& positions) override;
    void Stop() override;
    void Restart() override;

    bool GetDoc(TPersQueueDoc& doc) override;
    bool IsExhausted() const override;
    void LinkMirrorPartitions();

    NJson::TJsonValue GetStatus() override;

protected:
    void InitConsumerSettings() override;
    TString GetTopicName() const override;
    TString GetTopicPath(const TString& topicName) const;
    TString GetTopicNameFromPath(const TString& topicPath) const;

    TString GetCurrentServer() const override;
    TString GetCurrentConsumer() const override;

private:
    TVector<TString> GetTopicsList() const;

    void StartNoLock();
    void StopNoLock();
    static void* ReadProc(void* object);
    void ProcessPQMessage(const NPersQueue::TConsumerMessage& msg);

    TPersQueueDoc CreateDoc(const TString& partition, const NPersQueue::TReadResponse::TData::TMessage& pqMessage) const;

private:
    TIntrusivePtr<NPersQueue::ILogger> Logger;
    NPersQueue::TConsumerSettings ConsumerSettings;
    TPQLibPtr PQLib;
    TMutex Lock;

    TOffsets Offsets;
    TLockQueue Queue;
    TAtomic LastDataTimestamp = 0;
    NAtomic::TBool HasReadErrores = false;

    THolder<TThread> ReadThread;
};

IPersQueueReader::IPersQueueReader(const TPersQueueStreamConfig& config, TAtomicSharedPtr<TDatacenterChecker> datacenterChecker, TPersQueueStreamSignals* signals)
    : Config(config)
    , Datacenters(SplitString(Config.Datacenters, ","))
    , Signals(signals)
{
    Y_ASSERT(Signals != nullptr);
    if (datacenterChecker && Datacenters.size() > 1) {
        DatacenterChecker.Reset(datacenterChecker);
        DatacenterChecker->Subscribe(this);
    }
}

void IPersQueueReader::SetMinRank(const TMinRankCondition& rankCondition) {
    ItsOptions.MinRank = rankCondition;
}

void IPersQueueReader::SetMinTimestamp(TInstant minTimestamp) {
    INFO_LOG << "Set min_timestamp=" << minTimestamp.Seconds() << " (" << minTimestamp.ToStringLocal() << ") for stream=" << Config.Name << Endl;
    ItsOptions.MinTimestamp = minTimestamp;
}

void IPersQueueReader::BeforeStop() {
    NeedStop = true;
}

void IPersQueueReader::AsyncRestart() {
    TThread::TParams restartThreadParams([](void* object) -> void* {
        ThreadDisableBalloc();
        auto this_ = reinterpret_cast<IPersQueueReader*>(object);
        this_->Restart();
        return nullptr;
    }, this);

    TThread thread(restartThreadParams);
    thread.Start();
    thread.Detach();
}

TVector<TString> IPersQueueReader::GetShardsList() const {
    TVector<TString> shards;
    if (Config.TopicName) {
        shards.push_back(Config.TopicName);
    } else {
        if (Config.UseShardedLogtype) {
            shards.push_back(Join("-", "shard", Config.ShardMin, Config.ShardMax));
        } else {
            CHECK_EQ_WITH_LOG(Config.NumTopics % Config.NumShards, 0);

            ui32 topicsCount = (Config.NumTopics / Config.NumShards);
            ui32 topicMin = Config.NumTopics / Config.NumShards * Config.Shard;
            ui32 topicMax = Config.NumTopics / Config.NumShards * (Config.Shard + 1) - 1;

            for (ui32 t = topicMin; t <= topicMax; ++t) {
                shards.push_back("shard-" + ToString(t));
            }
            CHECK_WITH_LOG(topicsCount != 0);
            CHECK_EQ_WITH_LOG(topicsCount, shards.size());
        }
    }
    return shards;
}

TString IPersQueueReader::ChoseConsumer() {
    if (Config.Consumer) {
        return Config.Consumer;
    }
    if (ConsumerLocker) {
        auto consumer = DoWithRetry<TString, yexception>([this]() {
                if (this->NeedStop) {
                    return TString("");
                }
                TString consumer = this->ConsumerLocker->LockFreeConsumer();
                if (!consumer) {
                    Signals->ConsumerLocked(false);
                    ythrow yexception() << "no free consumer";
                }
                return consumer;
            },
            TRetryOptions(Max<ui32>(), TDuration::Seconds(10)),
            true
        );
        // we have either acquired a lock or need to stop. Anyway we do not need a lock.
        Signals->ConsumerLocked(true);
        return consumer.GetOrElse("");
    }
    TMaybe<TString> replicaId;
    for (auto& replicaIds : Config.ReplicaIdsList) {
        auto replicas = SplitString(replicaIds.second, ",");
        auto it = Find(replicas, Config.Replica);
        if (it != replicas.end()) {
            VERIFY_WITH_LOG(!replicaId, "Duplication replica id for %s", Config.Replica.data());
            replicaId = replicaIds.first;
        }
    }
    VERIFY_WITH_LOG(replicaId, "Not found replica id for %s", Config.Replica.data());
    return Config.UserPrefix + ToString(*replicaId);
}

TString IPersQueueReader::GetServerName(TStringBuf datacenter) const {
    return Join(".", datacenter, Config.Server);
}

TString IPersQueueReader::ChoseServer() const {
    TString server = Config.Server;
    if (DatacenterChecker) {
        auto datacenter = DatacenterChecker->GetCorrectDatacenter(Datacenters);
        server = GetServerName(datacenter);
    }
    return server;
}

TActionPtr IPersQueueReader::ParseMessage(const TString& data) {
    NRTYServer::TMessage message;
    if (!message.ParseFromString(data)) {
        return nullptr;
    }
    return MakeAtomicShared<NSaas::TAction>(message);
}

void IPersQueueReader::DatacenterChangedStatus(const TString& datacenter, bool isFailed) {
    bool needRestart = true;

    TString changedServer = GetServerName(datacenter);
    TString currentServer = GetCurrentServer();
    if (isFailed && changedServer != currentServer) {
        needRestart = false;
    }
    if (!isFailed) {
        for (auto dc : Datacenters) {
            TString server = GetServerName(dc);
            if (currentServer == server) {
                needRestart = false;
            }
            if (changedServer == server) {
                break;
            }
        }
    }

    if (needRestart) {
        INFO_LOG << "Restart stream=" << Config.StreamId << ", because the data center has changed status" << Endl;
        AsyncRestart();
    }
}

TPersQueueClient::TPersQueueClient(const TPersQueueStreamConfig& config, TAtomicSharedPtr<TDatacenterChecker> datacenterChecker, TPersQueueStreamSignals* signals, TLog& log)
    : IPersQueueReader(config, datacenterChecker, signals)
    , Logger(new NSaas::TPersQueueLogger<IPersQueue2Logger>(log))
    , Queue(Config.ReceiveDelay, Config.StartBlockTime)
{
    InitPQLib();

    InitConsumerSettings();

    if (Config.LockServers && Config.LockPath) {
        ConsumerLocker.Reset(new NSaasLB::TConsumerLocker(
            Config.LockServers,
            Config.LockPath,
            GetTopicName(),
            Config.Replica,
            [this]() {
                this->AsyncRestart();
            }
        ));
    }
}

void TPersQueueClient::StartConsumers(bool streamRead, size_t batchSize, ICallBackManager* callbackManager) {
    CHECK_WITH_LOG(callbackManager);
    {
        TGuard<TSpinLock> guard(CallbackManagerLock);
        CallbackManager.Reset(callbackManager);
    }

    Settings.ServerHostname = ChoseServer();
    Settings.FileName = ChoseConsumer();
    Settings.StreamRead = streamRead;

    if (NeedStop) {
        INFO_LOG << "Abort start of stream..." << Endl;
        return;
    }
    VERIFY_WITH_LOG(Settings.FileName, "Not specified the name of consumer");
    for (auto& shard : PQShards) {
        Settings.LogType = shard;
        {
            TGuard<TSpinLock> guard(CallbackManagerLock);
            Consumers.emplace_back(new NPersQueue2::TConsumer(
                *PQLib,
                Settings,
                batchSize,
                Config.BatchInBytes,
                Max<size_t>(),
                CallbackManager.Get(),
                Config.SleepTimeout,
                Config.LagTimeout,
                Logger.Get()
            ));
        }
        Consumers.back()->Start();
    }
    INFO_LOG << "Consumers started for stream: " << Config.StreamId << Endl;
}

void TPersQueueClient::StopConsumers() {
    BeforeStop();

    for (auto& consumer : Consumers) {
        consumer->Stop();
    }
    Consumers.clear();

    {
        TGuard<TSpinLock> guard(CallbackManagerLock);
        CallbackManager.Destroy();
    }
    if (ConsumerLocker) {
        ConsumerLocker->Release();
    }
}

void TPersQueueClient::BeforeStop() {
    IPersQueueReader::BeforeStop();
    TGuard<TSpinLock> guard(CallbackManagerLock);
    if (CallbackManager) {
        CallbackManager->Stop();
    }
}

void TPersQueueClient::Start(const TInstant& ts) {
    BeforeStop();

    TGuard<TMutex> guard(Lock);
    RestartNoLock(ts);
}

void TPersQueueClient::Start(const NRTYServer::TPositionsMap& positions) {
    BeforeStop();

    TGuard<TMutex> guard(Lock);
    RestartNoLock(positions);
}

void TPersQueueClient::Stop() {
    BeforeStop();

    TGuard<TMutex> guard(Lock);
    VERIFY_WITH_LOG(CallbackManager, "Client is not started, but need to stop");
    StartValue.Clear();
    StopNoLock();
}

template<class TStartValue>
void TPersQueueClient::Restart(const TStartValue& startValue) {
    BeforeStop();

    TGuard<TMutex> guard(Lock);
    if (!Running) {
        return;
    }
    RestartNoLock(startValue);
}

void TPersQueueClient::Restart() {
    BeforeStop();

    TGuard<TMutex> guard(Lock);
    if (!Running) {
        return;
    }
    if (StartValue) {
        std::visit([this](const auto& startValue) { this->RestartNoLock(startValue); }, *StartValue);
    } else {
        RestartNoLock();
    }
}

template<class TStartValue>
void TPersQueueClient::StartNoLock(const TStartValue& startValue) {
    NeedStop = false;
    Queue.Clear();
    ResetReadOffsets();

    Running = true;
    auto shapshot = GetSnapshot(startValue);
    if (!shapshot.Defined()) {
        StartValue.ConstructInPlace(startValue);
        return;
    }
    SetSnapshot(*shapshot);
    StartNoLock();
}

void TPersQueueClient::StartNoLock() {
    NeedStop = false;
    Running = true;
    StartConsumers(Config.StreamRead, Config.BatchSize, new TReadCallBackManager(this));
    StartValue.Clear();
    Queue.Init();
}

template <class... TArgs>
void TPersQueueClient::RestartNoLock(TArgs&... args) {
    StopNoLock();
    StartNoLock(args...);
}

void TPersQueueClient::StopNoLock() {
    Running = false;
    StopConsumers();
}

TMaybe<TPersQueueClient::TOffsetSnapshot> TPersQueueClient::GetSnapshot(const TInstant& ts) {
    TGetOffsetsCallBackManager* callbackManagerImpl = new TGetOffsetsCallBackManager(this, ts.Seconds());

    StartConsumers(false, 1, callbackManagerImpl);

    auto snapshot = callbackManagerImpl->GetSnapshot();
    StopConsumers();
    return snapshot;
}

TMaybe<TPersQueueClient::TOffsetSnapshot> TPersQueueClient::GetSnapshot(const NRTYServer::TPositionsMap& positions) {
    TOffsetSnapshot offsetSnapshot;

    for (const auto& position : positions) {
        const TString& topic = position.first;
        NPersQueue2::TConsumerRecId offset = position.second;
        offsetSnapshot[topic] = TTopicInfo(offset, offset);
    }
    return offsetSnapshot;
}

TString TPersQueueClient::GetTopicName() const {
    return Config.Ident + "--" + PQShards.front();
}

TString TPersQueueClient::GetCurrentServer() const {
    return Settings.ServerHostname;
}

TString TPersQueueClient::GetCurrentConsumer() const {
    return Settings.FileName;
}

void TPersQueueClient::InitConsumerSettings() {
    Settings.Entity = Config.Replica;
    Settings.Topic = Config.Ident;
    Settings.Timeout = Config.Timeout;
    Settings.ReadTimeout = Config.ReadTimeout;
    Settings.UseIpv6 = Config.UseIpv6;
    Settings.UseMirroredPartitions = Config.UseMirroredPartitions;
    Settings.Format = NPersQueue2::TSettings::FORMAT_RAW;

    PQShards = GetShardsList();
}

void TPersQueueClient::SetSnapshot(const TOffsetSnapshot& snapshot) {
    TGuard<TSpinLock> guard(LockTopicInfo);
    Snapshot = snapshot;
}

bool TPersQueueClient::QueueIsFull(const TString& topic /* = "" */) const {
    return (Queue.Size(topic) > Config.QueueSize);
}

NPersQueue2::TConsumerRecId TPersQueueClient::GetReadOffset(const TString& topic) {
    TGuard<TSpinLock> guard(LockTopicInfo);
    return Snapshot[topic].ReadOffset;
}

NPersQueue2::TConsumerRecId TPersQueueClient::GetCommitOffset(const TString& topic) {
    TGuard<TSpinLock> guard(LockTopicInfo);
    return Snapshot[topic].CommitOffset;
}

std::optional<IPersQueueReader::TPersQueueDoc> TPersQueueClient::CreateDoc(
    const TString& partition,
    const NPersQueue2::TOneProducerInfo& info,
    const NPersQueue2::TBlob& blob
) {
    if ((TInstant::Now() - TInstant::MilliSeconds(info.WriteTime)).Seconds() > Config.MaxAgeToGetSec) {
        TGuard<TSpinLock> guard(LockTopicInfo);
        Snapshot[partition].UpdateCommitOffset(info.Offset);
        return {};
    }
    TPersQueueDoc doc(
        partition,
        info.Offset,
        info.SourceId,
        TInstant::MilliSeconds(info.CreateTime),
        TInstant::MilliSeconds(info.WriteTime),
        ParseMessage(TString(blob.data(), blob.size()))
    );
    if (doc.Action) {
        doc.Action->GetDocument().SetStreamId(GetSubStream(Config.StreamId, SubStreamDocumentAux));
        doc.Action->SetPosition(doc.Partition, doc.Offset);
        TGuard<TSpinLock> guard(LockTopicInfo);
        for (auto&& topicInfo : Snapshot) {
            if (topicInfo.second.CommitOffset >= 0) {
                doc.Action->AddExtraPosition(topicInfo.first, topicInfo.second.CommitOffset);
            }
        }
    }
    return doc;
}

void TPersQueueClient::Apply(const TString& topic, NPersQueue2::TBlob* blob, const NPersQueue2::TProducerInfo& info, size_t size) {
    for (size_t i = 0; i < size; ++i) {
        if (auto doc = CreateDoc(topic, info[i], blob[i])) {
            Queue.Add(*doc);
            TGuard<TSpinLock> guard(LockTopicInfo);
            Snapshot[topic].Update(*doc);
        } else {
            TGuard<TSpinLock> guard(LockTopicInfo);
            Snapshot[topic].Update(info[i].Offset);
         }
    }
    if (size == 0) {
        TGuard<TSpinLock> guard(LockTopicInfo);
        Snapshot[topic].Exhausted = true;
    }
}

void TPersQueueClient::ResetReadOffsets() {
    TGuard<TSpinLock> guard(LockTopicInfo);
    for (auto& topicInfo : Snapshot) {
        topicInfo.second.ReadOffset = topicInfo.second.CommitOffset;
    }
}

void TPersQueueClient::SetReadOffset(const TString& topic, NPersQueue2::TConsumerRecId offset) {
    TGuard<TSpinLock> guard(LockTopicInfo);
    Snapshot[topic].ReadOffset = offset;
}

void TPersQueueClient::SetCommitOffset(const TString& topic, NPersQueue2::TConsumerRecId offset) {
    TGuard<TSpinLock> guard(LockTopicInfo);
    Snapshot[topic].CommitOffset = Max(Snapshot[topic].CommitOffset, offset);
}

bool TPersQueueClient::GetDoc(TPersQueueDoc& doc) {
    return Queue.Pop(doc);
}

bool TPersQueueClient::IsExhausted() const {
    TGuard<TSpinLock> guard(LockTopicInfo);
    for (auto& topic : Snapshot) {
        if (!topic.second.Exhausted) {
            return false;
        }
    }
    return true;
}

NJson::TJsonValue TPersQueueClient::GetStatus() {
    NJson::TJsonValue status, topics;
    status.InsertValue("Running", bool(Running));
    status.InsertValue("UseDatacenterChecker", bool(DatacenterChecker));
    status.InsertValue("Server", GetCurrentServer());
    status.InsertValue("User", GetCurrentConsumer());
    status.InsertValue("Ident", Config.Ident);
    status.InsertValue("CommonQueueSize", Queue.Size());

    TGuard<TSpinLock> guard(LockTopicInfo);
    for (auto& topic : Snapshot) {
        auto topicStatus = topic.second.GetStatus();
        topicStatus.InsertValue("QueueSize", Queue.Size(topic.first));
        topics.InsertValue(topic.first, topicStatus);
    }
    status.InsertValue("Topics", topics);
    return status;
}

TPersQueueClient::TPQLibPtr TPersQueueClient::GlobalPQLib;

void TPersQueueClient::InitPQLib() {
    if (Config.UseNewProtocol) {
        PQLib = MakeAtomicShared<NPersQueue2::TPQLib>(nullptr, Config.PQLibThreads, true);
    } else {
        if (!GlobalPQLib) {
            GlobalPQLib = MakeAtomicShared<NPersQueue2::TPQLib>(nullptr, Config.PQLibThreads, false);
        }
        PQLib = GlobalPQLib;
    }
}

TPersQueueReader::TPersQueueReader(const TPersQueueStreamConfig& config, TAtomicSharedPtr<TDatacenterChecker> datacenterChecker, TPersQueueStreamSignals* signals, TLog& log)
    : IPersQueueReader(config, datacenterChecker, signals)
    , Logger(new NSaas::TPersQueueLogger<NPersQueue::ILogger>(log))
{
    NPersQueue::TPQLibSettings pqLibSettings({ .ThreadsCount = Config.PQLibThreads, .DefaultLogger = Logger });
    PQLib.Reset(new NPersQueue::TPQLib(pqLibSettings));

    InitConsumerSettings();

    if (Config.LockServers && Config.LockPath) {
        ConsumerLocker.Reset(new NSaasLB::TConsumerLocker(
            Config.LockServers,
            Config.LockPath,
            GetTopicName(),
            Config.Replica,
            [this]() {
                this->AsyncRestart();
            }
        ));
    }
}

void TPersQueueReader::InitConsumerSettings() {
    ConsumerSettings.CommitsDisabled = true;
    ConsumerSettings.UseLockSession = true;
    ConsumerSettings.ReadMirroredPartitions = Config.UseMirroredPartitions;
    ConsumerSettings.MaxCount = Config.BatchSize;
    ConsumerSettings.MaxSize = Config.BatchInBytes;
    ConsumerSettings.Unpack = true;
    ConsumerSettings.MaxTimeLagMs = TInstant::Seconds(Config.MaxAgeToGetSec).MilliSeconds();

    ConsumerSettings.Topics = GetTopicsList();
}

TString TPersQueueReader::GetTopicName() const {
    return ConsumerSettings.Topics.front();
}

TString TPersQueueReader::GetCurrentServer() const {
    return ConsumerSettings.Server.Address;
}

TString TPersQueueReader::GetCurrentConsumer() const {
    return ConsumerSettings.ClientId;
}

TVector<TString> TPersQueueReader::GetTopicsList() const {
    auto shards = GetShardsList();
    const TString topicSeparator = Config.UseDirectoryTopicFormat ? "/" : "--";

    TVector<TString> topics;
    for (auto&& shard : shards) {
        topics.push_back(Join(topicSeparator, Config.Ident, shard));
    }
    INFO_LOG << "Topics to read: " << JoinSeq(",", topics) << Endl;
    return topics;
}

TString TPersQueueReader::GetTopicPath(const TString& topicName) const {
    TString shard;
    TString directory;
    StringSplitter(topicName).SplitByString("--").Limit(2).CollectInto(&directory, &shard);
    std::replace(directory.begin(), directory.vend(), '@', '/');
    return directory + "/" + shard;
}

TString TPersQueueReader::GetTopicNameFromPath(const TString& topicPath) const {
    TVector<TString> parts;
    StringSplitter(topicPath.begin(), topicPath.end()).Split('/').Collect(&parts);
    return JoinRange("@", parts.begin(), parts.end() - 1) + "--" + parts.back();
}

void TPersQueueReader::LinkMirrorPartitions() {
    Y_ENSURE(
        ConsumerSettings.Topics.size() == 1,
        "Remote mirror rules are not supported for multiple topics in one stream"
    );

    const TString token = NSaasLB::GetLogbrokerToken();
    const TString topicName = GetTopicName();
    TString topicPath = GetTopicPath(topicName);

    NSaasLB::TLogbrokerApi logbrokerAPI(NSaasLB::GetCMLogbrokerEndpointWithPort(Config.Server), token);

    const auto& entry = logbrokerAPI.DescribePath(topicPath);
    Y_ENSURE(entry.has_topic(), "The given path has no topics, path = " + topicPath);

    if (entry.topic().remote_mirror_rules().size() > 0) {
        THashMap<TLbTopicKey, TVector<TLbTopicInfo>> originToMirror;
        for (const auto& rule : entry.topic().remote_mirror_rules()) {
            const TString& cluster = rule.remote_mirror_rule().cluster().cluster();
            originToMirror[
                std::make_tuple(
                    NSaasLB::GetPropertyValue<TString>(rule.properties().src_cluster_endpoint()),
                    NSaasLB::GetPropertyValue<TString>(rule.properties().src_topic())
                )
            ].push_back(TLbTopicInfo(cluster, "rt3." + cluster + "--" + topicName));
        }
        Y_ENSURE(originToMirror.size() == 1, "Multiple source topics detected, fix mirror rules for " + topicPath);

        const TString originClusterEndpoint = std::get<0>(originToMirror.begin()->first);
        NSaasLB::ELogbrokerName originLogbroker = NSaasLB::GetLogbrokerNameByEndpointWithPort(originClusterEndpoint);
        NSaasLB::TLogbrokerApi originLogbrokerAPI(NSaasLB::GetCMLogbrokerEndpointWithPort(originLogbroker), token);

        const auto& originTopicPath = std::get<1>(originToMirror.begin()->first);
        const auto& originEntry = originLogbrokerAPI.DescribePath(originTopicPath);
        Y_ENSURE(originEntry.has_topic(), "The given path has no topics, path = " + originTopicPath);
        Y_ENSURE(
            originEntry.topic().instances().size() == 1,
            "Multiple instances of source topic detected, mirror rules are probably incorrect for " + topicPath
        );

        const TString originTopicName = GetTopicNameFromPath(originTopicPath);
        TVector<TLbTopicInfo>& topicsInfo = originToMirror.begin()->second;
        topicsInfo.push_back(
            TLbTopicInfo(
                originClusterEndpoint,
                "rt3." + originEntry.topic().instances()[0].origin_cluster() + "--" + originTopicName
            )
        );

        ui64 partitionsCnt = NSaasLB::GetPropertyValue<ui64>(originEntry.topic().properties().partitions_count());
        for(ui64 idx = 0; idx < partitionsCnt; idx++) {
            TSet<TString> partitions;
            for (const auto& topicInfo : topicsInfo) {
                partitions.insert(topicInfo.qualifiedTopicName + ":" + std::to_string(idx));
            }
            Offsets.DeclareMirrorPartitions(partitions);
        }
    }
}

void TPersQueueReader::ResetState() {
    StopNoLock();
    Queue.Clear();
    Offsets.Clear();

    try {
        LinkMirrorPartitions();
    } catch(...) {
        ERROR_LOG << "Unable to link mirror partitions, exc:\n" << CurrentExceptionMessage() << Endl;
    }
}

void TPersQueueReader::Start(const TInstant& ts) {
    BeforeStop();
    TGuard<TMutex> guard(Lock);

    ResetState();
    ConsumerSettings.ReadTimestampMs = Max(ItsOptions.MinTimestamp, ts).MilliSeconds();

    StartNoLock();
}

void TPersQueueReader::Start(const NRTYServer::TPositionsMap& positions) {
    BeforeStop();
    TGuard<TMutex> guard(Lock);

    ResetState();
    ConsumerSettings.ReadTimestampMs = ItsOptions.MinTimestamp.MilliSeconds();

    for (const auto& position : positions) {
        const TString& topic = position.first;
        ui64 offset = position.second;
        Offsets.SetOffset(topic, offset);
    }

    StartNoLock();
}

void TPersQueueReader::Stop() {
    BeforeStop();
    TGuard<TMutex> guard(Lock);
    StopNoLock();
}

void TPersQueueReader::Restart() {
    BeforeStop();
    TGuard<TMutex> guard(Lock);
    StopNoLock();
    StartNoLock();
}

bool TPersQueueReader::GetDoc(TPersQueueDoc& doc) {
    return Queue.Pop(doc);
}

bool TPersQueueReader::IsExhausted() const {
    return !HasReadErrores && TInstant::Now() - TInstant::Seconds(AtomicGet(LastDataTimestamp)) > Config.MaxIdleTime;
}

NJson::TJsonValue TPersQueueReader::GetStatus() {
    NJson::TJsonValue status, topics;
    status.InsertValue("UseDatacenterChecker", bool(DatacenterChecker));
    status.InsertValue("Server", GetCurrentServer());
    status.InsertValue("Consumer", GetCurrentConsumer());
    status.InsertValue("Ident", Config.Ident);
    status.InsertValue("QueueSize", Queue.Size());
    status.InsertValue("Partitions", Offsets.GetStatus());
    return status;
}

void TPersQueueReader::StartNoLock() {
    NeedStop = false;
    ConsumerSettings.Server.Address = ChoseServer();
    ConsumerSettings.ClientId = ChoseConsumer();

    if (NeedStop) {
        INFO_LOG << "Abort start of stream..." << Endl;
        return;
    }

    INFO_LOG << "Start reading with server=" << ConsumerSettings.Server.Address << " consumer=" << ConsumerSettings.ClientId << Endl;

    TThread::TParams readThreadParams(&ReadProc, this);
    readThreadParams.SetName("persqueue-read-proc");
    ReadThread.Reset(new TThread(readThreadParams));
    ReadThread->Start();
}

void TPersQueueReader::StopNoLock() {
    if (!ReadThread) {
        return;
    }
    NeedStop = true;
    ReadThread->Join();
    ReadThread.Destroy();
    NeedStop = false;

    ConsumerSettings.Server.Address = "";
    ConsumerSettings.ClientId = "";

    if (ConsumerLocker) {
        ConsumerLocker->Release();
    }
}

void* TPersQueueReader::ReadProc(void* object) {
    ThreadDisableBalloc();
    const TDuration waitInterval = TDuration::MilliSeconds(100);
    const TDuration connectionTimeout = TDuration::Minutes(1);
    auto this_ = reinterpret_cast<TPersQueueReader*>(object);

    if (this_->Config.TvmConfig) {
        auto tvmSettings = this_->Config.TvmConfig->GetSettings();
        while (!this_->NeedStop) {
            try {
                auto tvmClient = CreateTvmClient(tvmSettings);
                this_->ConsumerSettings.CredentialsProvider = CreateTVMCredentialsProvider(tvmClient, this_->Logger, tvmSettings.DestinationClients.begin()->first);
                break;
            } catch (...) {
                ERROR_LOG << "Can't create tvm credentials provider: " << CurrentExceptionMessage() << Endl;
                Sleep(connectionTimeout);
            }
        }
    } else {
        WARNING_LOG << "Stream without authorization in PQ: " << this_->Config.Name << Endl;
    }

    while (!this_->NeedStop) {
        try {
            auto consumer = this_->PQLib->CreateConsumer(this_->ConsumerSettings, nullptr, true);
            auto future = consumer->Start();

            future.Wait(connectionTimeout);
            if (!future.HasValue()) {
                ythrow yexception() << "Can't create consumer in " << connectionTimeout;
            }

            auto msg = consumer->GetNextMessage();
            while (!this_->NeedStop) {
                if (this_->Queue.Size() >= this_->Config.QueueSize) {
                    Sleep(waitInterval);
                    continue;
                }
                if (!msg.Wait(waitInterval)) {
                    continue;
                }
                this_->ProcessPQMessage(msg.GetValue());
                this_->HasReadErrores = false;
                msg = consumer->GetNextMessage();
            }
        } catch(...) {
            ERROR_LOG << "Error while reading: " << CurrentExceptionMessage() << Endl;
            this_->HasReadErrores = true;
        }
    }
    return nullptr;
}

void TPersQueueReader::ProcessPQMessage(const NPersQueue::TConsumerMessage& msg) {
    switch (msg.Type) {
    case NPersQueue::EMT_LOCK: {
        TString partition = Join(":", msg.Response.GetLock().GetTopic(), msg.Response.GetLock().GetPartition());
        ui64 offset = Offsets.GetOffset(partition) + 1;
        INFO_LOG << "Lock partition " << partition << " with offset " << offset << Endl;
        msg.ReadyToRead.SetValue(NPersQueue::TLockInfo{offset, 0, false});
        break;
    }
    case NPersQueue::EMT_RELEASE: {
        TString partition = Join(":", msg.Response.GetRelease().GetTopic(), msg.Response.GetRelease().GetPartition());
        INFO_LOG << "Release partition: " << partition << Endl;
        break;
    }
    case NPersQueue::EMT_DATA: {
        const auto& batch = msg.Response.GetData();
        for (auto messageBatch : batch.GetMessageBatch()) {
            TString partition = Join(":", messageBatch.GetTopic(), messageBatch.GetPartition());
            for (auto message : messageBatch.GetMessage()) {
                auto doc = CreateDoc(partition, message);
                TInstant receiveTimestamp = TInstant::Zero();
                if (doc.Action) {
                    receiveTimestamp = doc.Action->GetReceiveTimestamp();
                }
                if (!doc.Action || ItsOptions.MinTimestamp <= receiveTimestamp) {
                    Queue.Push(doc);
                }
                Offsets.SetOffset(partition, message.GetOffset(), receiveTimestamp, TInstant::Now());
            }
        }
        AtomicSet(LastDataTimestamp, TInstant::Now().Seconds());
        break;
    }
    case NPersQueue::EMT_COMMIT:
        break;
    case NPersQueue::EMT_ERROR:
        ythrow yexception() << msg.Response.GetError().GetDescription();
    default:
        ERROR_LOG << "Unknown message type. Response=" << msg.Response << Endl;
    }
}

IPersQueueReader::TPersQueueDoc TPersQueueReader::CreateDoc(const TString& partition, const NPersQueue::TReadResponse::TData::TMessage& pqMessage) const {
    auto& meta = pqMessage.GetMeta();
    TPersQueueDoc doc(
        partition,
        pqMessage.GetOffset(),
        meta.GetSourceId(),
        TInstant::MilliSeconds(meta.GetCreateTimeMs()),
        TInstant::MilliSeconds(meta.GetWriteTimeMs()),
        ParseMessage(pqMessage.GetData())
    );

    if (doc.Action) {
        doc.Action->SetPosition(partition, pqMessage.GetOffset());
        doc.Action->GetDocument().SetStreamId(GetSubStream(Config.StreamId, SubStreamDocumentAux));
        auto partitions = Offsets.GetSnapshot();
        for (auto partition : partitions) {
            doc.Action->AddExtraPosition(partition.first, partition.second);
        }
    }
    return doc;
}

TPersQueueStream::TPersQueueStream(
    const TDocFetcherModule& owner,
    const TPersQueueStreamConfig& config,
    const TRTYServerConfig& globalConfig,
    TAtomicSharedPtr<TDatacenterChecker> datacenterChecker,
    TLog& log,
    TLog& syslog
)
    : TBase(owner, config, globalConfig, log)
    , SignalsHolder<TPersQueueStreamSignals>(config)
    , Config(config)
{
    if (Config.UseNewPQLib) {
        Client.Reset(new TPersQueueReader(Config, datacenterChecker, Signals(), syslog));
    } else {
        Client.Reset(new TPersQueueClient(Config, datacenterChecker, Signals(), syslog));
    }
}

TPersQueueStream::~TPersQueueStream() = default;

void TPersQueueStream::OnDocStreamStart() {
    CHECK_WITH_LOG(Client);

    const auto& positionsSnapshot = GetPositionsSnapshot();
    const auto& timestampsSnapshot = GetTimestampSnapshot();
    auto subStreamId = GetSubStream(Config.StreamId, SubStreamDocumentAux);
    auto it = positionsSnapshot.find(subStreamId);
    auto positions = it == positionsSnapshot.end() ? NRTYServer::TPositionsMap() : it->second;

    TDuration delay = GetDelay(timestampsSnapshot);

    TInstant now = TInstant::Now();
    TDuration maxDelay = TDuration::Seconds(Config.MaxAgeToGetSec);
    if (ItsOptions.MinTimestamp <= now) {
        maxDelay = Min<TDuration>(maxDelay, now - ItsOptions.MinTimestamp);
    }

    if (!positions.empty() && delay <= maxDelay && !ItsOptions.DropPosition) {
        INFO_LOG << "Start stream=" << Config.StreamId << " from positions" << Endl;
        Client->Start(positions);
    } else {
        const auto it = timestampsSnapshot.find(subStreamId);
        ui32 streamTimestamp = 0;
        if (it != timestampsSnapshot.end()) {
            auto raw = it->second.MaxValue;
            if (SafeIntegerCast<ui64>(raw) > Max<ui32>()) {
                ERROR_LOG << "Incorrect timestamp for stream " << Config.StreamId << ": " << raw << Endl;
            } else {
                streamTimestamp = Max<ui32>(0, static_cast<ui32>(SafeIntegerCast<ui64>(raw)) - Config.OverlapAge);
            }
        }
        TInstant startTime = delay <= maxDelay ? TInstant::Seconds(streamTimestamp) : (now - maxDelay);
        INFO_LOG << "Start stream=" << Config.StreamId << " with time=" << startTime.ToStringLocal() << Endl;
        Client->Start(startTime);
    }
}

void TPersQueueStream::OnDocStreamStop() {
    CHECK_WITH_LOG(Client);
    Client->Stop();
}

bool TPersQueueStream::IsExhausted() const {
    return TBase::IsExhausted() && Client->IsExhausted();
}

TActionPtr TPersQueueStream::GetDoc(TMap<TString, TString>& propsToLog, bool& skipSleepOnError) {
    if (ItsOptions.MinRank.IsPaused())
        return nullptr;

    CHECK_WITH_LOG(Client);
    IPersQueueReader::TPersQueueDoc doc;
    if (!Client->GetDoc(doc)) {
        return nullptr;
    }

    propsToLog["partition"] = doc.Partition;
    propsToLog["position"] = ToString(doc.Offset);
    propsToLog["write_time"] = ToString(doc.WriteTime);
    propsToLog["source_id"] = doc.SourceId;

    if (!doc.Action) {
        ReportIncorrectDoc(TStringBuilder() << "Cannot parse message from " << doc.Partition << " offset " << doc.Offset);
        skipSleepOnError = true;
        return nullptr;
    }
    if (!doc.Action->HasDocument()) {
        ReportIncorrectDoc(TStringBuilder() << "Message has no document ", doc.Action);
        skipSleepOnError = true;
        return nullptr;
    }
    if (!doc.Action->GetDocument().HasTimestamp()) {
        ReportIncorrectDoc(TStringBuilder() << "Message has no timestamp ", doc.Action);
        skipSleepOnError = true;
        return nullptr;
    }

    if (!ItsOptions.MinRank.IsAllowedRank(doc.Action->GetDocument().GetFilterRank())) {
        skipSleepOnError = true;
        return nullptr;
    }

    LastReceiveTimestamp = doc.Action->GetReceiveTimestamp();
    return doc.Action;
}

bool TPersQueueStream::Process(IMessage* message_) {
    if (auto message = message_->As<TMessageSetDocfetcherTimestamp>()) {
        SetTimestamp(TInstant::Seconds(SafeIntegerCast<ui64>(message->Timestamp)), message->IncrementOnly);
        return true;
    }
    return TBase::Process(message_);
}

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

void TPersQueueStream::SetTimestamp(TInstant ts, bool incrementOnly) {
    CHECK_WITH_LOG(Client);
    if (incrementOnly) {
        if (ts <= LastReceiveTimestamp &&  LastIncoming != TInstant::Zero()) {
            WARNING_LOG << "Stream " << Config.Name << " age is not set due to increment_only" << Endl;
            return;
        }
    }
    Client->Start(ts);
    NOTICE_LOG << "Stream " << Config.Name << " age is set to " << TInstant::Now() - ts << Endl;
}

void TPersQueueStream::SetMinTimestamp(TInstant ts) {
    CHECK_WITH_LOG(Client);
    Client->SetMinTimestamp(ts);
    if (ts > LastReceiveTimestamp && LastIncoming != TInstant::Zero()) {
        Client->Start(ts);
    }
}

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

TDuration TPersQueueStream::GetDelay() const {
    return GetDelay(GetTimestampSnapshot());
}

void TPersQueueStream::ReportIncorrectDoc(const TString& expl, TActionPtr action) {
    NRTYServer::TReply reply;
    if (action && action->ToProtobuf().HasMessageId()) {
        reply.SetMessageId(action->ToProtobuf().GetMessageId());
    }
    reply.SetStatus(NRTYServer::TReply::INCORRECT_DOCUMENT);
    reply.SetStatusMessage(expl);
    OnReply(reply, action);
}

void TPersQueueStream::SubscribeToWatchdog(IWatchdogOptions& w) {
    Watchdog = w.Subscribe(this);
}

void TPersQueueStream::OnWatchdogOption(const TString& key, const TString& value) {
    if (key == "refresh.df.rank_threshold") {
        TMinRankCondition cond(value);
        if (ItsOptions.MinRank != cond) {
            ItsOptions.MinRank = cond;
            Client->SetMinRank(ItsOptions.MinRank);
        }
    } else if (key == "refresh.df.min_timestamp") {
        TInstant ts = TInstant::Seconds(FromString<ui32>(value));
        if (ItsOptions.MinTimestamp != ts) {
            ItsOptions.MinTimestamp = ts;
            SetMinTimestamp(ItsOptions.MinTimestamp);
        }
    } else if (key == "refresh.migrate.df.drop_position") {
        ItsOptions.DropPosition = IsTrue(value);
    }
}

}; // namespace NFusion
