#pragma once
#include <util/system/yassert.h>

#include <balancer/kernel/connection_manager/connection_pool.h>
#include <balancer/kernel/coro/waked.h>

namespace NConnectionManager {
    class TConnectionManager : public TThrRefBase {
    public:
        class TOptions: public NConfig::IConfig::IFunc {
          public:
            TOptions() = default;
            explicit TOptions(const NSrvKernel::ICtl& ctl, NConfig::IConfig& config);

            bool IgnorePerPoolLimit = false;
            size_t ConnectionsLimit = Max<size_t>();
            size_t WorkersCount = 0;

          private:
            START_PARSE {
                ON_KEY("capacity", ConnectionsLimit) {
                    return;
                }
                ON_KEY("ignore_per_proxy_limit", IgnorePerPoolLimit) {
                    return;
                }
            } END_PARSE
        };

        explicit TConnectionManager(TOptions options);
        TConnectionManager(NSrvKernel::ICtl const& ctl, NConfig::IConfig& config);

        size_t PreferredChannelCapacity() const noexcept;

        TConnectionPool* AddConnectionPool(size_t id, size_t limit, NSrvKernel::EPollMode pollMode, TDuration keepaliveTimeout, TDuration connectTimeout);
        void RemoveConnectionPool(size_t id);

        bool TrySendToRuntime(std::function<void()> func);
        bool TryIncConnections();
        void DecConnections();

        class TRuntime {
        public:
            TRuntime(TContExecutor* executor, TIntrusivePtr<TConnectionManager> parent);
            ~TRuntime();
            TColdStorage& EnsureColdStorage(TIntrusivePtr<TConnectionPool> pool, std::atomic<size_t>& counter);
            TColdStorage* TryGetColdStorage(TIntrusivePtr<TConnectionPool> pool);
            bool TrySend(std::function<void()> func);
            void AssignDisposal(TIntrusivePtr<TConnectionPool> pool);
        private:
            void Run();
            void ProcessDisposals();
            TContExecutor* Executor_;
            TIntrusivePtr<TConnectionManager> Parent_;
            NSrvKernel::TCoroutine Runtime_;
            TMutex Mutex_;
            TList<TIntrusivePtr<TConnectionPool>> ToDispose_;
            THashMap<TIntrusivePtr<TConnectionPool>, THolder<TColdStorage>> Storages_;
            NSrvKernel::TW2WChannel<std::function<void()>> Channel_;
        };

        THolder<TRuntime> StartRuntime(TContExecutor* executor);
        void ForgetRuntime();
        void MoveUnusedToColdStorage(TContExecutor* executor);
        TRuntime* RuntimeNoLock() const noexcept;

    private:
        TOptions Options_;
        std::atomic<size_t> ConnectionsCount_ = 0;
        TMutex Mutex_;
        THashMap<size_t, TIntrusivePtr<TConnectionPool>> ConnectionPoolMap_;
        TRuntime* Runtime_ = nullptr;
    };
}
