#pragma once

#include <balancer/yt/log-tool/controller/operations/convert.h>
#include <balancer/yt/log-tool/controller/operations/last.h>
#include <balancer/yt/log-tool/controller/operations/stat.h>
#include <balancer/yt/log-tool/proto/schema.pb.h>

#include <util/generic/string.h>

namespace NBalancerYt {
    static const TString LOG_NAME = "log";
    template<typename TInputFormat = NBalancerYt::NProto::TProdLog>
    class TLogProcessController {
        using TBackendStat = NBalancerYt::NProto::TBackendStat;
        using TLog = NBalancerYt::NProto::TLog;
        using TResolvedLog = NBalancerYt::NProto::TResolvedLog;
    public:
        TLogProcessController(
            TString server,
            TString inputTable,
            TString outputPrefix,
            NYT::TMapReduceOperationSpec mapReduceSpec = {}
        )
            : Server_(std::move(server))
            , InputTable_(std::move(inputTable))
            , OutputPrefix_(std::move(outputPrefix))
            , MapReduceSpec_(std::move(mapReduceSpec)) {
        }

        void InitSandbox()
        {
            auto client = NYT::CreateClient(Server_);
            client->Create(
                OutputPrefix_,
                NYT::NT_MAP,
                NYT::TCreateOptions()
                    .Recursive(true)
                    .IgnoreExisting(true)
            );
        }

        void PrepareLogs() {
            NYT::IClientPtr client = NYT::CreateClient(Server_);

            NYT::TMapReduceOperationSpec spec = MapReduceSpec_;
            spec.AddInput<TInputFormat>(InputTable_);
            spec.AddOutput<TResolvedLog>(PreparedLogPath());
            spec.ReduceBy("req-id");
            client->MapReduce(
                spec,
                new TConvertNode<TInputFormat, TResolvedLog>,
                new TLastTimeNode<TResolvedLog>
            );
        }

        void BuildStat() {
            NYT::IClientPtr client = NYT::CreateClient(Server_);

            NYT::TMapReduceOperationSpec spec = MapReduceSpec_;
            spec.AddInput<TResolvedLog>(PreparedLogPath());
            spec.AddOutput<TBackendStat>(LogStatPath());
            spec.ReduceBy("destination");
            client->MapReduce(
                spec,
                new TLogStatNode,
                new TLogStatMergeNode,
                new TLogStatMergeNode
            );
        }

        TString LogStatPath() {
            return OutputPrefix_ + "/stat";
        }

        TString PreparedLogPath() {
            return OutputPrefix_ + "/log";
        }

    private:
        TString Server_;
        TString InputTable_;
        TString OutputPrefix_;
        NYT::TMapReduceOperationSpec MapReduceSpec_;
    };
}
