#pragma once

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

#include <saas/library/yt/common/yt_blob.h>

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

#include <util/generic/size_literals.h>

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

    Y_SAVELOAD_JOB(ConfigBundle, DumperChunkSize, UseDatalessParts, RunModules);

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

class TMergeIndexPartsCommand : public TYTCommand {
public:
    TMergeIndexPartsCommand(TInputs inputs, TOutputs outputs, bool verbose, NSaas::TYTLaunchReport& report, const TRtyConfigBundle& configBundle,
        const TYTConfigFiles& configFiles, ui32 dumperChunkSize, bool useDatalessParts, bool runModules,
        ui32 ramGb, float ramReserveFactor, ui32 tmpfsGb, ui32 cpu, const NYT::TNode& acl)
        : TYTCommand(std::move(inputs), std::move(outputs), "TMergeIndexPartsCommand", verbose, report, acl)
        , ConfigBundle(configBundle)
        , ConfigFiles(configFiles)
        , DumperChunkSize(dumperChunkSize)
        , UseDatalessParts(useDatalessParts)
        , RunModules(runModules)
        , RamGb(ramGb)
        , RamReserveFactor(ramReserveFactor)
        , TmpfsGb(tmpfsGb)
        , Cpu(cpu)
    {
    }
    virtual ~TMergeIndexPartsCommand() = default;

protected:
    virtual void DoPrepareSpec() override {
        JobSpec.CpuLimit(Cpu);
        JobSpec.MemoryLimit(1_GB * RamGb).ExtraTmpfsSize(1_GB * TmpfsGb);
        JobSpec.MemoryReserveFactor(RamReserveFactor);
        ConfigFiles.FillSpecFiles(JobSpec);
        Spec.AddInput<NSaas::TYTBlobBase>(Inputs.at(0))
            .AddOutput<NSaas::TYTBlobBase>(Outputs.at(0))
            .AddOutput<NSaas::TYTBlobBase>(Outputs.at(1))
            .ReducerSpec(JobSpec)
            .ReduceBy(ShardIdKeyColumns);
        Opts.MountSandboxInTmpfs(true);
    }

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

private:
    const TRtyConfigBundle& ConfigBundle;
    const TYTConfigFiles& ConfigFiles;
    const ui32 DumperChunkSize;
    const bool UseDatalessParts;
    const bool RunModules;
    const ui32 RamGb;
    const float RamReserveFactor;
    const ui32 TmpfsGb;
    const ui32 Cpu;

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