#include "bundle.h"
#include "util.h"

#include <saas/rtyserver_jupi/library/extbuilder/wad_merge.h>
#include <saas/rtyserver_jupi/library/extbuilder/arc_merge.h>
#include <robot/jupiter/library/opt/shard_merge.h>
#include <robot/jupiter/tools/shards_prepare/lib/build_extinfo_arc.h>
#include <robot/jupiter/tools/shardmerge_utils/lib/document_index/arc_merge.h>

namespace NFusion {


void TBuilder::RunRemap(NJupiter::TTableBase* table,
                        const ::google::protobuf::Descriptor* descriptor,
                        const TString& srcName, const TVector<ui32>& mapping) {

    TString srcPath = "/" + srcName + "/" + table->GetName();

    if (TaskSetup.GetClientPtr()->Exists("//src" + srcPath)) {
        INFO_LOG << "Table " << table->GetPath().Path_ << "exists, ready to remap it" << Endl;
        NYT::TStructuredTablePath inputDescription(NYT::TRichYPath("//src" + srcPath), descriptor);
        NYT::TStructuredTablePath outputDescription(NYT::TRichYPath("//tmp" + srcPath + "_mapped"), descriptor);
        TIntrusivePtr<TRTYTDocIdRemapperBase> mapper = TaskSetup.GetRegistry().GetMapperByName(table->GetName());
        mapper->Reinitialize(table->GetName(), mapping);
        TaskSetup.GetClientPtr()->Map(
            NYT::TMapOperationSpec()
                .AddStructuredInput(inputDescription)
                .AddStructuredOutput(outputDescription),
            mapper
        );
        TaskSetup.GetClientPtr()->Sort(
            NYT::TSortOperationSpec()
                .AddInput("//tmp" + srcPath + "_mapped")
                .Output("//tmp" + srcPath)
                .SortBy({ "Shard", "Chunk", "LocalDocId" })
        );
    } else {
        INFO_LOG << "Table " << table->GetPath().Path_ << "doesn't exists, creating empty table instead of 'remapped'" << Endl;
        TaskSetup.GetClientPtr()->Create("//tmp" + srcPath, NYT::ENodeType::NT_TABLE, NYT::TCreateOptions().Recursive(true));
    }
}

void TBuilder::RunDocIdRemappings() {
    const auto tables = NJupiter::TShardsPrepareTables::Clientless();
    TTimerWrapper wrapper(*TimerLog, "RTYT.RunRemappings");
    
    for (const auto& srcDir : TaskSetup.GetInputs()) {
        if (IsPrepIndex(srcDir)) {
            TString prefix = srcDir.Basename() + "/shards_prepare/0";
            TaskSetup.GetClientPtr()->Create("//tmp/" + prefix, NYT::ENodeType::NT_MAP, NYT::TCreateOptions().Recursive(true));

            auto docidMap = ReadDocIdsFromRtyt(TaskSetup.GetClientPtr(), "//src/" + prefix + "/docid_map");
            for (const auto& tableName : TaskSetup.GetBuilderParams().GetLumpsToRemap().GetLump()) {
                RunRemap(TaskSetup.GetRegistry().GetTableByName(tableName),
                        TaskSetup.GetRegistry().GetProtoByName(tableName),
                        prefix, docidMap);
            }
        }
    }
}

void TBuilder::RunReducers() {
    TTimerWrapper wrapper(*TimerLog, "RTYT.RunReducers");
    for (const auto& srcDir : TaskSetup.GetInputs()) {
        if (IsPrepIndex(srcDir)) {
            if (TaskSetup.HasProcessor("BuildExtInfoArc")) {
                NJupiter::TProcArgs args = TaskSetup.GetArgsBundle("BuildExtInfoArc", srcDir);
                NJupiter::BuildExtInfoArc(args)->Do();
            }
        }
    }
}

void TBuilder::ExecuteBuilderCmd(const TString& processorName,
                                const std::function<void(NJupiter::TProcArgs& args)>& command) {
    THolder<NRtDoc::IWadMerger> merger = MakeHolder<NRtDoc::TWadMerger>();
    TTimerWrapper wrapper(*TimerLog, "RTYT.ExecuteBuilderCmd." + processorName);

    TString fullFileName = TaskSetup.GetBuildOperationParams(processorName).GetFullFileName();
    merger->Init(TaskSetup.GetOutput() / fullFileName);
    for (const auto& src : TaskSetup.GetInputs()) {
        const bool isPrep = IsPrepIndex(src);
        if (isPrep) {
            NJupiter::TProcArgs args = TaskSetup.GetArgsBundle(processorName, src);
            command(args);
        }
        TString prefix = "//tmp/" + src.Basename();
        THolder<NRtDoc::TDocIdMap> mapping = MakeHolder<NRtDoc::TDocIdMap>();
        *(mapping->MutableData()) = ReadDocIdsFromRtyt(TaskSetup.GetClientPtr(), prefix + "/docid_map");

        merger->Add(src / fullFileName, std::move(mapping), nullptr, src.Basename(), isPrep);
    }
    merger->Finish();
}


void TBuilder::RunBuildAndMerge() {
    if (TaskSetup.HasProcessor("MergeExtInfoArc")) {
        ExecuteBuilderCmd(
                        "MergeExtInfoArc",
                        [](NJupiter::TProcArgs& args) {
                            NJupiter::ExtInfoArcMerge(args);
                        });
    }
}

void TBuilder::BuildIndex() {
    RunDocIdRemappings();
    INFO_LOG << "Docids remapped!" << Endl;
    RunReducers();
    INFO_LOG << "Reducers reduced!" << Endl;
    RunBuildAndMerge();
}

} // namespace NFusion
