#pragma once

#include "cpu_usage.h"

#include <balancer/kernel/coro/coro_cond_var.h>
#include <balancer/kernel/coro/coroutine.h>
#include <balancer/kernel/fs/shared_files.h>
#include <balancer/kernel/stats/manager.h>

#include <library/cpp/cache/cache.h>
#include <library/cpp/config/sax.h>
#include <util/generic/list.h>


namespace NSrvKernel {

class TRejectedConnsCleaner {
public:
    struct TRejectedConn {
        TInstant Deadline;
        TSocketHolder S;
    };

    TRejectedConnsCleaner(
        TSharedCounter connHoldCounter,
        size_t holdCount,
        TDuration holdDuration
    )
        : ConnHoldCounter_(std::move(connHoldCounter))
        , HoldCount_(holdCount)
        , HoldDuration_(holdDuration)
    {}

    void OnInit(TContExecutor* executor);

    void Push(TSocketHolder&& socket);

    size_t Size() const {
        return RejectedConns_.size();
    }

private:
    void PushConn(TRejectedConn&& conn) noexcept {
        RejectedConns_.push_back(std::move(conn));
        ++ConnHoldCounter_;
    }

    void PopConn() noexcept {
        RejectedConns_.pop_front();
        --ConnHoldCounter_;
    }

    void RejectedConnsCleanerCont(TContExecutor* executor);

    TSharedCounter ConnHoldCounter_;

    TList<TRejectedConn> RejectedConns_;
    TCoroSingleCondVar NoRejectedConnsCV_;

    TCoroutine Cont_;

    const size_t HoldCount_;
    const TDuration HoldDuration_;
};

struct TCpuLimiterConfig {
    TString DisableFile;
    TString DisableHTTP2File;
    double CpuUsageCoeff = 0.05;

    size_t ConnHoldCount = 10000;
    TDuration ConnHoldDuration = TDuration::Seconds(10);
    size_t CheckerAddressCacheSize = 4096;

    double ConnRejectLo = 0.95;
    double ConnRejectHi = 1;
    double HTTP2DropLo = 0.9;
    double HTTP2DropHi = 0.95;
    double KeepAliveCloseLo = 0.95;
    double KeepAliveCloseHi = 1;

    bool EnableConnReject = false;
    bool EnableHTTP2Drop = false;
    bool EnableKeepAliveClose = false;
};

struct TCpuLimiterStat {
    explicit TCpuLimiterStat(TSharedStatsManager& statsManager)
        : AverageCpuUsageCounter(statsManager.MakeGauge("worker-average_cpu_usage").Build())
        , CpuLimiterConnRejectedCounter(statsManager.MakeCounter("worker-cpu_limiter_conn_rejected").Build())
        , CpuLimiterHTTP2DisabledCounter(statsManager.MakeCounter("worker-cpu_limiter_http2_disabled").Build())
        , CpuLimiterHTTP2ClosedCounter(statsManager.MakeCounter("worker-cpu_limiter_http2_closed").Build())
        , CpuLimiterKeepAliveClosedCounter(statsManager.MakeCounter("worker-cpu_limiter_keepalive_closed").Build())
        , ConnHoldCounter(statsManager.MakeCounter("worker-cpu_limiter_conn_hold").Build())
        , CheckerCacheMissCounter(statsManager.MakeCounter("worker-cpu_limiter_checker_cache_miss").Build())
        , CheckerCacheSizeCounter(statsManager.MakeGauge("worker-cpu_limiter_checker_cache_size").Build())
    {}

    TCpuLimiterStat(TCpuLimiterStat& stat, size_t workerId)
        : AverageCpuUsageCounter(stat.AverageCpuUsageCounter, workerId)
        , CpuLimiterConnRejectedCounter(stat.CpuLimiterConnRejectedCounter, workerId)
        , CpuLimiterHTTP2DisabledCounter(stat.CpuLimiterHTTP2DisabledCounter, workerId)
        , CpuLimiterHTTP2ClosedCounter(stat.CpuLimiterHTTP2ClosedCounter, workerId)
        , CpuLimiterKeepAliveClosedCounter(stat.CpuLimiterKeepAliveClosedCounter, workerId)
        , ConnHoldCounter(stat.ConnHoldCounter, workerId)
        , CheckerCacheMissCounter(stat.CheckerCacheMissCounter, workerId)
        , CheckerCacheSizeCounter(stat.CheckerCacheSizeCounter, workerId)
    {}

    TSharedCounter AverageCpuUsageCounter;
    TSharedCounter CpuLimiterConnRejectedCounter;
    TSharedCounter CpuLimiterHTTP2DisabledCounter;
    TSharedCounter CpuLimiterHTTP2ClosedCounter;
    TSharedCounter CpuLimiterKeepAliveClosedCounter;

    TSharedCounter ConnHoldCounter;

