#pragma once

#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/map.h>

class TSaasYTRenamingMapper
    : public NYT::IMapper<NYT::TTableReader<NYT::Message>, NYT::TTableWriter<NYT::Message>>
{
public:
    using TRenameMap = TMap<std::pair<TShardId, TString>, TString>;

public:
    TSaasYTRenamingMapper() = default;
    TSaasYTRenamingMapper(TRenameMap renameMap, EYTBlobFormat format = EYTBlobFormat::Old);
    TSaasYTRenamingMapper(NYT::TTableReader<NSaas::TYTBlobBase>* renameMapTable, EYTBlobFormat format = EYTBlobFormat::Old);
    void Do(TReader* input, TWriter* output) override;

    Y_SAVELOAD_JOB(RenameMap);

private:
    TRenameMap RenameMap;
    EYTBlobFormat Format;
};

class TRenameIndexPartsDataCommand : public TYTCommand {
public:
    TRenameIndexPartsDataCommand(TInputs inputs, TOutputs outputs, bool verbose, NSaas::TYTLaunchReport& report, const NYT::TNode& acl)
        : TYTCommand(std::move(inputs), std::move(outputs), "TRenameIndexPartsDataCommand", verbose, report, acl)
    {
    }
    virtual ~TRenameIndexPartsDataCommand() = default;

protected:
    virtual void DoPrepareSpec() override {
        Spec.AddInput<NSaas::TYTBlobBase>(Inputs.at(0))
            .AddOutput<NSaas::TYTBlobBase>(Outputs.at(0))
            .Ordered(true);
    }

    virtual void DoRun(NYT::IClientBase* client) override {
        EYTBlobFormat blobFormat = GetBlobFormatFromAttrs(client, Inputs.at(0));
        auto renameMapTable = client->CreateTableReader<NSaas::TYTBlobBase>(Inputs.at(1));
        Map(client, Spec, new TSaasYTRenamingMapper(renameMapTable.Get(), blobFormat));
    }

private:
    NYT::TMapOperationSpec Spec;
};
