#pragma once
#include <balancer/kernel/connection_manager_helpers/helpers.h>

namespace NConnectionManager {
    struct TDebugMetrics;
    class TSocketContainer {
    public:
        explicit TSocketContainer(THolder<NModProxy::TKeepAliveData> connection);
        TSocketContainer() {}

        THolder<NModProxy::TKeepAliveData> BuildKeepaliveData(TContExecutor* executor, NModProxy::EPollMode pollMode);

        TInstant Deadline() const noexcept {
            return KeepAliveReleaseDeadline_;
        }
        bool IsEmpty() const noexcept {
            return !SocketHolder_;
        }
        void SetDeadline(TInstant keepAliveReleaseDeadline) noexcept {
            KeepAliveReleaseDeadline_ = keepAliveReleaseDeadline;
        }
    private:
        NModProxy::TSockAddr SockAddr_;
        THolder<TSocketHolder> SocketHolder_;
        NModProxy::TSockAddrInfo SockAddrInfo_;
        TInstant KeepAliveReleaseDeadline_;
        THolder<NSrvKernel::TSslIo> SslIo_;
        TString SniHost_;
    };

    class TWorkerConnectionPool;
    class TConnectionPool;

    class TColdStorage {
    public:
        TColdStorage(TContExecutor* executor, TConnectionPool* owner, std::atomic<size_t>& stored);
        ~TColdStorage();
        void Add(THolder<NModProxy::TKeepAliveData> connection);
        THolder<NModProxy::TKeepAliveData> TryPull();
    private:
        void BackgroundTask();
        void ProcessSignalled();
        TContExecutor* Executor_;
        TConnectionPool* Owner_;
        std::atomic<size_t>& Stored_;
        NSrvKernel::TCoroSingleCondVar BackgroundCondVar_;
        TList<THolder<NModProxy::TKeepAliveData>> Connections_;
        THashMap<NModProxy::TKeepAliveData*, typename decltype(Connections_)::iterator> Signalled_;
        NSrvKernel::TCoroutine Background_;
    };

    class TConnectionPool : public TThrRefBase {
    public:
        TConnectionPool(TConnectionManager& parent, size_t tlsLimit, size_t connectionLimit, NSrvKernel::EPollMode pollMode, TDuration keepaliveTimeout, TDuration connectTimeout);
        ~TConnectionPool();
        TWorkerConnectionPool* GetPoolForWorker(NSrvKernel::IWorkerCtl* process);
        void ReleasePoolForWorker(NSrvKernel::IWorkerCtl* process);
        void SetOnExpire(std::function<void()> onExpire);

        bool AddConnection(THolder<NModProxy::TKeepAliveData>& connection);
        TSocketContainer TryPullConnection();
        TSocketContainer TryPullHotConnection();
        NSrvKernel::EPollMode PollMode() const noexcept;
        void MoveUnusedToColdStorage(TContExecutor* executor);
        void OnExpire();
    private:
        bool TryIncConnections();
        void DecConnections();
        std::deque<TSocketContainer> Storage_;
        TMutex Mutex_;
        TConnectionManager& Parent_;
        TVector<THolder<TWorkerConnectionPool>> Tls_;
        const size_t ConnectionLimit_;
        NSrvKernel::EPollMode PollMode_;
        TDuration KeepaliveTimeout_;
        TDuration ConnectTimeout_;
        std::atomic<size_t> Count_ = 0;
        std::atomic<size_t> StoredInCold_ = 0;
        std::atomic<size_t> RequestedFromCold_ = 0;
        std::function<void()> OnExpire_;
    };

    class TWorkerConnectionPool {
    public:
        TWorkerConnectionPool(TContExecutor* executor, TConnectionPool& parent);
        bool AddConnection(THolder<NModProxy::TKeepAliveData>& connection);
        THolder<NModProxy::TKeepAliveData> TryPullConnection();
    private:
        TContExecutor* Executor_;
        TConnectionPool& Parent_;
    };
}
