#include "common.h"
#include "mapper.h"
#include "reducer.h"
#include <saas/tools/dump_fullarc_to_yt/params.pb.h>
#include <saas/library/yt/common/yt_blob.h>

#include <kernel/multipart_archive/multipart.h>

#include <mapreduce/yt/interface/client.h>
#include <library/cpp/getoptpb/getoptpb.h>
#include <google/protobuf/messagext.h>
#include <util/system/env.h>


void LocalArchiveDump(const NUtils::TParams& params) {
    INFO_LOG << "Initiating local fullarchive dumping" << Endl;
    auto compression = params.GetCompression().empty() ? "COMPRESSED" : params.GetCompression();
    auto shardDispatcher = CreateShardDispatcher(params);
    NRTYArchive::TMultipartConfig layerConfig;
    layerConfig.Compression = FromString<NRTYArchive::IArchivePart::TType>(compression);
    TArchiveOwner::TPtr archive = TArchiveOwner::Create(TFsPath{params.GetIndexDir()} / GetLayerPrefix(params), layerConfig, 0, true);

    NRTYArchive::TMultipartArchive::TIteratorPtr archiveIterator = archive->CreateIterator();

    auto ytClient = NYT::CreateClient(params.GetYt());
    auto tx = ytClient->StartTransaction();
    tx->Create(params.GetOutput(), NYT::NT_TABLE, NYT::TCreateOptions().IgnoreExisting(true).Recursive(true));
    auto writer = tx->CreateTableWriter<NYT::TNode>(NYT::TRichYPath(params.GetOutput()).Append(true));
    THolder<IFormatDumper> dumper;
    TString dumpFormat = params.GetDumpFormat().empty() ? "Message" : params.GetDumpFormat();
    if (dumpFormat == "Message") {
        dumper = MakeHolder<TMessageDumper>(writer.Get(), shardDispatcher.Get());
    } else if (dumpFormat == "UrlsOnly") {
        dumper = MakeHolder<TDirDumper>(writer.Get(), shardDispatcher.Get());
    } else {
        dumper = MakeHolder<TDebugDumper>(writer.Get(), shardDispatcher.Get());
    }

    size_t progress = 0;
    ReadArchive(*archiveIterator, dumper.Get());
    Cout << "\n" << progress << " lines written." << Endl;

    writer->Finish();

    tx->Commit();
}

void YTArchiveDump(const NUtils::TParams& params) {
    INFO_LOG << "Initiating YT fullarchive dumping" << Endl;

    auto client = NYT::CreateClient(params.GetYt());
    TString tmpDir = params.GetTmpDir().empty() ? "//tmp/saas" : params.GetTmpDir();
    bool tmpDirCreated = !client->Exists(tmpDir);
    NYT::TRichYPath mapperOutput(tmpDir + "/secondary_segments");
    NYT::TRichYPath sortedOutput(tmpDir + "/secondary_segments.sorted");
    auto reduceKeys = ShardIdKeyColumns;
    reduceKeys.Add("segment_id");
    reduceKeys.Add("secondary_segment_id");
    auto sortKeys = reduceKeys;
    sortKeys.Add("name");
    sortKeys.Add("part_index");
    {
        auto tx = client->StartTransaction();
        NYT::TMapOperationSpec mapperSpec;
        mapperSpec.AddInput<NYT::TNode>(params.GetIndexDir());
        mapperSpec.AddOutput<NYT::TNode>(mapperOutput);
        tx->Map(mapperSpec, new TSegmentsMapper());
        tx->Commit();
    }
    {
        auto tx = client->StartTransaction();
        NYT::TSortOperationSpec spec;
        spec.AddInput(mapperOutput);
        spec.Output(sortedOutput);
        spec.SortBy(sortKeys);
        tx->Sort(spec);
        tx->Commit();
    }
        ui64 memoryLimit = 10ULL * 1024 * 1024 * 1024;
    {
        auto tx = client->StartTransaction();
        NYT::TReduceOperationSpec spec;
        NYT::TUserJobSpec userSpec;
        spec.AddInput<NSaas::TYTBlobBase>(sortedOutput);
        spec.AddOutput<NYT::TNode>(params.GetOutput());
        spec.ReduceBy(reduceKeys);
        spec.SortBy(sortKeys);
        userSpec.MemoryLimit(memoryLimit);
        spec.ReducerSpec(userSpec);
        INFO_LOG << "Spawning Reduce operation..." << Endl;
        NYT::IOperationPtr reduceJob = tx->Reduce(spec, new TFullarcReducer(params));
        tx->Commit();
        INFO_LOG << "Reduce operation completed" << Endl;
    }
    if (tmpDirCreated) {
        client->Remove(tmpDir, NYT::TRemoveOptions().Recursive(true).Force(true));
    }
}

int main(int argc, const char* argv[]) {
    if (!GetEnv("YT_LOG_LEVEL")) {
        SetEnv("YT_LOG_LEVEL", "INFO");
    }
    try {
        NYT::Initialize(argc, argv);
        NUtils::TParams params = NGetoptPb::GetoptPbOrAbort(argc, argv);
        auto mode = params.GetMode();
        if (mode == "local") {
            LocalArchiveDump(params);
        } else {
            YTArchiveDump(params);
        }
    } catch (...) {
        ERROR_LOG << CurrentExceptionMessage() << Endl;
        return 1;
    }

    return 0;
}
