#pragma once

#include <saas/library/sharding/sharding.h>
#include <saas/protos/rtyserver.pb.h>

#include <mapreduce/yt/interface/client.h>

#include <util/generic/vector.h>

namespace NSaas {
    namespace NYTPull {

        class TSaasSlaveTableWriter : public TNonCopyable {
            NYT::TTableWriter<NYT::TYaMRRow>* SlaveWriterPtr;
            TAtomicSharedPtr<IShardDispatcher> ShardDispatcherPtr;
            bool ExtractKiwiObject = false;
            bool WritePrototext = false;

        public:
            TSaasSlaveTableWriter(NYT::TTableWriter<NYT::TYaMRRow>* slaveWriterPtr,
                                  TAtomicSharedPtr<IShardDispatcher> shardDispatcherPtr,
                                  bool extractKiwiObject, bool writePrototext = false);

            bool WriteMessage(const NRTYServer::TMessage& message);
        };

        class TSaasSlaveTableHoldingWriter : public TNonCopyable {
        private:
            NYT::IClientBasePtr Client;
            TString TablePath;
            NYT::TTableWriterPtr<NYT::TYaMRRow> SlaveTableWriterPtr;
            TSaasSlaveTableWriter Writer;

        public:
            TSaasSlaveTableHoldingWriter(NYT::IClientBasePtr client, const TString& tablePath,
                                         TAtomicSharedPtr<IShardDispatcher> shardDispatcherPtr, bool extractKiwiObject);

            virtual ~TSaasSlaveTableHoldingWriter();

            bool WriteMessage(const NRTYServer::TMessage& message);
            void Finish();
       };

        void CreateMasterTable(const TString& masterTablePath, NYT::IClientBasePtr client);
        void SortSlaveTable(const TString& slaveTablePath, NYT::IClientBasePtr client);
        void RegisterSlaveTable(const TString& masterTablePath, const TString& slaveTablePath, const ui64 timestamp, NYT::IClientBasePtr client);

        void CreateMasterTable(const TString& masterTablePath, NYT::IClientBase& client);
        void SortSlaveTable(const TString& slaveTablePath, NYT::IClientBase& client);
        void RegisterSlaveTable(const TString& masterTablePath, const TString& slaveTablePath, const ui64 timestamp, NYT::IClientBase& client);
        void RegisterSlaveTables(const TString& masterTablePath, const TVector<std::pair<TString, ui64>>& tables, NYT::IClientBase& client);

    } // namespace NYTPull
} // namespace NSaas