    TSharedCounter CheckerCacheMissCounter;
    TSharedCounter CheckerCacheSizeCounter;
};

class TCpuLimiter final : public ICpuMeasureCallback {
public:
    TCpuLimiter(const TCpuLimiterConfig& config, TCpuLimiterStat& stat, size_t workerId);

    class TCacheKey {
    public:
        explicit TCacheKey(const NAddr::IRemoteAddr& addr) {
            const sockaddr* const sa = addr.Addr();

            switch (sa->sa_family) {
                case AF_INET:{
                    Address_[0] = AF_INET;
                    const auto& saddr = (reinterpret_cast<const sockaddr_in*>(sa))->sin_addr.s_addr;
                    static_assert(sizeof(saddr) + 1 <= SIZE);
                    memcpy(Address_.data() + 1, &saddr, sizeof(saddr));
                    break;
                }
                case AF_INET6: {
                    Address_[0] = AF_INET6;
                    const auto& saddr = (reinterpret_cast<const sockaddr_in6*>(sa))->sin6_addr.s6_addr;
                    static_assert(sizeof(saddr) + 1 <= SIZE);
                    memcpy(Address_.data() + 1, &saddr, sizeof(saddr));
                    break;
                }
            }
        }

        TStringBuf GetAddress() const {
            return {Address_.begin(), Address_.end()};
        }

        auto operator<=>(const TCacheKey&) const = default;

        static constexpr size_t SIZE = 17;

    private:
        std::array<char, SIZE> Address_{};
    };

    class TCacheValue {};

    bool CheckConnRejected(TSocketHolder& socket) noexcept {
        if (!FileControlDisabled()) {
            if (Config_.EnableConnReject && IsDropped(Config_.ConnRejectLo, Config_.ConnRejectHi)) {
                RejectedConnsCleaner_.Push(std::move(socket));
                ++Stat_.CpuLimiterConnRejectedCounter;
                return true;
            }
        }
        return false;
    }

    bool CheckHTTP2Disabled() noexcept {
        if (!FileControlDisabled() && !HTTP2Enabled_ || FileControlHTTP2Disabled()) {
            ++Stat_.CpuLimiterHTTP2DisabledCounter;
            return true;
        }
        return false;
    }

    bool CheckHTTP2Closed() noexcept {
        if (!FileControlDisabled() && !HTTP2Enabled_ || FileControlHTTP2Disabled()) {
            ++Stat_.CpuLimiterHTTP2ClosedCounter;
            return true;
        }
        return false;
    }

    bool CheckKeepAliveClosed() noexcept {
        if (!FileControlDisabled()) {
            if (Config_.EnableKeepAliveClose && IsDropped(Config_.KeepAliveCloseLo, Config_.KeepAliveCloseHi)) {
                ++Stat_.CpuLimiterKeepAliveClosedCounter;
                return true;
            }
        }
        return false;
    }

    void PushCheckerAddress(const NAddr::IRemoteAddr& addr) noexcept {
        ++Stat_.CheckerCacheMissCounter;

        TCacheKey key(addr);
        size_t oldSize = Cache_.Size();
        Cache_.Insert({key, {}});
        if (oldSize < Cache_.Size()) {
            ++Stat_.CheckerCacheSizeCounter;
        }
    }

    bool IsCheckerAddress(const NAddr::IRemoteAddr& addr) noexcept {
        if (Cache_.Empty()) {
            return false;
        }

        TCacheKey key(addr);
        return Cache_.Find(key) != Cache_.End();
    }

    void OnMeasure(TCpuAndTimeDelta delta) noexcept {
        if (delta.MeasureTime.MilliSeconds() > 0) {
            double usage = delta.TotalUsage.MilliSeconds();
            double measure = std::max(delta.TotalUsage.MilliSeconds(), delta.MeasureTime.MilliSeconds());
            SetCpuUsage((1 - Config_.CpuUsageCoeff) * CpuUsage_ + Config_.CpuUsageCoeff * usage / measure);

            if (Config_.EnableHTTP2Drop) {
                if (HTTP2Enabled_ && CpuUsage_ > Config_.HTTP2DropHi) {
                    HTTP2Enabled_ = false;
                } else if (!HTTP2Enabled_ && CpuUsage_ < Config_.HTTP2DropLo) {
                    HTTP2Enabled_ = true;
                }
            }
        }
    }

    void OnInit(TContExecutor* executor, TSharedFiles* sharedFiles);

private:

    bool FileControlDisabled() const noexcept {
        return DisableFileChecker_.Exists();
    }

    bool FileControlHTTP2Disabled() const noexcept {
        return DisableHTTP2FileChecker_.Exists();
    }

    bool IsDropped(double lo, double hi) const noexcept {
         if (CpuUsage_ > lo) {
            if (CpuUsage_ >= hi || RandomNumber<double>() * (hi - lo) < (CpuUsage_ - lo)) {
                return true;
            }
        }
        return false;
    }

