#pragma once

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

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

#include <util/generic/hash.h>
#include <util/generic/string.h>

namespace NBalancerYt {
    class TLogStatNode: public NYT::IMapper<NYT::TTableReader<NBalancerYt::NProto::TResolvedLog>, NYT::TTableWriter<NBalancerYt::NProto::TBackendStat>> {
        using TResolvedLog = NBalancerYt::NProto::TResolvedLog;
        using TBackendStat = NBalancerYt::NProto::TBackendStat;
    public:
        virtual void Do(
            NYT::TTableReader<TResolvedLog>* input,
            NYT::TTableWriter<TBackendStat>* output) override
        {
            THashMap<TString, TBackendStat> stats;
            auto pushStat = [] (auto& stat, const TResolvedLog& row) {
                stat.SetRequests(stat.GetRequests() + 1);
                stat.SetMeanDuration(stat.GetMeanDuration() + row.GetDuration());
            };
            for (; input->IsValid(); input->Next()) {
                const TResolvedLog& row = input->GetRow();
                TBackendStat& stat = stats[row.GetDestination()];
                stat.SetDestination(row.GetDestination());
                pushStat((*stat.MutableStat()), row);
                pushStat((*stat.MutableByError()->MutableReduced())[row.GetError()], row);
                pushStat((*stat.MutableByErrorType()->MutableReduced())[row.GetErrorType()], row);
                pushStat((*stat.MutableByResponseCode()->MutableReduced())[row.GetResponseCode()], row);
            }
            auto normalize = [] (auto& stat) {
                stat.SetMeanDuration(stat.GetMeanDuration() / stat.GetRequests());
            };
            for (auto& statPair : stats) {
                TBackendStat& stat = statPair.second;
                normalize(*stat.MutableStat());
                for (auto* statGroup : { stat.MutableByError(), stat.MutableByErrorType(), stat.MutableByResponseCode() }) {
                    for (auto& groupItem : *statGroup->MutableReduced()) {
                        normalize(groupItem.second);
                    }
                }
                output->AddRow(stat);
            }
        }
    };

    class TLogStatMergeNode: public NYT::IReducer<NYT::TTableReader<NBalancerYt::NProto::TBackendStat>, NYT::TTableWriter<NBalancerYt::NProto::TBackendStat>> {
        using TBackendStat = NBalancerYt::NProto::TBackendStat;
    public:
        virtual void Do(
            NYT::TTableReader<TBackendStat>* input,
            NYT::TTableWriter<TBackendStat>* output) override
        {
            THashMap<TString, TBackendStat> stats;
            auto pushStat = [] (TBackendStat::TStat& destStat, const TBackendStat::TStat& srcStat) {
                destStat.SetRequests(destStat.GetRequests() + srcStat.GetRequests());
                destStat.SetMeanDuration(destStat.GetMeanDuration() + srcStat.GetMeanDuration() * srcStat.GetRequests());
            };
            auto pushReducedStat = [&pushStat] (TBackendStat::TStatDict& destStat, const TBackendStat::TStatDict& srcStat) {
                for (const auto& pair : srcStat.GetReduced()) {
                    pushStat((*destStat.MutableReduced())[pair.first], pair.second);
                }
            };
            for (; input->IsValid(); input->Next()) {
                const TBackendStat& row = input->GetRow();
                TBackendStat& stat = stats[row.GetDestination()];
                if (!stat.GetDestination()) {
                    stat.SetDestination(row.GetDestination());
                }
                pushStat(*stat.MutableStat(), row.GetStat());
                pushReducedStat(*stat.MutableByError(), row.GetByError());
                pushReducedStat(*stat.MutableByErrorType(), row.GetByErrorType());
                pushReducedStat(*stat.MutableByResponseCode(), row.GetByResponseCode());
            }
            auto normalize = [] (auto& stat) {
                stat.SetMeanDuration(stat.GetMeanDuration() / stat.GetRequests());
            };
            for (auto& statPair : stats) {
                TBackendStat& stat = statPair.second;
                normalize(*stat.MutableStat());
                for (auto* statGroup : { stat.MutableByError(), stat.MutableByErrorType(), stat.MutableByResponseCode() }) {
                    for (auto& groupItem : *statGroup->MutableReduced()) {
                        normalize(groupItem.second);
                    }
                }
                output->AddRow(stat);
            }
        }
    };
    REGISTER_MAPPER(TLogStatNode);
    REGISTER_REDUCER(TLogStatMergeNode);
}
