#include <balancer/kernel/rpslimiter/quota.h>

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

using namespace NSrvKernel::NRpsLimiter;

Y_UNIT_TEST_SUITE(TLocalQuotaStateTest) {
    Y_UNIT_TEST(TestOneHost) {
        const auto now = TInstant::Seconds(1600000000);
        const auto q0 = TQuota{
            .Name="q0",
            .Start=now,
            .Window=TDuration::Seconds(1),
            .Limit=5,
        };
        const auto q1 = TQuota{
            .Name="q1",
            .Start=now,
            .Window=TDuration::Seconds(10),
            .Limit=15,
        };
        const auto q2 = TQuota{
            .Name="q2",
            .Start=now,
            .Window=TDuration::Seconds(30),
            .Limit=20,
        };
        const auto ms = TDuration::MilliSeconds(1);
        TQuotaState state("self", now, {"self"}, {q2, q1, q0});
        // Ordered by name
        UNIT_ASSERT_VALUES_EQUAL(state.QuotaIdx(q0.Name), 0ull);
        UNIT_ASSERT_VALUES_EQUAL(state.QuotaIdx(q1.Name), 1ull);
        UNIT_ASSERT_VALUES_EQUAL(state.QuotaIdx(q2.Name), 2ull);
        UNIT_ASSERT(!state.QuotaIdx("qX"));
        UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(0, now), 0);
        UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(1, now), 0);
        UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(2, now), 0);
        UNIT_ASSERT_VALUES_EQUAL(state.QuotaInfo(0), q0);
        UNIT_ASSERT_VALUES_EQUAL(state.QuotaInfo(1), q1);
        UNIT_ASSERT_VALUES_EQUAL(state.QuotaInfo(2), q2);

        for (auto i : xrange(1, 6)) {
            const auto t = now + 5 * i * ms;

            state.IncLocalQuota(0, t, 1);
            UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(0, t), i);
            UNIT_ASSERT_VALUES_EQUAL(state.LocalQuotaRate(0, t), i);

            state.IncLocalQuota(1, t, 2);
            UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(1, t), 2*i);
            UNIT_ASSERT_VALUES_EQUAL(state.LocalQuotaRate(1, t), 2*i);
        }

        UNIT_ASSERT_VALUES_EQUAL(state.GenPeerQuotasFromLocal(now + q0.Window / 2), (TPeerQuotas{
            .Name="self",
            .Time=now + q0.Window / 2,
            .Quotas={
                {.Name=q0.Name, .Window=q0.Window, .Rate=5},
                {.Name=q1.Name, .Window=q1.Window, .Rate=10}
            },
        }));
        UNIT_ASSERT_VALUES_EQUAL(state.GenPeerQuotasFromLocal(now + q0.Window + q0.Window / 2), (TPeerQuotas{
            .Name="self",
            .Time=now + q0.Window + q0.Window / 2,
            .Quotas={
                {.Name=q0.Name, .Window=q0.Window, .Rate=2.5},
                {.Name=q1.Name, .Window=q1.Window, .Rate=10}
            },
        }));
        UNIT_ASSERT_VALUES_EQUAL(state.GenPeerQuotasFromLocal(now + q1.Window + q1.Window / 2), (TPeerQuotas{
            .Name="self",
            .Time=now + q1.Window + q1.Window / 2,
            .Quotas={
                {.Name=q1.Name, .Window=q1.Window, .Rate=5}
            },
        }));
    }

    Y_UNIT_TEST(TestMultipleHosts) {
        const auto now = TInstant::Seconds(1600000000);
        const auto q0 = TQuota{
            .Name="q0",
            .Start=now,
            .Window=TDuration::Seconds(1),
            .Limit=5,
        };
        const auto q1 = TQuota{
            .Name="q1",
            .Start=now,
            .Window=TDuration::Seconds(10),
            .Limit=15,
        };
        const auto us = TDuration::MicroSeconds(1);
        const auto ms = TDuration::MilliSeconds(1);
        TQuotaState state("self", now, {"self", "peer0", "peer1", "peer2"}, {q0, q1});

        auto localTime = now + 250*ms;
        auto peer1Time = now + 50*ms;
        auto peer2Time = now;

        // An invalid update
        UNIT_ASSERT(!state.UpdatePeerQuotas({
            .Name="peer2",
            .Time=peer2Time,
            .Quotas={
                {.Name="qX", .Window=TDuration::Seconds(1), .Rate=100},
                {.Name="q0", .Rate=100},
            },
        }, now));

        UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(0, now), 0);

        peer2Time += us;
        // A window rescale
        UNIT_ASSERT(state.UpdatePeerQuotas({
           .Name="peer2",
           .Time=peer2Time,
           .Quotas={
                {.Name="q1", .Window=TDuration::Seconds(1), .Rate=1},
            },
        }, localTime));
        UNIT_ASSERT_VALUES_EQUAL(state.PeerQuotaRate(1, localTime), 10);
        UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(1, localTime), 10);
        UNIT_ASSERT_VALUES_EQUAL(state.PeerQuotaRate(0, localTime), 0);
        UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(0, localTime), 0);

        localTime += q1.Window / 2;
        UNIT_ASSERT(state.UpdatePeerQuotas({
           .Name="peer1",
           .Time=peer1Time,
           .Quotas={
                {.Name="q1", .Window=TDuration::Seconds(10), .Rate=5},
            },
        }, localTime));

        // A stale update, ignored
        UNIT_ASSERT(!state.UpdatePeerQuotas({
            .Name="peer1",
            .Time=peer1Time,
            .Quotas={
                {.Name="q1", .Window=TDuration::Seconds(10), .Rate=10},
            },
        }, localTime));

        state.IncLocalQuota(1, localTime, 1);
        UNIT_ASSERT_VALUES_EQUAL(state.PeerQuotaRate(1, localTime), 15);
        UNIT_ASSERT_VALUES_EQUAL(state.LocalQuotaRate(1, localTime), 1);
        UNIT_ASSERT_VALUES_EQUAL(state.TotalQuotaRate(1, localTime), 16);

        // Forgetting about stale remote hosts
        localTime += q1.Window / 2;
        UNIT_ASSERT_DOUBLES_EQUAL(state.LocalQuotaRate(1, localTime), 0.975, 0.000001);
        UNIT_ASSERT_VALUES_EQUAL(state.PeerQuotaRate(1, localTime), 5);
        UNIT_ASSERT_DOUBLES_EQUAL(state.TotalQuotaRate(1, localTime), 5.975, 0.000001);

        peer1Time += 1000 * ms;
        peer2Time += 1000 * ms;
        // New updates
        UNIT_ASSERT(state.UpdatePeerQuotas({
           .Name="peer1",
           .Time=peer1Time,
           .Quotas={
                {.Name="q1", .Window=TDuration::Seconds(10), .Rate=15},
            },
        }, localTime));
        UNIT_ASSERT(state.UpdatePeerQuotas({
           .Name="peer2",
           .Time=peer2Time,
           .Quotas={
                {.Name="q1", .Window=TDuration::Seconds(10), .Rate=10},
            },
        }, localTime));

        UNIT_ASSERT_DOUBLES_EQUAL(state.LocalQuotaRate(1, localTime), 0.975, 0.000001);
        UNIT_ASSERT_VALUES_EQUAL(state.PeerQuotaRate(1, localTime), 25);
        UNIT_ASSERT_DOUBLES_EQUAL(state.TotalQuotaRate(1, localTime), 25.975, 0.000001);
    }

    Y_UNIT_TEST(TestParseRender) {
        const TString productionData =
            "zAMKIDJhMDI6NmI4OmMwODphMjE3OjEwZToyYmQ6MTZiODowEImS3OelqO0CIh8KDnJlcXdpemFy"
            "ZC1tYWluEMCEPSIJER7rPce8poZAIhsKCnNlYXJjaC12bGEQwIQ9IgkRuZhcwG7LS0AiHQoMbDdf"
            "ZmFzdC1iZWxsEMCEPSIJEZXcO4AgwJJAIiMKEnJlcXdpemFyZC1tZWdhbWluZBDAhD0iCRFtGdvF"
            "eIlhQCIgCg9sN19mYXN0LWRlZmF1bHQQwIQ9IgkRGbVoPdfcpUAiKwoacmVxd2l6YXJkLW1hcmtl"
            "dC1saW5nYm9vc3QQwIQ9IgkRUxygDTWhNEAiIwoSY29uZmxhZ2V4cC1kZWZhdWx0EMCEPSIJEe6k"
            "EFQ1IWNAIiAKD21hcHN1Z2dlc3QtbWFpbhDAhD0iCRHDBGghtrp4QCIbCgpzZWFyY2gtc2FzEMCE"
            "PSIJEQHr2PVFvE1AIhsKCnNlYXJjaC1tYW4QwIQ9IgkRVIv49WTBSUAiLgodcmVxd2l6YXJkLW1h"
            "cmtldC1uby1saW5nYm9vc3QQwIQ9IgkR7RGvpp8fVEAiHQoMaXRkaXRwLXR1cmJvEMCEPSIJER99"
            "EplfAGBA";
        TPeerQuotas hq = ParsePeerQuotas(Base64Decode(productionData)).GetOrThrow();

        UNIT_ASSERT_VALUES_EQUAL(hq, (TPeerQuotas{
            .Name="2a02:6b8:c08:a217:10e:2bd:16b8:0",
            .Time=TInstant::MicroSeconds(1606671515715849ull),
            .Quotas={
                {.Name="reqwizard-main", .Window=TDuration::Seconds(1), .Rate=724.84217689870616},
                {.Name="search-vla", .Window=TDuration::Seconds(1), .Rate=55.589317364895983},
                {.Name="l7_fast-bell", .Window=TDuration::Seconds(1), .Rate=1200.0317391732581},
                {.Name="reqwizard-megamind", .Window=TDuration::Seconds(1), .Rate=140.29599278252445},
                {.Name="l7_fast-default", .Window=TDuration::Seconds(1), .Rate=2798.4203903885905},
                {.Name="reqwizard-market-lingboost", .Window=TDuration::Seconds(1), .Rate=20.629715777959358},
                {.Name="conflagexp-default", .Window=TDuration::Seconds(1), .Rate=153.03775981188613},
                {.Name="mapsuggest-main", .Window=TDuration::Seconds(1), .Rate=395.66946545250192},
                {.Name="search-sas", .Window=TDuration::Seconds(1), .Rate=59.470885020176873},
                {.Name="search-man", .Window=TDuration::Seconds(1), .Rate=51.510893579825876},
                {.Name="reqwizard-market-no-lingboost", .Window=TDuration::Seconds(1), .Rate=80.494119330374545},
                {.Name="itditp-turbo", .Window=TDuration::Seconds(1), .Rate=128.01166967021069},
            },
        }));

        UNIT_ASSERT_VALUES_EQUAL(ParsePeerQuotas(RenderPeerQuotas(hq)).GetOrThrow(), hq);
    }
}
