#pragma once

#include <balancer/kernel/helpers/common_parsers.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/helpers/ranges.h>
#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/memory/shared.h>
#include <balancer/kernel/net/config_addrs.h>
#include <balancer/kernel/stats/common/balancer_stats.h>

#include <library/cpp/config/config.h>
#include <library/cpp/config/sax.h>
#include <library/cpp/coroutine/engine/impl.h>
#include <library/cpp/monlib/metrics/labels.h>
#include <library/cpp/monlib/metrics/metric_type.h>

#include <logbroker/unified_agent/client/cpp/logger/backend.h>

#include <util/generic/array_ref.h>
#include <util/generic/maybe.h>

namespace NSrvKernel {
    namespace {
        using TUnifiedAgentLogsCountersExternal = THashMap<TString, TIntrusivePtr<NUnifiedAgent::TClientCounters>>;
        using TUnifiedAgentStatsCounters = THashMap<TString, THashMap<TString, TCounter>>;
    }

    class TSharedAtomicsDispenser {
    public:
        TSharedAtomicsDispenser(TSharedAllocator& alloc)
                : Allocator_(alloc)
        {}

        void SetChildrenCount(size_t count) {
            Atomics_.resize(count + 1);
        }

        TAtomic* AllocateAtomic(size_t child);

    private:
        TSharedAllocator& Allocator_;
        TVector<TVector<TAtomic*>> Atomics_;
    };

    enum class TWorkerAggregation {
        Max,
        Min,
        Sum,
        Avg,
    };

    struct TSharedCounterOpts {
        TWorkerAggregation Aggregation;
        TMaybe<std::function<ui64(ui64)>> PostProcess;
        NMonitoring::EMetricType MetricType;
        TString Sigopt;
        bool IsGauge = false;
        ui64 DefaultValue = 0;
        bool AggregateMaster = true;
    };

    class TSharedCounter;

    class TSharedCounterHolder : public TMoveOnly {
    public:
        TSharedCounterHolder(TSharedAtomicsDispenser& dispenser, size_t workerCount, size_t extraWorkersCount,
                             TSharedCounterOpts opts) noexcept;

        TAtomic* GetAtomicForCounter(TMaybe<size_t> workerId = Nothing()) noexcept;

        void SetAtomicInSharedCounter(TSharedCounter& counter, TMaybe<size_t> workerId = Nothing()) noexcept;

        ui64 AggregateValues() const noexcept;

        ui64 AggregateSum() const noexcept;

        ui64 AggregateMin() const noexcept;

        ui64 AggregateMax() const noexcept;

        ui64 AggregateAvg() const noexcept;

        TAtomic*& PointerToAtomic() {
            return PointerToAtomic_;
        }

        const TSharedCounterOpts& Opts() const noexcept {
            return Opts_;
        }

    private:
        TVector<TAtomic*> Atomics_;
        TAtomic* PointerToAtomic_ = nullptr;

        const TSharedCounterOpts Opts_;
        const size_t WorkerCount_;
    };

    class TSharedCounter : public TMoveOnly {
    public:
        static_assert(sizeof(ui64) == sizeof(TAtomic), "");

        TSharedCounter(TAtomic* atomic = nullptr);

        explicit TSharedCounter(TSharedCounterHolder& counterHolder, TMaybe<size_t> workerId = Nothing());

        TSharedCounter(const TSharedCounter& counter, size_t workerId);

        void Set(ui64 x) noexcept;

        void Add(ui64 x) noexcept;

        void Sub(ui64 x) noexcept;

        void Inc() noexcept {
            Add(1);
        }

        void Dec() noexcept {
            Sub(1);
        }

        ui64 Value() const noexcept;

        TSharedCounter& operator++() noexcept {
            Add(1);
            return *this;
        }

        TSharedCounter& operator--() noexcept {
            Sub(1);
            return *this;
        }

