#include <balancer/kernel/coro/coroutine.h>
#include <balancer/kernel/cpu/cpu_usage.h>
#include <balancer/kernel/cpu/cpu_limiter.h>
#include <balancer/kernel/custom_io/chunkio.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/http/parser/http.h>

#include <library/cpp/testing/unittest/registar.h>

#include <util/generic/xrange.h>

Y_UNIT_TEST_SUITE(TCpuLimiterTest) {
    Y_UNIT_TEST(TestCpuMeasurerIdle) {
        using namespace NSrvKernel;
        TContExecutor exec(32000);
        auto loadGen = TCoroutine(ECoroType::Service, "loadgen", &exec, [&exec]() {
            auto before = TInstant::Now();
            while (!exec.Running()->Cancelled() && TInstant::Now() < before + TDuration::Seconds(3)) {
                  exec.Running()->SleepT(TDuration::Seconds(1));
            }
            exec.Abort();
        });

        class TCpuMeasureCallback final : public ICpuMeasureCallback {
            void OnMeasure(TCpuAndTimeDelta delta) noexcept {
                UNIT_ASSERT(delta.MeasureTime.MilliSeconds() > 500);
                UNIT_ASSERT(delta.MeasureTime.MilliSeconds() < 1500);
                UNIT_ASSERT(delta.TotalUsage < (delta.MeasureTime / 4));
            }
        };

        TCpuMeasureCallback callback;

        auto cpuMeasurer = StartCpuAndTimeMeasurer(exec, &callback, nullptr);

        exec.Execute();
    }

    Y_UNIT_TEST(TestCpuMeasurerFull) {
        using namespace NSrvKernel;
        TContExecutor exec(32000);
        auto loadGen = TCoroutine(ECoroType::Service, "loadgen", &exec, [&exec]() {
            auto before = TInstant::Now();
            while (!exec.Running()->Cancelled() && TInstant::Now() < before + TDuration::Seconds(3)) {
                exec.Running()->Yield();
            }
            exec.Abort();
        });

        class TCpuMeasureCallback final : public ICpuMeasureCallback {
            void OnMeasure(TCpuAndTimeDelta delta) noexcept {
                UNIT_ASSERT(delta.MeasureTime.MilliSeconds() > 500);
                UNIT_ASSERT(delta.MeasureTime.MilliSeconds() < 1500);
                UNIT_ASSERT(delta.TotalUsage > (delta.MeasureTime / 2));
                UNIT_ASSERT(delta.TotalUsage < delta.MeasureTime * 1.1);
            }
        };

        TCpuMeasureCallback callback;

        auto cpuMeasurer = StartCpuAndTimeMeasurer(exec, &callback, nullptr);

        exec.Execute();
    }

    struct TLimiterEnv {
        NSrvKernel::TSharedAllocator Alloc;
        NSrvKernel::TSharedStatsManager Manager{Alloc};
        TContExecutor Exec{32000};
        NSrvKernel::TCpuLimiterConfig Config;
        TMaybe<NSrvKernel::TCpuLimiterStat> Stat;
        TMaybe<NSrvKernel::TCpuLimiter> Limiter;

        TLimiterEnv(TString cfg)
        {
            Manager.SetWorkersCount(1, 1);
            TStringStream rawCfg(cfg);
            NSrvKernel::TCpuLimiterBuilder builder;
            builder.ParseConfig(NConfig::ConfigParser(rawCfg).Get());
            Config = builder.GetConfig();
            Stat.ConstructInPlace(Manager);
            Limiter.ConstructInPlace(Config, *Stat, 0);
            Alloc.Freeze();
            Limiter->OnInit(&Exec, nullptr);
        }

        ~TLimiterEnv() {
            Exec.Abort();
            Exec.Execute();
        }
    };

    Y_UNIT_TEST(TestDefaultSettings) {
        using namespace NSrvKernel;

        TLimiterEnv env("");

        for (auto i : xrange(100)) {
            Y_UNUSED(i);
            UNIT_ASSERT(!env.Limiter->CheckHTTP2Closed());
            UNIT_ASSERT(!env.Limiter->CheckHTTP2Disabled());
            UNIT_ASSERT(!env.Limiter->CheckKeepAliveClosed());
            env.Limiter->OnMeasure({TDuration::Seconds(1), TDuration::Seconds(1)});
        }
        UNIT_ASSERT(!env.Limiter->CheckHTTP2Closed());
        UNIT_ASSERT(!env.Limiter->CheckHTTP2Disabled());
        UNIT_ASSERT(!env.Limiter->CheckKeepAliveClosed());
    }

    void DoCheckStats(TLimiterEnv& env, TString expected) {
        using namespace NSrvKernel;
        TChunksOutputStream out;
        env.Manager.WriteResponseNoWait(TSharedStatsManager::TResponseType::Unistat, &out);
        UNIT_ASSERT_VALUES_EQUAL(StrInplace(out.Chunks()), expected);
    }

    void DoCheckHttp2Drop(TLimiterEnv& env, ui32 toDisable, ui32 toReeable) {
        for (auto i : xrange(toDisable)) {
            Y_UNUSED(i);
            UNIT_ASSERT(!env.Limiter->CheckHTTP2Closed());
            UNIT_ASSERT(!env.Limiter->CheckHTTP2Disabled());
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(1),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        for (auto i : xrange(toReeable)) {
            Y_UNUSED(i);
            UNIT_ASSERT(env.Limiter->CheckHTTP2Closed());
            UNIT_ASSERT(env.Limiter->CheckHTTP2Disabled());
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(0),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        UNIT_ASSERT(!env.Limiter->CheckHTTP2Closed());
        UNIT_ASSERT(!env.Limiter->CheckHTTP2Disabled());
    }

    Y_UNIT_TEST(TestHttp2Drop) {
        using namespace NSrvKernel;

        TLimiterEnv env(R"(
            instance = {
                enable_http2_drop = 1;
                http2_drop_lo = 0.3;
                http2_drop_hi = 0.5;
            };
        )");

        DoCheckHttp2Drop(env, 14, 11);
        DoCheckStats(env, R"([["worker-cpu_limiter_conn_rejected_summ",0],)"
                          R"(["worker-cpu_limiter_http2_closed_summ",11],)"
                          R"(["worker-cpu_limiter_conn_hold_summ",0],)"
                          R"(["worker-average_cpu_usage_ammv",291410],)"
                          R"(["worker-cpu_limiter_checker_cache_size_ammv",0],)"
                          R"(["worker-cpu_limiter_checker_cache_miss_summ",0],)"
                          R"(["worker-cpu_limiter_keepalive_closed_summ",0],)"
                          R"(["worker-cpu_limiter_http2_disabled_summ",11]])");
    }

    Y_UNIT_TEST(TestHttp2DropInvalid) {
        using namespace NSrvKernel;

        TLimiterEnv env(R"(
            instance = {
                http2_drop_lo = 0.5;
                http2_drop_hi = 0.75;
                enable_http2_drop = 1;
            };
        )");

        env.Limiter->OnMeasure({
            .TotalUsage = {},
            .MeasureTime = {},
        });
        env.Limiter->OnMeasure({
            .TotalUsage = TDuration::Seconds(1),
            .MeasureTime = {},
        });
        env.Limiter->OnMeasure({
            .TotalUsage = TDuration::Max(),
            .MeasureTime = TDuration::Seconds(1),
        });
        DoCheckHttp2Drop(env, 27, 9);
        DoCheckStats(env, R"([["worker-cpu_limiter_conn_rejected_summ",0],)"
                          R"(["worker-cpu_limiter_http2_closed_summ",9],)"
                          R"(["worker-cpu_limiter_conn_hold_summ",0],)"
                          R"(["worker-average_cpu_usage_ammv",480359],)"
                          R"(["worker-cpu_limiter_checker_cache_size_ammv",0],)"
                          R"(["worker-cpu_limiter_checker_cache_miss_summ",0],)"
                          R"(["worker-cpu_limiter_keepalive_closed_summ",0],)"
                          R"(["worker-cpu_limiter_http2_disabled_summ",9]])");
    }


    void DoCheckKeepaliveClose(TLimiterEnv& env, ui32 beginClose, ui32 sureClose, ui32 reOpen) {
        for (auto i : xrange(beginClose)) {
            Y_UNUSED(i);
            UNIT_ASSERT(!env.Limiter->CheckKeepAliveClosed());
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(1),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        for (auto i : xrange(sureClose)) {
            Y_UNUSED(i);
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(1),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        UNIT_ASSERT(env.Limiter->CheckKeepAliveClosed());
        for (auto i : xrange(reOpen)) {
            Y_UNUSED(i);
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(0),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        UNIT_ASSERT(!env.Limiter->CheckKeepAliveClosed());
    }

    Y_UNIT_TEST(TestKeepaliveClose) {
        using namespace NSrvKernel;

        TLimiterEnv env(R"(
            instance = {
                keepalive_close_lo = 0.3;
                keepalive_close_hi = 0.9;
                enable_keepalive_close = 1;
            };
        )");

        DoCheckKeepaliveClose(env, 7, 38, 22);
        DoCheckStats(env, R"([["worker-cpu_limiter_conn_rejected_summ",0],)"
                          R"(["worker-cpu_limiter_http2_closed_summ",0],)"
                          R"(["worker-cpu_limiter_conn_hold_summ",0],)"
                          R"(["worker-average_cpu_usage_ammv",291361],)"
                          R"(["worker-cpu_limiter_checker_cache_size_ammv",0],)"
                          R"(["worker-cpu_limiter_checker_cache_miss_summ",0],)"
                          R"(["worker-cpu_limiter_keepalive_closed_summ",1],)"
                          R"(["worker-cpu_limiter_http2_disabled_summ",0]])");
    }

    void DoCheckConnReject(TLimiterEnv& env, ui32 beginReject, ui32 sureReject, ui32 reAccept) {
        for (auto i : xrange(beginReject)) {
            Y_UNUSED(i);
            TSocketHolder sock;
            UNIT_ASSERT(!env.Limiter->CheckConnRejected(sock));
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(1),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        for (auto i : xrange(sureReject)) {
            Y_UNUSED(i);
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(1),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        {
            TSocketHolder sock;
            UNIT_ASSERT(env.Limiter->CheckConnRejected(sock));
        }
        for (auto i : xrange(reAccept)) {
            Y_UNUSED(i);
            env.Limiter->OnMeasure({
                .TotalUsage = TDuration::Seconds(0),
                .MeasureTime = TDuration::Seconds(1),
            });
        }
        {
            TSocketHolder sock;
            UNIT_ASSERT(!env.Limiter->CheckConnRejected(sock));
        }
    }

    Y_UNIT_TEST(TestConnReject) {
        using namespace NSrvKernel;

        TLimiterEnv env(R"(
            instance = {
                conn_reject_lo = 0.3;
                conn_reject_hi = 0.9;
                enable_conn_reject = 1;
            };
        )");

        DoCheckConnReject(env, 7, 38, 22);
        DoCheckStats(env, R"([["worker-cpu_limiter_conn_rejected_summ",1],)"
                          R"(["worker-cpu_limiter_http2_closed_summ",0],)"
                          R"(["worker-cpu_limiter_conn_hold_summ",1],)"
                          R"(["worker-average_cpu_usage_ammv",291361],)"
                          R"(["worker-cpu_limiter_checker_cache_size_ammv",0],)"
                          R"(["worker-cpu_limiter_checker_cache_miss_summ",0],)"
                          R"(["worker-cpu_limiter_keepalive_closed_summ",0],)"
                          R"(["worker-cpu_limiter_http2_disabled_summ",0]])");
    }
}
