#pragma once

#include "blacklist.h"
#include "dynamic_backend.h"
#include "logger.h"

#include <balancer/kernel/balancing/config.cfgproto.pb.h>
#include <balancer/kernel/log/log.h>
#include <balancer/kernel/pinger/backend_info_helper.h>
#include <util/generic/set.h>
#include <util/generic/hash_set.h>


namespace NSrvKernel::NDynamicBalancing {
    using NSrvKernel::TSystemLog;

    struct TStateAggregate {
        ui64 TotalRequests = 0;
        double TotalResponseTimeAverage = 0;
    };

    bool IsBackendHealthy(TDynamicBackend& backend) noexcept;

    class TPessimizationTracker final : public TMoveOnly {
    public:
        TPessimizationTracker(const TConfig& config, TDynamicLog& log) noexcept
            : Config_(config)
            , Log_(log)
        {}

        // IsOverPessimized returns true if by applying all pessimizations of this tracker, max_pessimized_share would be exceeded; false otherwise.
        bool IsOverPessimized() const noexcept;

        // DiscoverNewBackends updates backend set of this tracker.
        void DiscoverNewBackends(const TVector<TBackendDescriptor::TRef>& newBackends, const TPessimizationTracker* donor);

        void UpdateNotReadySet(TIntrusivePtr<TNotReadyBackendSet> notReadySet) noexcept;
        void UpdateBlacklist(TIntrusivePtr<TBlacklist> blacklist) noexcept;

        // Backends returns all registered backends since last DiscoverNewBackends.
        const THashMap<TString, THolder<TDynamicBackendHolder>>& Backends() const noexcept {
            return Backends_;
        }

        // CalculateStateAggregate calculates statistics over all backends.
        TStateAggregate CalculateStateAggregate() noexcept;

        TVector<TIntrusivePtr<TDynamicBackend>> Step() noexcept;

        void LogInitialInfo() const noexcept;

        // RunBackendChecks calculates new backend status.
        TDynamicBackend::TStatus RunBackendChecks(TDynamicBackend& backend, TStateAggregate stateAggregate, bool& hasFiredCheck, bool& hasSkippedCheck) noexcept;

        void Dump(NJson::TJsonWriter& out) const noexcept;

    private:

        // UpdateBackend calculates new backend status using RunBackendChecks and updates internal structures.
        void UpdateBackend(TDynamicBackendHolder& backend, const TStateAggregate& aggregate) noexcept;

        void UpdateWeightCoeff(TDynamicBackendHolder& backend, bool canIncrease = true) noexcept;

        bool TryPessimizeBackend(TDynamicBackendHolder& backend, bool needReset = true) noexcept;
        void RecoverBackend(TDynamicBackendHolder& backend) noexcept;
        void UnpessimizeBackend(TDynamicBackendHolder& backend) noexcept;

        bool TrySupersedBackendFor(TDynamicBackend::TStatus newStatus) noexcept;
        void UpdatePessimizedBackends(TStateAggregate& stateAggregate) noexcept;
        bool IsBackendReady(const TDynamicBackend& backend) noexcept;
        void UpdateBlacklistedBackends() noexcept;

        void SetBackendStatusAndLog(TDynamicBackendHolder& backend, TDynamicBackend::TStatus newStatus) noexcept;

        bool StateAggregationsNeeded() noexcept {
            return Config_.has_response_time_check();
        }

        bool PessimizeIfNotReady(TDynamicBackendHolder& backend) noexcept;
        bool PessimizeIfNotHealthy(TDynamicBackendHolder& backend) noexcept;
        bool PessimizeIfDegraded(TDynamicBackendHolder& backend) noexcept;

    private:
        const TConfig& Config_;
        TDynamicLog& Log_;
        size_t MaxPessimizedBackendsCount_ = 0;
        THashMap<TString, THolder<TDynamicBackendHolder>> Backends_;
        size_t PessimizedBackendsCount_ = 0;
        // TODO(carzil): maybe separate list for RECOVERING backends for readability?
        std::array<TIntrusiveList<TDynamicBackendHolder>, static_cast<size_t>(TDynamicBackend::TStatus::Max)> PessimizedBackends_;
        TIntrusivePtr<TBlacklist> Blacklist_;
        TIntrusivePtr<TNotReadyBackendSet> NotReadySet_;
        bool OverPessimized_ = false;
        size_t Step_ = 0;
    };
}
