#include "component.h"

#include "config.h"
#include "builder.h"
#include "index_verifier.h"
#include "manager.h"
#include "parsed_entity.h"
#include "prep_normalizer.h"

#include <saas/rtyserver/config/config.h>
#include <saas/rtyserver/indexer_core/index_dir.h>
#include <saas/rtyserver/indexer_core/merger_interfaces.h>

#include <saas/rtyserver_jupi/library/extbuilder/wad_merge.h>
#include <saas/rtyserver_jupi/library/rtyt/bundle.h>
#include <saas/rtyserver_jupi/library/rtyt/task_setup.h>
#include <saas/rtyserver_jupi/library/tables/tables.h>

#include <robot/jupiter/library/rtdoc/file/docidmap_io.h>
#include <robot/jupiter/library/rtdoc/file/model/files.h>
#include <robot/jupiter/library/rtdoc/protos/builder_task.pb.h>

#include <google/protobuf/text_format.h>

namespace NRTYServer {

TJupiComponent::TJupiComponent(const TRTYServerConfig& config)
    : IIndexComponent(config.ComponentsSet.contains(JupiComponentName))
    , Config(config)
    , ComponentConfig(*(config.ComponentsConfig.Get<TJupiComponentConfig>(JupiComponentName)))
    , Normalizer(MakeHolder<TPrepNormalizer>(config))
{
    if (ComponentConfig.GetRtytEnabled()) {
        const auto& configPath = ComponentConfig.GetRtytBuilderConfigPath();
        if (configPath.Defined()) {
            INFO_LOG << "Parsing rtyt builder config from path " << configPath->GetPath();
            TUnbufferedFileInput reader(configPath->GetPath());
            google::protobuf::TextFormat::Parser protoParser;
            protoParser.AllowUnknownField(true);
            AssertCorrectConfig(protoParser.ParseFromString(reader.ReadAll(), &RTYTConfig), "Failed to parse RTYT config %s: not a valid protobuf text", configPath->GetPath().data());
        } else {
            RTYTConfig = GetDefaultRTYTConfig();
        }
    }
}

TString TJupiComponent::GetName() const {
    return JupiComponentName;
}

const IIndexFilesGroup::TIndexFiles& TJupiComponent::GetIndexFiles() const {
    return IndexFiles;
}

NFusion::TRTYTBuilderConfig TJupiComponent::GetDefaultRTYTConfig() {
    INFO_LOG << "Default config was requested!" << Endl;
    NFusion::TRTYTBuilderConfig config;
    config.MutableLumpsToIndex()->AddLump("walrus");
    config.MutableLumpsToIndex()->AddLump("external_relev_attrs");
    config.MutableLumpsToIndex()->AddLump("content_attrs");
    config.MutableLumpsToIndex()->AddLump("url_heavy_data");
    config.MutableLumpsToIndex()->AddLump("calculated_attrs");
    config.MutableLumpsToIndex()->AddLump("mobile_url");

    config.MutableLumpsToRemap()->AddLump("walrus");
    config.MutableLumpsToRemap()->AddLump("external_relev_attrs");
    config.MutableLumpsToRemap()->AddLump("content_attrs");
    config.MutableLumpsToRemap()->AddLump("url_heavy_data");
    config.MutableLumpsToRemap()->AddLump("calculated_attrs");
    config.MutableLumpsToRemap()->AddLump("mobile_url");

    NFusion::TRTYTBuilderConfig::TBuildOperation *buildExtInfoArc = config.MutableOperations()->AddOperation();
    buildExtInfoArc->SetName("BuildExtInfoArc");
    buildExtInfoArc->SetFilePrefix(""); // no need in this parameter
    buildExtInfoArc->SetFactory(NFusion::TRTYTBuilderConfig::SHARDS_PREPARE);

    NFusion::TRTYTBuilderConfig::TBuildOperation *mergeExtInfoArc = config.MutableOperations()->AddOperation();
    mergeExtInfoArc->SetName("MergeExtInfoArc");
    mergeExtInfoArc->SetFilePrefix("index.extinfo");
    mergeExtInfoArc->SetFullFileName("index.extinfoarc.wad");
    mergeExtInfoArc->SetFactory(NFusion::TRTYTBuilderConfig::SHARDMERGE_UTILS);

    INFO_LOG << "Total " << config.MutableLumpsToIndex()->GetLump().size() << "lumps to index" << Endl;
    return config;
}

//
// Creator methods
//

THolder<IIndexComponentBuilder> TJupiComponent::CreateBuilder(const TBuilderConstructionContext& context) const {
    TString tempPrefix = TFsPath(context.TempDir.PathName()) / NRtDoc::TBuilderFiles::MainPrepFile;
    return MakeHolder<TJupiBuilder>(tempPrefix, ComponentConfig.GetRtytEnabled(), RTYTConfig);
}

THolder<IIndexComponentManager> TJupiComponent::CreateManager(const TManagerConstructionContext& context) const {
    switch(context.IndexType) {
        case IIndexController::MEMORY:
            return nullptr;
        case IIndexController::FINAL:
        case IIndexController::DISK:
        case IIndexController::PREPARED:
            return MakeHolder<TJupiComponentManager>();
        default:
            FAIL_LOG("Incorrect index type");
    }
}

IComponentParser::TPtr TJupiComponent::BuildParser() const {
    return new TJupiDocParser{};
}

IParsedEntity::TPtr TJupiComponent::BuildParsedEntity(IParsedEntity::TConstructParams& params) const {
    return new TJupiParsedEntity(params);
}

//
// merger & normalizer
//
class TJupiMerger final : protected NRtDoc::TDocIdMapIo {
private:
    static TString FormatMergerMapFile(const TString& tempDir, const TPathName& srcDir) {
        return TString(TFsPath(tempDir) / TString(srcDir.BaseName())) + NRtDoc::TBuilderFiles::MergerMapSuffix;
    }
private:
    NRtDoc::TBuilderTask* BuilderTask; // TJupiMerger works as Builder for TBuilderTask
    bool RtytEnabled = false;

public:
    TJupiMerger(NRtDoc::TBuilderTask* task, bool rtytEnabled = false)
        : BuilderTask(task)
        , RtytEnabled(rtytEnabled)
    {
    }

