#pragma once

#include "config_files.h"
#include "yt_command.h"

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

#include <saas/rtyserver/components/fullarchive/config.h>
#include <saas/rtyserver/config/config.h>
#include <saas/library/yt/common/yt_blob.h>
#include <saas/library/yt/common/yt_document.h>

#include <mapreduce/yt/interface/operation.h>

#include <util/generic/size_literals.h>

class TSaasYTIndexingMapper
    : public NYT::IMapper<NYT::TTableReader<TYTDocument>, NYT::TTableWriter<NYT::Message>>
{
public:
    TSaasYTIndexingMapper() = default;
    TSaasYTIndexingMapper(const TRtyConfigBundle& configBundle, bool failOnIndexError, ui32 dumperChunkSize, bool useDatalessParts, bool runModules)
        : ConfigBundle(configBundle)
        , FailOnIndexError(failOnIndexError)
        , DumperChunkSize(dumperChunkSize)
        , UseDatalessParts(useDatalessParts)
        , RunModules(runModules)
    {
    }
    void Do(TReader* input, TWriter* output) override;

    Y_SAVELOAD_JOB(ConfigBundle, FailOnIndexError, DumperChunkSize, UseDatalessParts, RunModules);

private:
    TRtyConfigBundle ConfigBundle;
    bool FailOnIndexError = false;
    ui32 DumperChunkSize = 0;
    bool UseDatalessParts = false;
    bool RunModules = false;
};

class TIndexDocumentsCommand : public TYTCommand {
public:
    TIndexDocumentsCommand(TInputs inputs, TOutputs outputs, bool verbose, NSaas::TYTLaunchReport& report, const TRtyConfigBundle& configBundle,
        const TYTConfigFiles& configFiles, bool failOnIndexError, ui32 dumperChunkSize, bool useDatalessParts, bool runModules, const NYT::TNode& acl)
        : TYTCommand(std::move(inputs), std::move(outputs), "TIndexDocumentsCommand", verbose, report, acl)
        , ConfigBundle(configBundle)
        , ConfigFiles(configFiles)
        , FailOnIndexError(failOnIndexError)
        , DumperChunkSize(dumperChunkSize)
        , UseDatalessParts(useDatalessParts)
        , RunModules(runModules)
    {
    }
    virtual ~TIndexDocumentsCommand() = default;

protected:
    virtual void DoPrepareSpec() override {
        // The indexing job is expected to produce one index segment with one FULLARC part in it, but with multiple
        // FULLARC layers and with DDK index. We use the part size limit to calculate the upper limit on the incoming
        // data and to estimate needed RAM and TmpFS.
        ui64 fullLayerPartSizeLimit;
        ui32 layers;
        bool compressionEnabled;
        {
            DoInitGlobalLog("null", 0, false, false);   // suppress logging for config parsing
            auto config = ConfigBundle.Parse();
            auto faConfig = config->ComponentsConfig.Get<TRTYFullArchiveConfig>(FULL_ARCHIVE_COMPONENT_NAME);
            layers = faConfig->GetLayers().size();
            const auto& fullLayerConfig = faConfig->GetLayer(NRTYServer::NFullArchive::FullLayer);
            fullLayerPartSizeLimit = (fullLayerConfig.PartSizeLimit <= 1_GB) ? fullLayerConfig.PartSizeLimit : 1_GB;
            compressionEnabled = (fullLayerConfig.Compression == NRTYArchive::IArchivePart::COMPRESSED);
        }
        // memory indexer is turned off, so not much RAM is required
        const ui64 memoryLimit = 1_GB;
        const ui64 tmpfsSize = 1.5f * (fullLayerPartSizeLimit * (/* FULLARC */ layers + /*DDK*/ 1));
        const ui64 maxDataSizePerJob = compressionEnabled ? fullLayerPartSizeLimit : fullLayerPartSizeLimit / 2;

        JobSpec.MemoryLimit(memoryLimit).ExtraTmpfsSize(tmpfsSize);
        ConfigFiles.FillSpecFiles(JobSpec);
        Spec.AddInput<TYTDocument>(Inputs.at(0))
            .AddOutput<NSaas::TYTBlobBase>(Outputs.at(0))
            .AddOutput<NSaas::TYTBlobBase>(Outputs.at(1))
            .MapperSpec(JobSpec)
            .Ordered(true);
        Opts.MountSandboxInTmpfs(true).Spec(NYT::TNode()("max_data_size_per_job", maxDataSizePerJob));
        Opts.InferOutputSchema(true);
    }

    virtual void DoRun(NYT::IClientBase* client) override {
        Map(client, Spec, new TSaasYTIndexingMapper(ConfigBundle, FailOnIndexError, DumperChunkSize, UseDatalessParts, RunModules), Opts);
    }

private:
    const TRtyConfigBundle& ConfigBundle;
    const TYTConfigFiles& ConfigFiles;
    const bool FailOnIndexError;
    const ui32 DumperChunkSize;
    const bool UseDatalessParts;
    const bool RunModules;

    NYT::TMapOperationSpec Spec;
    NYT::TUserJobSpec JobSpec;
    NYT::TOperationOptions Opts;
};
