#include "banned_addresses.h"

namespace NSrvKernel {

    TBannedAddresses::TBannedAddresses(TSharedStatsManager& statsManager)
        : AtomicAddressesCount_(0)
        , BannedAddressesCount_(statsManager.MakeGauge("ban-addresses").AllowDuplicate().Build())
        , BannedConnectionsCount_(statsManager.MakeCounter("ban-conns").AllowDuplicate().Build())
    {}

    bool TBannedAddresses::IsBanned(const NAddr::IRemoteAddr& addr) noexcept {
        if (!AtomicAddressesCount_) {
            return false;
        }
        with_lock(Mutex_) {
            bool isBanned = Banned_.contains(TCacheKey(addr));
            if (isBanned) {
                BannedConnectionsCount_.Inc();
            }
            return isBanned;
        }
    }

    bool TBannedAddresses::Add(const NAddr::IRemoteAddr& addr, TInstant deadline, size_t limit) noexcept {
        TCacheKey key = TCacheKey(addr);
        with_lock(Mutex_) {
            if (Banned_.size() >= limit && !Banned_.contains(key)) {
                return false;
            }
            auto res = Banned_.emplace(key, deadline);
            if (res.second) { // new address
                Ttls_[deadline].insert(key);
            } else { // probably, update deadline
                if (res.first->second < deadline) {
                    auto it = Ttls_.find(res.first->second);
                    if (it != Ttls_.end()) {
                        it->second.erase(key);
                        if (it->second.empty()) {
                            Ttls_.erase(it);
                        }
                    }
                    Ttls_[deadline].insert(key);
                    res.first->second = deadline;
                }
            }
            AtomicAddressesCount_ = Banned_.size();
            BannedAddressesCount_.Set(Banned_.size());
        }
        return true;
    }

    void TBannedAddresses::Clear() noexcept {
        with_lock(Mutex_) {
            Ttls_.clear();
            Banned_.clear();
            AtomicAddressesCount_ = Banned_.size();
            BannedAddressesCount_.Set(Banned_.size());
        }
    }

    void TBannedAddresses::ProcessTtl(TInstant deadline) noexcept {
        while (true) {
            with_lock(Mutex_) {
                auto it = Ttls_.begin();
                if (it == Ttls_.end() || it->first > deadline) {
                    break;
                }
                auto kit = it->second.begin();
                Banned_.erase(*kit);
                it->second.erase(kit);
                if (it->second.empty()) {
                    Ttls_.erase(it);
                }
                AtomicAddressesCount_ = Banned_.size();
                BannedAddressesCount_.Set(Banned_.size());
            }
        }
    }

}