    void PrepareInputs(const TRTYMergerContextBase& mc, ui32 destCl) {
        CHECK_WITH_LOG(mc.Decoder);
        const IRTYMergerDocIdDecoder& decoder = *mc.Decoder;
        ui32 totalDocCount = 0;
        ui32 newDocCount = 0;

        TString mountPrefix = "//" + TFsPath(mc.TempDir).Basename();

        TIntrusivePtr<NRTYT::TClientBase> rtytClient;
        if (RtytEnabled) {
            rtytClient = new NRTYT::TClientBase();
            rtytClient->Mount(mc.TempDir, mountPrefix);
        }

        // create and store the mappings
        for (ui32 sourceCl = 0; sourceCl < mc.Sources.size(); ++sourceCl) {
            NRtDoc::TDocIdMap mergerMap;
            TPathName srcDir(mc.Sources[sourceCl]);
            TFsPath mergerMapFile = FormatMergerMapFile(mc.TempDir, srcDir);
            TString onMergeMappingPath = FormatRTYTMergeMapPath(mc.TempDir, mc.Sources[sourceCl]) + "/docid_map";

            const ui32 sourceSize = decoder.GetSizeOfCluster(sourceCl);
            ui32 docCount = sourceSize;
            const bool isPrep = HasIndexDirPrefix(srcDir.BaseName(), DIRPREFIX_PREP);

            TVector<ui32>& remap = *mergerMap.MutableData();
            remap.resize(sourceSize);
            for (ui32 sourceId = 0; sourceId < sourceSize; ++sourceId) {
                TRTYMergerAddress target = decoder.Decode(sourceCl, sourceId);
                if (target.IsRemoved() || target.ClusterId != destCl) {
                    remap[sourceId] = NRtDoc::TDocIdMap::DeletedDocument();
                    --docCount;
                } else {
                    remap[sourceId] = target.DocId; // pruning support ends here, generally it should not be used
                }
            }
            if (docCount > 0) {
                totalDocCount += docCount;
                if (isPrep) {
                    newDocCount += docCount;
                }

                NRtDoc::TDocIdMapIo::Save(mergerMapFile, &mergerMap);

                if (BuilderTask) {
                    NRtDoc::TBuilderTask::TBuilderInput* input = BuilderTask->AddInputs();
                    input->SetSrcDir(srcDir.PathName());
                    input->SetSrcMapping(mergerMapFile);
                    input->SetIsFinalIndex(!isPrep);
                    input->SetDocCount(docCount);
                }
            } else {
                WARNING_LOG << "Skip segment " << srcDir.PathName() << " which consists of removed docs only. Segments size = " << sourceSize << " docs" << Endl;
            }


            // rtyt processing
            if (RtytEnabled) {
                INFO_LOG << "writing merge remap for " << srcDir.BaseName() << " ..." << Endl;
                try {
                    rtytClient->Mount(srcDir.PathName(), "//" + srcDir.BaseName());
                } catch (...) {
                    NOTICE_LOG << CurrentExceptionMessage() << Endl;
                    continue;
                }
                NYT::TCreateOptions options;
                options.Attributes(NYT::TNode()("row_count", remap.size())).Recursive(true);
                rtytClient->Create(onMergeMappingPath, NYT::ENodeType::NT_TABLE, options);
                auto writer = rtytClient->CreateTableWriter(NYT::TRichYPath(onMergeMappingPath).Append(true));
                for (size_t i = 0; i < remap.size(); i++) {
                    writer->AddRow(NYT::TNode()("old", i)("new", remap[i]));
                }
                rtytClient->Umount(rtytClient->GetNodeId("//" + srcDir.BaseName()));
                INFO_LOG << "writing merge remap for " << srcDir.BaseName() << " Success!" << Endl;
            }

            // end of rtyt processing
        }
        if (RtytEnabled) {
            rtytClient->Sync(rtytClient->GetNodeId(mountPrefix));
        }
        INFO_LOG << "Total doc count: " << totalDocCount << ", docs in prep segments: " << newDocCount << Endl;

        if (!newDocCount) {
            INFO_LOG << "Cancelling merge. All prep segments are empty" << Endl;
            ythrow TMergeCancelException();
        }

        if (BuilderTask) {
            BuilderTask->MutableOutput()->SetDocCountStat(totalDocCount);
        }
    }

