#include <saas/tools/message2json/params.pb.h>
#include <saas/protos/rtyserver.pb.h>
#include <saas/api/action.h>

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

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

struct TMessageToJsonConverter : NYT::IMapper<NYT::TTableReader<NYT::TNode>, NYT::TTableWriter<NYT::TNode>> {
    TMessageToJsonConverter() = default;
    TMessageToJsonConverter(TString inputColumn, TString outputColumn)
        : InputColumn(std::move(inputColumn))
        , OutputColumn(std::move(outputColumn))
    {
    }

    Y_SAVELOAD_JOB(InputColumn, OutputColumn);

    void Do(TReader* reader, TWriter* writer) final {
        for (; reader->IsValid(); reader->Next()) {
            try {
                const auto& inRow = reader->GetRow();
                const auto& inMap = inRow.AsMap();
                const auto it = inMap.find(InputColumn);
                if (it != inMap.end()) {
                    NRTYServer::TMessage message;
                    Y_ENSURE(message.ParseFromString(it->second.AsString()));
                    NSaas::TAction action(message);
                    NYT::TNode outRow;
                    outRow[OutputColumn] = action.BuildJsonMessage();
                    writer->AddRow(outRow);
                }
            } catch (...) {
                ERROR_LOG << CurrentExceptionMessage() << Endl;
            }
        }
    }

private:
    TString InputColumn;
    TString OutputColumn;
};

REGISTER_MAPPER(TMessageToJsonConverter);

int main(int argc, const char* argv[]) {
    try {
        NYT::Initialize(argc, argv);
        NUtils::TParams params = NGetoptPb::GetoptPbOrAbort(argc, argv);

        auto yt = NYT::CreateClient(params.GetYt());
        auto tx = yt->StartTransaction();

        TVector<NYT::ILockPtr> locks(Reserve(params.InputSize()));
        TVector<NThreading::TFuture<void>> futures(Reserve(params.InputSize()));
        for (auto& input : params.GetInput()) {
            auto lock = tx->Lock(input, NYT::LM_SNAPSHOT, NYT::TLockOptions().Waitable(true));
            futures.emplace_back(lock->GetAcquiredFuture());
            locks.emplace_back(std::move(lock));
        }
        NThreading::WaitExceptionOrAll(futures).GetValue(TDuration::Minutes(1));

        NYT::TMapOperationSpec spec;
        for (auto& lock : locks) {
            spec.AddInput<NYT::TNode>(TString::Join('#', GetGuidAsString(lock->GetLockedNodeId())));
        }
        spec.AddOutput<NYT::TNode>(params.GetOutput());

        tx->Map(spec, new TMessageToJsonConverter(params.GetInputColumn(), params.GetOutputColumn()));

        tx->Commit();

    } catch (...) {
        ERROR_LOG << CurrentExceptionMessage() << Endl;
        return 1;
    }
    return 0;
}
