#include "yt_merging_reducer.h"

#include "fullarc_toolkit.h"

#include <saas/tools/standalone_indexer/lib/standalone_indexer.h>

#include <saas/rtyserver/config/config.h>
#include <saas/rtyserver/config/indexer_config.h>
#include <saas/rtyserver/config/realm_config.h>

#include <saas/library/yt/file/yt_file_dumper.h>
#include <saas/library/yt/stream/yt_chunked_input.h>

#include <util/folder/filelist.h>
#include <util/folder/path.h>

REGISTER_REDUCER(TSaasYTMergingReducer);

static bool DirIsEmpty(const TFsPath& dir);
static void RenameSourceSegments(const TFsPath& indexDir, bool isPrep);
static void DumpIndexSegment(const TFsPath& resultDir, const TString& segment, const TYTChunkedOutput::TConstructionContext& dataContext, const TMaybe<TYTChunkedOutput::TConstructionContext>& renameContext);
static void YTRenameFile(const TYTChunkedOutput::TConstructionContext& context, const TString& from, const TString& to);

void TSaasYTMergingReducer::Do(TReader* input, TWriter* output) {
    auto config = ConfigBundle.Parse();
    TFsPath indexDir = config->IndexDir;
    indexDir.ForceDelete();
    indexDir.MkDir();

    TShardId shardId;
    {
        TYTChunkedInputToFiles inputProcessor(input, indexDir);
        inputProcessor.ProcessAll();
        shardId = inputProcessor.GetShardId();
    }

    TYTChunkedOutput::TConstructionContext dataContext({output, shardId, 0});
    TMaybe<TYTChunkedOutput::TConstructionContext> renameContext;

    if (UseDatalessParts) {
        CreateFakeFullArcParts(indexDir);
        renameContext.ConstructInPlace(TYTChunkedOutput::TConstructionContext{output, shardId, 1});
    }

    RenameSourceSegments(indexDir, config->GetRealmListConfig().GetMainRealmConfig().GetIndexerConfigDisk().PreparatesMode);

    TFsPath resultsDir = TFsPath::Cwd() / "results";
    resultsDir.ForceDelete();
    resultsDir.MkDir();

    TStandaloneIndexer standaloneIndexer(ConfigBundle.Text, ConfigBundle.Variables, RunModules);
    standaloneIndexer.EnableMerging();
    ui32 destSegmentId = 0;
    while (TVector<TString> segments = standaloneIndexer.GetSegmentsForMerge()) {
        TString destSegmentName = Sprintf("index_%010d_%010d", 0, destSegmentId);
        standaloneIndexer.MergeTo(std::move(segments), resultsDir, destSegmentName);
        DumpIndexSegment(resultsDir, destSegmentName, dataContext, renameContext);
        TFsPath(resultsDir / destSegmentName).ForceDelete();
        ++destSegmentId;
    }
    standaloneIndexer.DisableMerging();

    // We assume that all segments were merged and moved to the target directory. In theory some input segments may be
    // bigger than MaxDocumentsToMerge and therefore not covered by GetSegmentsForMerge(). We consider that being
    // ruled-out by IndexingMapper job spec that limits the size of input. If we need to handle such segments in the
    // future then they could be moved as-is from indexDir to resultsDir after ensuring they're valid index segments,
    // not some left-over garbage from a merge operation or whatever.
    Y_VERIFY(DirIsEmpty(indexDir), "merge is complete but index directory is not empty");

    Y_VERIFY(DirIsEmpty(resultsDir), "merge results were dumped into YT output but results/ directory is not empty");
}

static bool DirIsEmpty(const TFsPath& dir) {
    // TODO: move to util
    TFileEntitiesList fl(TDirsList::EM_FILES_DIRS_SLINKS);
    fl.Fill(dir);
    return (fl.Size() == 0);
}

static void RenameSourceSegments(const TFsPath& indexDir, bool isPrep) {
    const char* const prefix = isPrep ? "prep_" : "";
    try {
        TDirsList segmentsList;
        segmentsList.Fill(indexDir);
        ui32 segmentId = 0;
        while (TString segment = segmentsList.Next()) {
            TString segmentName = Sprintf("%sindex_%010d_%010d", prefix, 0, segmentId);
            TFsPath(indexDir / segment).RenameTo(indexDir / segmentName);
            ++segmentId;
        }
    } catch (yexception& e) {
        e << " in RenameSourceSegments('" << indexDir << "')";
        throw;
    }
}

static void DumpIndexSegment(const TFsPath& resultDir, const TString& segment, const TYTChunkedOutput::TConstructionContext& dataContext, const TMaybe<TYTChunkedOutput::TConstructionContext>& renameContext) {
    try {
        TFileList fl;
        fl.Fill(resultDir / segment, true);
        TVector<TString> filesToDump;
        while (TString fileName = fl.Next()) {
            if (renameContext.Defined() && IsFullArcPartFile(fileName)) {
                TString origin = GetFullArcPartOrigin(resultDir / segment / fileName);
                YTRenameFile(*renameContext, origin, TFsPath(segment) / fileName);
            } else {
                filesToDump.push_back(fileName);
            }
        }
        TYTFileDumper dataDumper(segment, dataContext);
        dataDumper.DumpFileList(resultDir / segment, filesToDump);
    } catch (yexception& e) {
        e << " in DumpIndexSegment('" << resultDir << "', '" << segment << "', ...)";
        throw;
    }
}

static void YTRenameFile(const TYTChunkedOutput::TConstructionContext& context, const TString& from, const TString& to) {
    TYTChunkedOutput out(context, from);
    out << to;
}
