#pragma once

#include "backend.h"
#include "blacklist.h"
#include "pessimization.h"

#include <balancer/kernel/balancing/config.cfgproto.pb.h>
#include <balancer/kernel/coro/coroutine.h>
#include <balancer/kernel/ctl/children_process_common.h>
#include <balancer/kernel/ctl/ctl.h>
#include <balancer/kernel/fs/shared_files.h>
#include <balancer/kernel/log/log.h>
#include <balancer/kernel/module/events.h>
#include <balancer/kernel/pinger/pinger2.h>
#include <balancer/kernel/pinger/pinger_backend_group_manager.h>
#include <balancer/kernel/stats/manager.h>
#include <balancer/kernel/thread/threadedqueue.h>

#include <library/cpp/coroutine/engine/impl.h>
#include <library/cpp/logger/file.h>
#include <util/digest/murmur.h>
#include <util/string/builder.h>
#include <util/string/cast.h>
#include <util/system/mutex.h>
#include <util/system/rwlock.h>


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

    class TBackendGroupUpdater : public TMoveOnly {
    private:
        struct TPingerHolder {
            THolder<TPingerV2> Pinger;
            TIntrusivePtr<TDynamicBackend> Backend;
            TCoroutine Coroutine;
        };

    public:
        TBackendGroupUpdater(
            const TVector<TBackendDescriptor::TRef>& initialBackends,
            TConfig config,
            TSharedStatsManager& statsManager,
            TPingerStatsCounters* pingerStats,
            THotSwap<TBlacklist>& blacklist,
            TBackendGroupUpdater* old
        );

        ~TBackendGroupUpdater();

        TCoroutine Start(IWorkerCtl* process) noexcept;

        void ChangeNotReadySet(TIntrusivePtr<TNotReadyBackendSet> set) noexcept {
            NotReadySet_.AtomicStore(set);
        }

        TIntrusivePtr<TNotReadyBackendSet> GetNotReadySet() const noexcept {
            return NotReadySet_.AtomicLoad();
        }

        TIntrusivePtr<TBlacklist> GetCurrentBlacklist() const noexcept {
            return CurrentBlacklist_.AtomicLoad();
        }

        TMaybe<TPingerConfig> PingerConfig() const noexcept {
            return PingerConfig_;
        }

        void Stop() noexcept {
            // This is required, because Pingers_ is cleared only in destructor of dynamic module,
            // which will be called after TLS disposal.
            Pingers_.clear();
        }

        TIntrusivePtr<TBalancingState> AcquireState() const noexcept {
            return CurrentBalancingState_.AtomicLoad();
        }

        const TString& Name() {
            return Config_.backends_name();
        }

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

    private:
        void Run() noexcept;
        void CreatePingerFor(TIntrusivePtr<TDynamicBackend> backend) noexcept;
        void UpdateNow() noexcept;
        void ToggleOverPessimizedMark() noexcept;

    private:
        TConfig Config_;
        TMaybe<TPingerConfig> PingerConfig_;
        TDynamicLog Log_;
        TPessimizationTracker PessimizationTracker_;

        // Set in Start().
        IWorkerCtl* UpdaterProcess_ = nullptr;

        THashMap<TString, THolder<TPingerHolder>> Pingers_;

        THotSwap<TBalancingState> CurrentBalancingState_;
        THotSwap<TNotReadyBackendSet> NotReadySet_;

        bool WasOverPessimized_ = false;

        // Counters.
        TSharedCounter OverPessimizedCount_;
        TSharedHistogram UpdateTimes_;
        TPingerStatsCounters* PingerStats_ = nullptr;

        THotSwap<TBlacklist>& CurrentBlacklist_;
        bool Outdated_ = false;
        TMutex Mutex_;
    };

    class TUpdaterManager : public TMoveOnly {
    public:
        TUpdaterManager(
            TSharedStatsManager& statsManager,
            TString backendsBlacklistFile,
            TPingerStatsCounters* pingerStats
        )
            : SharedStatsManager_(statsManager)
            , BackendsBlacklistFile_(std::move(backendsBlacklistFile))
            , PingerStats_(pingerStats)
        {}

        THolder<TBackendGroupUpdater> RegisterBackends(
            size_t,
            const TVector<TBackendDescriptor::TRef>& backends,
            TConfig config,
            TBackendGroupUpdater* old
        );

        void Start(IWorkerCtl* process) noexcept;

        bool RequiresWorkerStart() const noexcept {
            return RequiresWorkerStart_;
        }

        void SetSharedFiles(TSharedFiles* sharedFiles) noexcept {
            SharedFiles_ = sharedFiles;
        }

        void DumpBalancingState(TEventData&) const noexcept {}

    private:
        TSystemLog Log_;
        TSharedStatsManager& SharedStatsManager_;
        TSharedFiles* SharedFiles_ = nullptr;
        TString BackendsBlacklistFile_;
        TSharedFileReReader BlacklistReader_;
        bool RequiresWorkerStart_ = false;
        TPingerStatsCounters* PingerStats_ = nullptr;
        THotSwap<TBlacklist> CurrentBlacklist_;

        TCoroutine BlacklistFileUpdater_;
    };
}
