#pragma once

#include <balancer/kernel/cpu/cpu_limiter.h>
#include <balancer/kernel/stats/manager.h>
#include <util/datetime/base.h>
#include <util/system/mutex.h>
#include <util/generic/hash.h>
#include <util/generic/map.h>
#include <atomic>

namespace NSrvKernel {

class TBannedAddresses {
public:
    explicit TBannedAddresses(TSharedStatsManager& statsManager);
    bool Add(const NAddr::IRemoteAddr& addr, TInstant deadline, size_t limit) noexcept;
    bool IsBanned(const NAddr::IRemoteAddr& addr) noexcept;
    void Clear() noexcept;
    void ProcessTtl(TInstant deadline) noexcept;
private:
    using TCacheKey = TCpuLimiter::TCacheKey;
    TMutex Mutex_;
    THashMap<TCacheKey, TInstant> Banned_;
    TMap<TInstant, THashSet<TCacheKey>> Ttls_;
    std::atomic<size_t> AtomicAddressesCount_;
    TSharedCounter BannedAddressesCount_;
    TSharedCounter BannedConnectionsCount_;
};

}  // namespace NSrvKernel

