#include "yt_renaming_mapper.h"

#include <saas/library/yt/common/yt_blob.h>
#include <library/cpp/protobuf/yt/proto2yt.h>

REGISTER_MAPPER(TSaasYTRenamingMapper);

TSaasYTRenamingMapper::TSaasYTRenamingMapper(TRenameMap renameMap, EYTBlobFormat format)
    : RenameMap(std::move(renameMap))
    , Format(format)
{
}

TSaasYTRenamingMapper::TSaasYTRenamingMapper(NYT::TTableReader<NSaas::TYTBlobBase>* renameMapTable, EYTBlobFormat format)
    : Format(format)
    {
    for (; renameMapTable->IsValid(); renameMapTable->Next()) {
        TYTBlob row(renameMapTable->GetRow()); // assuming that rename table always in old format
        RenameMap.emplace(std::make_pair(row.GetShardId(), row.GetBlobRef().GetName()), row.GetBlobRef().GetData());
    }
}

void TSaasYTRenamingMapper::Do(TReader* input, TWriter* output) {
    for (; input->IsValid(); input->Next()) {
        TYTBlob row(input->GetRow<NSaas::TYTBlobBase>());
        auto key = std::make_pair(row.GetShardId(), row.GetBlobRef().GetName());
        TString newName = RenameMap[key];
        if (newName) {
            row.GetBlobRef().SetName(newName);
        }
        if (Format == EYTBlobFormat::Old) {
            output->AddRow(static_cast<NSaas::TYTBlobBase>(row));
        } else {
            output->AddRow(static_cast<NSaas::TYTNewBlobBase>(row));
        }
    }
}