        void SetAtomic(TAtomic* atomic) { Atomic_ = atomic; }

    private:
        TSharedCounterHolder* CounterHolder_ = nullptr;
        TAtomic* Atomic_ = nullptr;
    };

    class TCombinedCounter : public TMoveOnly {
    public:
        TCombinedCounter() = default;

        TCombinedCounter(const TCombinedCounter& other, size_t workerId);

        TCombinedCounter& operator++() noexcept {
            Inc();
            return *this;
        }

        TCombinedCounter& operator--() noexcept {
            Dec();
            return *this;
        }

        TCombinedCounter& operator+=(ui64 val) noexcept {
            Add(val);
            return *this;
        }

        TCombinedCounter& operator-=(ui64 val) noexcept {
            Sub(val);
            return *this;
        }

        void SetSharedCounter(TSharedCounter counter) {
            Y_ENSURE(!Counter_, "counter reassignment is forbidden");
            Counter_ = std::move(counter);
            Value_ = 0;
        }

        void Set(ui64 x) noexcept {
            Value_ = x;
            if (Counter_) {
                Counter_->Set(x);
            }
        }

        void Add(ui64 x) noexcept {
            if (Counter_) {
                Counter_->Add(x);
                Value_ = Counter_->Value();
            } else {
                Value_ += x;
            }
        }

        void Sub(ui64 x) noexcept {
            if (Counter_) {
                Counter_->Sub(x);
                Value_ = Counter_->Value();
            } else {
                Value_ -= x;
            }
        }

        void Inc() noexcept {
            Add(1);
        }

        void Dec() noexcept {
            Sub(1);
        }

        ui64 Value() const noexcept {
            if (Counter_) {
                return Counter_->Value();
            } else {
                return Value_;
            }
        }

    private:
        ui64 Value_ = 0;
        TMaybe<TSharedCounter> Counter_ = Nothing();
    };

    class TSharedStatsManager;

    class TSharedCounterBuilder {
    private:
        TMaybe<TWorkerAggregation> Aggregation_ = Nothing();
        TMaybe<std::function<ui64(ui64)>> PostProcess_ = Nothing();
        TMaybe<TString> Sigopt_ = Nothing();
        TMaybe<NMonitoring::EMetricType> MetricType_ = Nothing();
        TMaybe<NMonitoring::TLabels> SolomonLabels_ = Nothing();
        TMaybe<bool> AllowDuplicate_ = Nothing();
        TMaybe<ui64> DefaultValue_ = Nothing();
        TMaybe<bool> AggregateMaster_ = Nothing();

        TString Name_;
        TSharedStatsManager& StatsManager_;
        bool IsGauge_ = false;

    public:
        TSharedCounterBuilder(TSharedStatsManager& manager, TString name, bool gauge);
        TSharedCounter Build();

        TSharedCounterBuilder& Aggregation(TWorkerAggregation aggregation);

        TSharedCounterBuilder& AllowDuplicate(bool value = true);

        TSharedCounterBuilder& Default(ui64 value);

        TSharedCounterBuilder& Labels(NMonitoring::TLabels labels);

        TSharedCounterBuilder& PostProcess(std::function<ui64(ui64)> f);

        TSharedCounterBuilder& Sigopt(TString sigopt);

        TSharedCounterBuilder& Kind(NMonitoring::EMetricType kind);

        TSharedCounterBuilder& AggregateMaster(bool value);
    };

    struct TSharedHistogramOpts {
        NMonitoring::EMetricType MetricType;
        TString Sigopt;
        double Scale = 1;
    };

    class TSharedHistogram;

    class TSharedHistogramHolder : public TMoveOnly {
    public:
        TSharedHistogramHolder(TSharedAtomicsDispenser& dispenser, size_t workersCount,
                               TVector<ui64> intervals, TSharedHistogramOpts opts);

        TVector<TAtomic*> GetAtomicsForHistogram(TMaybe<size_t> workerId = Nothing()) noexcept;