    bool IsAnythingEnabled() const noexcept {
        return Config_.EnableConnReject || Config_.EnableHTTP2Drop || Config_.EnableKeepAliveClose;
    }

    void SetCpuUsage(double cpuUsage) noexcept {
        CpuUsage_ = cpuUsage;
        Stat_.AverageCpuUsageCounter.Set(cpuUsage * 1000000);
    }

    const TCpuLimiterConfig& Config_;
    TCpuLimiterStat Stat_;

    TLRUCache<TCacheKey, TCacheValue> Cache_;

    double CpuUsage_ = 0;
    TRejectedConnsCleaner RejectedConnsCleaner_;

    TSharedFileExistsChecker DisableFileChecker_;
    TSharedFileExistsChecker DisableHTTP2FileChecker_;

    bool HTTP2Enabled_ = true;
};

class TCpuLimiterBuilder : public NConfig::IConfig::IFunc {
public:
    TCpuLimiterConfig GetConfig() {
        return std::move(Config_);
    }

    void ParseConfig(NConfig::IConfig* config) {
        config->ForEach(this);

        if (Config_.CpuUsageCoeff < 0 || Config_.CpuUsageCoeff > 1) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: cpu usage coeff should be in [0,1]";
        }

        if (Config_.ConnRejectLo < 0 || Config_.ConnRejectLo > 1) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: conn drop lower bound should be in [0,1]";
        }

        if (Config_.ConnRejectHi < 0 || Config_.ConnRejectHi > 1) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: conn drop upper bound should be in [0,1]";
        }

        if (Config_.ConnRejectLo > Config_.ConnRejectHi) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: conn drop lower bound shouldn't be above than upper bound";
        }

        if (Config_.HTTP2DropLo < 0 || Config_.HTTP2DropLo > 1) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: http2 drop lower bound should be in [0,1]";
        }

        if (Config_.HTTP2DropHi < 0 || Config_.HTTP2DropHi > 1) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: http2 drop upper bound should be in [0,1]";
        }

        if (Config_.HTTP2DropLo > Config_.HTTP2DropHi) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: http2 drop lower bound shouldn't be above than upper bound";
        }

        if (Config_.KeepAliveCloseLo < 0 || Config_.KeepAliveCloseLo > 1) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: keepalive close lower bound should be in [0,1]";
        }

        if (Config_.KeepAliveCloseHi < 0 || Config_.KeepAliveCloseHi > 1) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: keepalive close upper bound should be in [0,1]";
        }

        if (Config_.KeepAliveCloseLo > Config_.KeepAliveCloseHi) {
            ythrow NConfig::TConfigParseError{} << "cpu limiter: keepalive close lower bound shouldn't be above than upper bound";
        }
    }

private:
    START_PARSE {
        ON_KEY("cpu_usage_coeff", Config_.CpuUsageCoeff) {
            return;
        }

        ON_KEY("enable_conn_reject", Config_.EnableConnReject) {
            return;
        }

        ON_KEY("enable_http2_drop", Config_.EnableHTTP2Drop) {
            return;
        }

        ON_KEY("enable_keepalive_close", Config_.EnableKeepAliveClose) {
            return;
        }

        ON_KEY("conn_reject_lo", Config_.ConnRejectLo) {
            return;
        }

        ON_KEY("conn_reject_hi", Config_.ConnRejectHi) {
            return;
        }

        ON_KEY("conn_hold_count", Config_.ConnHoldCount) {
            return;
        }

        ON_KEY("conn_hold_duration", Config_.ConnHoldDuration) {
            return;
        }

        ON_KEY("checker_address_cache_size", Config_.CheckerAddressCacheSize) {
            return;
        }

        ON_KEY("http2_drop_lo", Config_.HTTP2DropLo) {
            return;
        }

        ON_KEY("http2_drop_hi", Config_.HTTP2DropHi) {
            return;
        }

        ON_KEY("keepalive_close_lo", Config_.KeepAliveCloseLo) {
            return;
        }

        ON_KEY("keepalive_close_hi", Config_.KeepAliveCloseHi) {
            return;
        }

        ON_KEY("disable_file", Config_.DisableFile) {
            return;
        }

        ON_KEY("disable_http2_file", Config_.DisableHTTP2File) {
            return;
        }

        if (key == "active_check_subnet") {
            return;
        }

        if (key == "active_check_subnet_default") {
            return;
        }

        if (key == "active_check_subnet_file") {
            return;
        }
    } END_PARSE

    TCpuLimiterConfig Config_;
};

}  // namespace NSrvKernel

template <>
struct THash<NSrvKernel::TCpuLimiter::TCacheKey> {
    size_t operator()(const NSrvKernel::TCpuLimiter::TCacheKey& i) const {
        return THash<TStringBuf>()(i.GetAddress());
    }
};
