#include "saas_yt_writer.h"

#include <google/protobuf/messagext.h>
#include <google/protobuf/text_format.h>

#include <util/stream/format.h>

namespace NSaas {
    namespace NYTPull {

        TSaasSlaveTableWriter::TSaasSlaveTableWriter(NYT::TTableWriter<NYT::TYaMRRow>* slaveWriterPtr,
                                                     TAtomicSharedPtr<IShardDispatcher> shardDispatcherPtr,
                                                     bool extractKiwiObject, bool writePrototext)
            : SlaveWriterPtr(slaveWriterPtr)
            , ShardDispatcherPtr(std::move(shardDispatcherPtr))
            , ExtractKiwiObject(extractKiwiObject)
            , WritePrototext(writePrototext)
        {
            Y_ENSURE(SlaveWriterPtr, "SlaveWriter is null");
            Y_ENSURE(ShardDispatcherPtr, "ShardDispatcher is null");
        }

        bool TSaasSlaveTableWriter::WriteMessage(const NRTYServer::TMessage& message) {
            NYT::TYaMRRow row;
            if (!message.HasDocument()) {
                return false;
            }
            const auto& doc = message.GetDocument();
            TStringStream ss;
            const NSearchMapParser::TShardIndex idx = ShardDispatcherPtr->GetShard(message);
            TString shardString = NSaas::GetShardString(idx);
            row.Key = shardString;
            row.SubKey = doc.GetUrl();
            TString messageStr;
            if (ExtractKiwiObject && !WritePrototext) {
                messageStr = doc.GetIndexedDoc().GetKiwiObject();
            } else {
                if (Y_UNLIKELY(WritePrototext)) {
                    using namespace ::google::protobuf;
                    TextFormat::Printer printer;
                    printer.SetHideUnknownFields(true);
                    if (!printer.PrintToString(message, &messageStr)) {
                        return false;
                    }
                } else {
                    if (!message.SerializeToString(&messageStr)) {
                        return false;
                    }
                }
            }
            row.Value = messageStr;
            SlaveWriterPtr->AddRow(row);
            return true;
        }

        TSaasSlaveTableHoldingWriter::TSaasSlaveTableHoldingWriter(NYT::IClientBasePtr client, const TString& tablePath,
                                                                   TAtomicSharedPtr<IShardDispatcher> shardDispatcherPtr,
                                                                   bool extractKiwiObject)
            : Client(client)
            , TablePath(tablePath)
            , SlaveTableWriterPtr(Client->CreateTableWriter<NYT::TYaMRRow>(tablePath))
            , Writer(SlaveTableWriterPtr.Get(), std::move(shardDispatcherPtr), extractKiwiObject)
        {
        }

        TSaasSlaveTableHoldingWriter::~TSaasSlaveTableHoldingWriter() {
            if (!!SlaveTableWriterPtr) {
                Finish();
            }
        }

        bool TSaasSlaveTableHoldingWriter::WriteMessage(const NRTYServer::TMessage& message) {
            Y_ENSURE(!!SlaveTableWriterPtr, "Slave table writer already closed");
            return Writer.WriteMessage(message);
        }

        void TSaasSlaveTableHoldingWriter::Finish() {
            SlaveTableWriterPtr->Finish();
            SlaveTableWriterPtr.Drop();
            SortSlaveTable(TablePath, Client);
            Client.Drop();
        }

        void SortSlaveTable(const TString& slaveTablePath, NYT::IClientBasePtr client) {
            SortSlaveTable(slaveTablePath, *client);
        }

        void CreateMasterTable(const TString& masterTablePath, NYT::IClientBasePtr client) {
            CreateMasterTable(masterTablePath, *client);
        }

        void RegisterSlaveTable(const TString& masterTablePath, const TString& slaveTablePath, const ui64 timestamp, NYT::IClientBasePtr client) {
            RegisterSlaveTable(masterTablePath, slaveTablePath, timestamp, *client);
        }

        void SortSlaveTable(const TString& slaveTablePath, NYT::IClientBase& client) {
            client.Sort(NYT::TSortOperationSpec().AddInput(slaveTablePath).Output(slaveTablePath).SortBy({"key", "subkey"}));
        }

        void CreateMasterTable(const TString& masterTablePath, NYT::IClientBase& client) {
            NYT::TNode key = NYT::TNode::CreateMap();
            key["name"] = "key";
            key["required"] = false;
            key["sort_order"] = "ascending";
            key["type"] = "any";
            NYT::TNode subkey = NYT::TNode::CreateMap();
            subkey["name"] = "subkey";
            subkey["required"] = false;
            subkey["sort_order"] = "ascending";
            subkey["type"] = "any";
            NYT::TNode value = NYT::TNode::CreateMap();
            value["name"] = "value";
            value["required"] = false;
            value["type"] = "any";
            NYT::TNode columns = NYT::TNode::CreateList();
            columns.Add(key);
            columns.Add(subkey);
            columns.Add(value);
            NYT::TNode attributes = NYT::TNode::CreateMap();
            attributes["schema"] = columns;
            client.Create(masterTablePath, NYT::NT_TABLE, NYT::TCreateOptions().IgnoreExisting(true).Recursive(true).Attributes(attributes));
        }

        void AddSlaveTable(NYT::TTableWriter<NYT::TYaMRRow>& masterWriter, TStringBuf slaveTablePath, ui64 timestamp) {
            NYT::TYaMRRow masterRow;
            TString key = ToString(timestamp);
            masterRow.Key = key;
            masterRow.Value = slaveTablePath;
            masterWriter.AddRow(masterRow);
        }

        void RegisterSlaveTable(const TString& masterTablePath, const TString& slaveTablePath, const ui64 timestamp, NYT::IClientBase& client) {
            auto masterWriter = client.CreateTableWriter<NYT::TYaMRRow>(NYT::TRichYPath(masterTablePath).Append(true).SortedBy({"key", "subkey"}));
            AddSlaveTable(*masterWriter, slaveTablePath, timestamp);
            masterWriter->Finish();
        }

        void RegisterSlaveTables(const TString& masterTablePath, const TVector<std::pair<TString, ui64>>& tables, NYT::IClientBase& client) {
            auto masterWriter = client.CreateTableWriter<NYT::TYaMRRow>(NYT::TRichYPath(masterTablePath).Append(true).SortedBy({"key", "subkey"}));
            for (auto& table : tables) {
                AddSlaveTable(*masterWriter, table.first, table.second);
            }
            masterWriter->Finish();
        }

    } // namespace NYTPull
} // namespace NSaas
