#pragma once

#include "table_registry.h"
#include <robot/jupiter/library/opt/proc_args.h>
#include <saas/rtyserver_jupi/library/rtyt/config.pb.h>
#include <robot/jupiter/library/opt/mropt.h>
#include <robot/jupiter/library/opt/common/common.h>
#include <robot/jupiter/library/opt/arc_merge.h>
#include <robot/jupiter/library/opt/shard_merge.h>

namespace NFusion {

class TRTYTTaskSetup {
    TRTYTBuilderConfig BuilderParams;
    THashMap<TString, int> ProcessorIdByName;
    TIntrusivePtr<NRTYT::TClient> ClientPtr;
    TVector<TFsPath> Inputs;
    TFsPath Output;
    TFsPath TempDir;
    NJupiter::TShardsPrepareTables ShardsPrepareTables;
    TTableRegistry Registry;

public:
    TRTYTTaskSetup(
        TRTYTBuilderConfig builderParams,
        const TVector<TString>& inputs,
        const TString& output,
        const TString& tmp)
        : BuilderParams(std::move(builderParams))
        , ClientPtr(NRTYT::MakeClient())
        , Output(output)
        , TempDir(tmp)
        , ShardsPrepareTables(ClientPtr, "//src/", "", 1) {
        for (const auto& input : inputs) {
            Inputs.emplace_back(input);
        }

        ClientPtr->Create("//src", NYT::ENodeType::NT_MAP);
        ClientPtr->Mount(TempDir, "//tmp");
        for (const auto& input : Inputs) {
            ClientPtr->Create("//src/" + input.Basename() + "/shards_prepare",
                                NYT::ENodeType::NT_MAP, NYT::TCreateOptions().Recursive(true));
            ClientPtr->Mount(input, "//src/" + input.Basename() + "/shards_prepare/0");
        }
        INFO_LOG << "register tables start" << Endl; 
        ShardsPrepareTables.RegisterTables(Registry);

        const auto& operations = BuilderParams.GetOperations().GetOperation();
        for (int i = 0; i < operations.size(); i++) {
            ProcessorIdByName[operations[i].GetName()] = i;
        }
    }
    
    TRTYTTaskSetup(TRTYTTaskSetup&& other)
        : BuilderParams(std::move(other.BuilderParams))
        , ProcessorIdByName(std::move(other.ProcessorIdByName))
        , ClientPtr(std::move(other.ClientPtr))
        , Inputs(std::move(other.Inputs))
        , Output(std::move(other.Output))
        , TempDir(std::move(other.TempDir))
        , ShardsPrepareTables(std::move(other.ShardsPrepareTables))
        , Registry(std::move(other.Registry))
    {
        other.ClientPtr.Reset();
    }

    ~TRTYTTaskSetup() {
        if (!ClientPtr) {
            return;
        }
        try {
            INFO_LOG << "Unmounting tmp node..." << Endl;
            ClientPtr->Umount(ClientPtr->GetNodeId("//tmp"));
            for (const auto& input : Inputs) {
                INFO_LOG << "Unmounting input node //src/" << input.Basename() << "/shards_prepare/0" << Endl;
                ClientPtr->Umount(ClientPtr->GetNodeId("//src/" + input.Basename() + "/shards_prepare/0"));
            }
        } catch (...) {
            VERIFY_WITH_LOG(false, "%s", CurrentExceptionMessage().data());
        }
    }
    TIntrusivePtr<NRTYT::TClient> GetClientPtr() const;
    NJupiter::TCommonOpts GetCommonOpts(const TString& processorName, const TFsPath& inputPath) const;
    NJupiter::TArcMergeOpts GetArcMergeOpts() const;
    NJupiter::TShardMergerOpts GetShardMergerOpts(const TString& processorName, const TFsPath& inputPath) const;
    NJupiter::TProcArgs GetArgsBundle(const TString& processorName, const TFsPath& inputPath) const;
    const TRTYTBuilderConfig::TBuildOperation& GetBuildOperationParams(const TString& processorName) const;

    const TRTYTBuilderConfig& GetBuilderParams() const {
        return BuilderParams;
    }

    TTableRegistry& GetRegistry() {
        return Registry;
    }

    const TVector<TFsPath>& GetInputs() const {
        return Inputs;
    }

    const TFsPath& GetTempDir() const {
        return TempDir;
    }
    
    const TFsPath& GetOutput() const {
        return Output;
    }
    bool HasProcessor(const TString& name) const {
        return ProcessorIdByName.contains(name);
    }
};

} // namespace NFusion
