#include <saas/tools/map_to_readable/doc_adapter.pb.h>

#include <saas/protos/rtyserver.pb.h>
#include <mapreduce/yt/interface/client.h>
#include <google/protobuf/text_format.h>


class TSimpleDocMapper : public NYT::IMapper<NYT::TTableReader<NSaas::TYTDocAdapter>, NYT::TTableWriter<::google::protobuf::Message>> {
    bool AddZoneView = false;

public:
    Y_SAVELOAD_JOB(AddZoneView);
    TSimpleDocMapper() = default;
    TSimpleDocMapper(bool addZoneView)
        : AddZoneView(addZoneView) {
    }
        
    void Do(TReader* in, TWriter* out) override {
        NSaas::TZoneViewer zoneView;
        TString textView;
        google::protobuf::TextFormat::Printer printer;
        printer.SetUseUtf8StringEscaping(true);
        while (in->IsValid()) {
            auto row = in->GetRow();
            if (row.HasDocument()) {
                out->AddRow<NRTYServer::TMessage::TDocument>(row.GetDocument(), 0);
                if (AddZoneView) {
                    zoneView.SetUrl(row.GetDocument().GetUrl());
                    printer.PrintToString(row.GetDocument().GetRootZone(), &textView);
                    zoneView.SetZoneView(textView);
                    out->AddRow<NSaas::TZoneViewer>(zoneView, 1);
                }
            } else if (row.HasValue()) {
                out->AddRow<NRTYServer::TMessage::TDocument>(row.GetValue().GetDocument(), 0);
                if (AddZoneView) {
                    zoneView.SetUrl(row.GetValue().GetDocument().GetUrl());
                    printer.PrintToString(row.GetValue().GetDocument().GetRootZone(), &textView);
                    zoneView.SetZoneView(textView);
                    out->AddRow<NSaas::TZoneViewer>(zoneView, 1);
                }
            }
            in->Next();
        }
        out->Finish();
    }
};

REGISTER_MAPPER(TSimpleDocMapper);

int main(int argc, char **argv) {
    NYT::Initialize(argc, argv);
    if (argc < 4 || argc > 5) {
        Cerr << "Usage: ./map_to_readable $PROXY $SRC_TABLE $DST_TABLE [$ZONES_TEXT_TABLE]" << Endl;
        return 0;
    }
    auto client = NYT::CreateClient(argv[1]);
    client->Create(argv[3], NYT::NT_TABLE,
                NYT::TCreateOptions().Force(true).
                        Attributes(NYT::TNode()
                            ("schema", NYT::CreateTableSchema<NRTYServer::TMessage::TDocument>({}).ToNode())));

    NYT::TMapOperationSpec spec;
    spec.AddInput<NSaas::TYTDocAdapter>(argv[2]);
    spec.AddOutput<NRTYServer::TMessage::TDocument>(argv[3]);

    bool addZoneView = false;
    if (argc == 5) {
        addZoneView = true;
        spec.AddOutput<NSaas::TZoneViewer>(argv[4]);
    }
    client->Map(spec, new TSimpleDocMapper(addZoneView));
}
