#pragma once

#include "function.h"

#include <balancer/serval/contrib/cone/mun.h>

#include <library/cpp/json/writer/json.h>

#include <util/generic/algorithm.h>
#include <util/generic/string.h>
#include <util/generic/vector.h>

#include <atomic>

namespace NSv {
    // An unistat writer. Must always output a `[name, value]` pair, where `value`
    // is either a single number or a list of `[upper bound, count]` histogram buckets.
    using TSignal = NSv::TFunction<void(NJsonWriter::TBuf&)>;

    // Write all unistat signals into a single JSON array. Each argument must be either
    // a single signal or a container of signals or a container of containers of ... etc.
    template <typename... Ts>
    static inline TString SerializeSignals(Ts&&... sigs) {
        NJsonWriter::TBuf b{NJsonWriter::HEM_UNSAFE};
        b.BeginList();
        auto write = [&](auto&& x, auto&& write) {
            if constexpr (std::is_invocable<decltype(x), NJsonWriter::TBuf&>::value) {
                x(b);
            } else {
                for (const auto& item : x) {
                    write(item, write);
                }
            }
        };
        int unpack[] = { 0, (write(sigs, write), 0)... };
        Y_UNUSED(unpack);
        return b.EndList().Str();
    }

    // Write a single number in a `[name, value]` format. Zeroes are skipped.
    template <typename T>
    struct TNumber : std::atomic<T> {
    public:
        TNumber(TString name, T initial = T{}) noexcept
            : std::atomic<T>(initial)
            , K_(std::move(name))
        {
        }

        template <typename F>
        void Update(F&& f) noexcept(noexcept(f(*this))) {
            for (T v = this->load(std::memory_order_acquire); !this->compare_exchange_weak(v, f(v)); ) {
            }
        }

        void operator()(NJsonWriter::TBuf& out) const {
            if (auto value = this->load(std::memory_order_relaxed)) {
                if constexpr (std::is_same<T, double>::value || std::is_same<T, float>::value) {
                    out.BeginList().WriteString(K_).WriteDouble(value).EndList();
                } else {
                    out.BeginList().WriteString(K_).WriteLongLong(value).EndList();
                }
            }
        }

    private:
        TString K_;
    };

    struct THistogram : TNonCopyable {
    public:
        // Construct a list of boundaries from a sequence of {stop, step} pairs.
        // For example, {{100, 20}, {300, 100}} -> {0, 20, 40, 60, 80, 100, 200, 300}.
        static TVector<double> MakeBoundaries(const std::initializer_list<ui64[2]>& ranges);

        // A set of boundaries spanning the range [0, 100_000_000] in some way.
        static const TVector<double> DefaultBoundaries;

        THistogram(TString name, const TVector<double> boundaries = DefaultBoundaries)
            : K_(std::move(name))
            , V_(new std::atomic<ui64>[boundaries.size()]())
            , B_(std::move(boundaries))
        {}

        std::atomic<ui64>& operator[](double value) {
            Y_ASSERT(value >= B_[0]);
            return V_[UpperBound(B_.begin(), B_.end(), value) - B_.begin() - 1];
        }

        void operator()(NJsonWriter::TBuf&) const;

    private:
        TString K_;
        THolder<std::atomic<ui64>, TDeleteArray> V_;
        const TVector<double> B_;
    };

    static inline auto Timer(THistogram* hgram) {
        auto record = [t = mun_usec_monotonic()](THistogram* h) {
            (*h)[mun_usec_monotonic() - t]++;
        };
        return std::unique_ptr<THistogram, decltype(record)>{hgram, std::move(record)};
    }

    static inline auto Timer(THistogram& hgram) {
        return Timer(&hgram);
    }
}