        void SetAtomicsInSharedHistogram(TSharedHistogram& histogram, TMaybe<size_t> workerId = Nothing());

        TVector<ui64> AggregateValues() const noexcept;

        TVector<TVector<TAtomic*>>& WorkerAtomics() noexcept {
            return Atomics_;
        }

        const TVector<ui64>& Intervals() const noexcept {
            return Intervals_;
        }

        TVector<TAtomic*>& PointerToAtomics() {
            return PointerToAtomics_;
        }

        const TSharedHistogramOpts& Opts() const noexcept {
            return Opts_;
        }

    private:
        TVector<TVector<TAtomic*>> Atomics_;
        TVector<TAtomic*> PointerToAtomics_;

        TVector<ui64> Intervals_;
        TSharedHistogramOpts Opts_;
    };

    class TSharedHistogram : public TMoveOnly {
    public:
        explicit TSharedHistogram(TSharedHistogramHolder& histogramHolder, TMaybe<size_t> workerId = Nothing());

        TSharedHistogram(const TSharedHistogram& histogram, size_t workerId);

        void AddValue(ui64 value) noexcept;

        void SetHistogram(TVector<TAtomic*> Histogram) { Histogram_ = std::move(Histogram); }

    private:
        TSharedHistogramHolder* HistogramHolder_ = nullptr;
        TVector<TAtomic*> Histogram_;
    };

    class TSharedHistogramBuilder {
        friend class TSharedStatsManager;
    private:
        TMaybe<TString> Sigopt_ = Nothing();
        TMaybe<NMonitoring::EMetricType> MetricType_ = Nothing();
        TMaybe<NMonitoring::TLabels> SolomonLabels_ = Nothing();
        TMaybe<bool> AllowDuplicate_ = Nothing();
        TMaybe<double> Scale_ = Nothing();

        TString Name_;
        TSharedStatsManager& StatsManager_;
        TVector<ui64> Intervals_;

    public:
        TSharedHistogramBuilder(TSharedStatsManager& manager, TString name, TVector<ui64> intervals);
        TSharedHistogram Build();

        TSharedHistogramBuilder& AllowDuplicate(bool value = true) {
            Y_VERIFY(!AllowDuplicate_.Defined(), "Allow duplicate parameter is already set for histogram");
            AllowDuplicate_ = value;
            return *this;
        }

        TSharedHistogramBuilder& Labels(NMonitoring::TLabels labels) {
            Y_VERIFY(!SolomonLabels_.Defined(), "Solomon labels have been already set for histogram");
            SolomonLabels_ = std::move(labels);
            return *this;
        }

        TSharedHistogramBuilder& Scale(double value) {
            Y_VERIFY(!Scale_.Defined(), "Scale was already set for histogram");
            Scale_ = value;
            return *this;
        }

        TSharedHistogramBuilder& Sigopt(TString sigopt) {
            Y_VERIFY(!Sigopt_.Defined(), "Sigopt is already set for histogram");
            Sigopt_ = std::move(sigopt);
            return *this;
        }

        TSharedHistogramBuilder& Kind(NMonitoring::EMetricType kind) {
            Y_VERIFY(!MetricType_.Defined(), "Metric type is already set for histogram");
            MetricType_ = kind;
            return *this;
        }
    };

    class TSharedStatsManager : public TMoveOnly {
        friend class TSharedCounterBuilder;
        friend class TSharedHistogramBuilder;
    public:
        class TConfig : public NConfig::IConfig::IFunc {
        public:
            TConfigAddresses Addrs;
            TMaybe<size_t> MaxConns;
            bool HideLegacySignals = false;

