#include "common.h"

#include <mapreduce/yt/interface/client.h>
#include <util/stream/mem.h>
#include <library/cpp/logger/global/global.h>
#include <google/protobuf/messagext.h>

void TMessageDumper::ProcessMessage(const NRTYServer::TMessage& message) {
    NYT::TNode node;
    node["subkey"] = message.GetDocument().GetUrl();
    node["value"] = message.SerializeAsString();
    node["key"] = GetShardIndex(message, *ShardDispatcher);
    Writer->AddRow(node);
    if (++Progress % 100000 == 0) {
        Cerr << "\r" << Progress << " lines written so far...";
        Cerr.Flush();
    }
}

void TDebugDumper::ProcessMessage(const NRTYServer::TMessage& message) {
    NYT::TNode node;
    node["kps"] = message.GetDocument().GetKeyPrefix();
    node["value"] = message.DebugString();
    node["url"] = message.GetDocument().GetUrl();
    node["timestamp"] = message.GetDocument().GetModificationTimestamp();
    node["shard_id"] = GetShardIndex(message, *ShardDispatcher);
    Writer->AddRow(node);
    if (++Progress % 100000 == 0) {
        Cerr << "\r" << Progress << " lines written so far...";
        Cerr.Flush();
    }
}

void TDirDumper::ProcessMessage(const NRTYServer::TMessage& message) {
    NYT::TNode node;
    node["kps"] = message.GetDocument().GetKeyPrefix();
    node["url"] = message.GetDocument().GetUrl();
    node["timestamp"] = message.GetDocument().GetModificationTimestamp();
    node["shard_id"] = GetShardIndex(message, *ShardDispatcher);
    Writer->AddRow(node);
    if (++Progress % 100000 == 0) {
        Cerr << "\r" << Progress << " lines written so far...";
        Cerr.Flush();
    }
}

bool UnpackMessage(const TBlob& blob, NRTYServer::TMessage& message) {
    NRTYServer::TParsedDoc parsedDoc;
    try {
        TMemoryInput inp(blob.Data(), blob.Size());
        ::google::protobuf::io::TProtoSerializer::Load(&inp, parsedDoc);
    } catch (...) {
        ERROR_LOG << "Cannot parse document: " << CurrentExceptionMessage() << Endl;
        return false;
    }
    message.SetMessageType(NRTYServer::TMessage::MODIFY_DOCUMENT);
    message.MutableDocument()->Swap(parsedDoc.MutableDocument());
    return true;
}

TString GetShardIndex(const NRTYServer::TMessage& message, NSaas::IShardDispatcher& shardDispatcher) {
    return NSaas::GetShardString(shardDispatcher.GetShard(message));
}

THolder<NSaas::IShardDispatcher> CreateShardDispatcher(const NUtils::TParams& params) {
    NSaas::TShardsDispatcher::TContext ctx(
        params.GetShardingType().empty() ? NSaas::UrlHash : FromString<NSaas::ShardingType>(params.GetShardingType()),
        params.GetShardingKpsShift()
    );
    return MakeHolder<NSaas::TShardsDispatcher>(ctx);
}

TString GetLayerPrefix(const NUtils::TParams& params) {
    return params.GetLayerPrefix().empty() ? "indexfullarc.full" : params.GetLayerPrefix();
}