    void PrepareOutputs(const TRTYMergerContextBase& mc, ui32 destCl, const TString& realm) {
        if (BuilderTask) {
            NRtDoc::TBuilderTask::TBuilderOutput* output = BuilderTask->MutableOutput();
            output->SetTrgDir(mc.Dests.at(destCl));
            output->SetTempDir(mc.TempDir);
            output->SetTmpfsDir(mc.TmpfsDests.at(destCl));
            // output->SetDocCountStat is done in PrepareInputs

            NRtDoc::TBuilderTaskId* id = BuilderTask->MutableId();
            id->SetBuilderStatName(realm);
        }
    }

    // call the plugin & wait for completion
    bool ExecuteSync(IExtBuilderClient& client, const TString &realm) {
        if (BuilderTask) {
            BuilderTask->SetRealm(realm);
            NRtDoc::TBuilderTaskResult result;
            client.Run(result, *BuilderTask);

            auto errorCode = result.GetErrorCode();

            if (Y_UNLIKELY(errorCode == 10)) {
                ythrow TMergeCancelException();
            }

            if (Y_UNLIKELY(errorCode != 0)) {
                ERROR_LOG << "IExtBuilderClient task failed: rc=" << result.GetErrorCode() << Endl;
                return false;
            }
        }
        return true;
    }
};

bool TJupiComponent::DoMerge(const TMergeContext& context) const {
    CHECK_WITH_LOG(context.Context.Decoder != nullptr);
    const TRTYMergerContextBase& mc = context.Context;
    const IRTYMergerDocIdDecoder& decoder = *mc.Decoder;
    if (mc.Sources.empty() || mc.Dests.empty() || mc.Decoder->GetNewDocsCount() == 0)
        return true; // empty


    if (!ExtBuilder) {
        auto conn = TExtBuilderEngine::WaitConnection(TDuration::Hours(4).ToDeadLine());
        VERIFY_WITH_LOG(!!conn, "TJupiComponent has no TExtBuilderEngine");
        ExtBuilder = conn->CreateClient();

        if (!ExtBuilder) {
            ERROR_LOG << "TJupiComponent cannot connect to TExtBuilderEngine" << Endl;
            return false;
        }
    }

    for (ui32 destCl = 0; destCl < mc.Dests.size(); ++destCl) {
        if (decoder.GetNewDocsCount(destCl) == 0)
            continue;

        NRtDoc::TBuilderTask task;
        TJupiMerger merger(&task, ComponentConfig.GetRtytEnabled());

        // Dump decoder tables and fill TBuilderTask::Inputs
        merger.PrepareInputs(mc, destCl);

        // fill TBuilderTask::Output
        TString realm = context.RealmName;
        merger.PrepareOutputs(mc, destCl, realm);

        //Run the external builder (& wait for completion)
        if (!merger.ExecuteSync(*ExtBuilder, realm)) {
            return false;
        }

        // Launching rtyt processing
        if (ComponentConfig.GetRtytEnabled()) {
            //read rtdoc results
            TJupiIndexVerifier indexVerifier(mc.Dests[destCl], ComponentConfig.GetVerifyIndexesRegex());

            //use rtyt to build indexes
            NFusion::TBuilder rtytBuilder(NFusion::TRTYTTaskSetup(RTYTConfig, mc.Sources, mc.Dests[destCl], mc.TempDir), ComponentConfig.GetTimerLogPath(), ToString(context.Realm));
            rtytBuilder.BuildIndex();

            //compare rtyt results with rtdoc
            indexVerifier.Verify(ComponentConfig.GetCrashOnDiff());
        }
        // End of rtyt processing
    }

    return true;
}

bool TJupiComponent::DoAllRight(const TNormalizerContext& context) const {
    const TFsPath frq = TFsPath(context.Dir.PathName()) / "indexfrq";
    THolder<TFileMap> fileIndexFrq = MakeHolder<TFileMap>(frq, TFileMap::oRdWr);
    return Normalizer->AllRight(context, fileIndexFrq);
}

void TJupiComponent::CheckAndFix(const TNormalizerContext& context) const {
    const TFsPath frq = TFsPath(context.Dir.PathName()) / "indexfrq";
    THolder<TFileMap> fileIndexFrq = MakeHolder<TFileMap>(frq, TFileMap::oRdWr);
    return Normalizer->Fix(context, fileIndexFrq);
}

}