        private:
            START_PARSE {
                if (key == "addrs") {
                    value->AsSubConfig()->ForEach(&Addrs);
                    return;
                }

                if (key == "max_conns") {
                    MaxConns = 0;
                    NConfig::ParseFromString(value->AsString(), MaxConns.GetRef());
                    return;
                }

                if (key == "hide_legacy_signals") {
                    HideLegacySignals = value->AsBool();
                    return;
                }
            } END_PARSE
        };
    public:
        TSharedStatsManager(TSharedAllocator& sharedAllocator)
            : SharedAllocator_(sharedAllocator)
            , AtomicsDispenser_(sharedAllocator)
        {}

        TSharedCounterBuilder MakeCounter(TString name) {
            return TSharedCounterBuilder{*this, std::move(name), false};
        }

        TSharedCounterBuilder MakeCounter(std::initializer_list<TStringBuf> subnames) {
            TString name;
            for (auto&& n : subnames) {
                if (n && name) {
                    name.append('-');
                }
                name.append(n);
            }
            return MakeCounter(name);
        }

        TSharedCounterBuilder MakeGauge(TString name) {
            return TSharedCounterBuilder{*this, std::move(name), true};
        }

        TSharedHistogramBuilder MakeHistogram(TString name, TVector<ui64> intervals) {
            return TSharedHistogramBuilder{*this, std::move(name), intervals};
        }

        void SetConfig(TConfig&& config) {
            Config_ = std::move(config);
        }

        bool HideLegacySignals() {
            return Config_.HideLegacySignals;
        }

        void SetWorkersCount(size_t workersCount, size_t extraWorkersCount) noexcept {
            WorkersCount_ = workersCount;
            ExtraWorkersCount_ = extraWorkersCount;
            AtomicsDispenser_.SetChildrenCount(workersCount + extraWorkersCount);
        }

        enum class TResponseType {
            Unistat,
            Solomon,
        };

        void WriteResponse(TCont* cont, TResponseType type, IOutputStream* out) noexcept;
        void WriteResponseNoWait(TResponseType type, IOutputStream* out);

        void DumpSignals() const;

        const TConfigAddresses& Addrs() const noexcept {
            return Config_.Addrs;
        }

        TMaybe<TSharedCounter> GetExistingCounter(TString signalName, NMonitoring::TLabels labels = {}, TMaybe<size_t> workerId = Nothing());

        TBalancerStats GetBalancerStats() const;

        /**
         * Get counters by name for specific log
         *
         * @param[in] name                   log name
         *
         * @return                           ref to Unified Agent counters
         */
        const TIntrusivePtr<NUnifiedAgent::TClientCounters>& GetUnifiedAgentCounters(const TStringBuf name);

        /**
         * Append counter from Unified Agent library
         *
         * @param[in] stats                  object with statistics
         */
        void AppendUnifiedAgentCounters(TBalancerStats& stats);
    private:
        TSharedCounter MakeIntSignal(TString signalName, NMonitoring::TLabels labels,
            bool allowDuplicate, TSharedCounterOpts opts);
        TSharedHistogram MakeHistogramImpl(TString signalName, NMonitoring::TLabels labels,
            bool allowDuplicate, TVector<ui64> intervals, TSharedHistogramOpts opts);
        void CheckState();
        TVector<TCounter> GetCounters() const;
        TVector<TCounter> TryGetCounters() const;
        TVector<THistogram> GetHistograms() const;
        TVector<THistogram> TryGetHistograms() const;

    private:
        TMutex Mutex_;
        TSharedAllocator& SharedAllocator_;
        TSharedAtomicsDispenser AtomicsDispenser_;
        THashMap<NMonitoring::TLabels, TSharedCounterHolder> Counters_;
        THashMap<NMonitoring::TLabels, TSharedHistogramHolder> Histograms_;
        TMaybe<size_t> WorkersCount_;
        TMaybe<size_t> ExtraWorkersCount_;
        TConfig Config_;
        TUnifiedAgentLogsCountersExternal UnifiedAgentLogsCounters_;
        TUnifiedAgentStatsCounters UnifiedAgentStatsCounters_;
    };
}

