#include <balancer/serval/core/config.h>

#include <array>

static NSv::TAction CountCodes(const YAML::Node&, NSv::TAuxData& aux) {
    struct TCodeCountingStream : NSv::TStreamProxy {
    public:
        TCodeCountingStream(NSv::IStreamPtr s, const std::array<NSv::TNumber<ui64>*, 5>& counters) noexcept
            : NSv::TStreamProxy(std::move(s))
            , Counters(counters)
        {
        }

        bool WriteHead(NSv::THead& head) noexcept override {
            if (100 <= head.Code && head.Code < 600) {
                (*Counters[head.Code / 100 - 1])++;
            }
            return NSv::TStreamProxy::WriteHead(head);
        }

    private:
        const std::array<NSv::TNumber<ui64>*, 5>& Counters;
    };

    // TODO allow to list desired ranges and specific codes in config
    std::array<NSv::TNumber<ui64>*, 5> counters = {
        &aux.Signal("code-1xx_dmmm"),
        &aux.Signal("code-2xx_dmmm"),
        &aux.Signal("code-3xx_dmmm"),
        &aux.Signal("code-4xx_dmmm"),
        &aux.Signal("code-5xx_dmmm"),
    };
    return [counters](NSv::IStreamPtr& s) {
        s = std::make_shared<TCodeCountingStream>(std::move(s), counters);
        return true;
    };
}

SV_DEFINE_ACTION("count-codes", CountCodes);
