#include <balancer/kernel/custom_io/chunkio.h>
#include <balancer/kernel/http/parser/httpdecoder.h>
#include <balancer/kernel/http/parser/httpencoder.h>
#include <balancer/kernel/stats/manager.h>
#include <balancer/kernel/testing/stats.h>
#include <library/cpp/json/json_reader.h>
#include <library/cpp/testing/unittest/registar.h>
#include <util/thread/pool.h>

#ifdef _linux_
#include <sys/wait.h>
#endif

using namespace NSrvKernel;
using namespace NSrvKernel::NTesting;


Y_UNIT_TEST_SUITE(TTestSharedStatsManager) {
    Y_UNIT_TEST(TestMakeCounterSameSignal) {
        TStatsFixture fixture;
        auto& statsManager = fixture.Manager();

        Y_UNUSED(statsManager.MakeCounter("test_counter").Build());
        UNIT_ASSERT_EXCEPTION(statsManager.MakeCounter("test_counter").Build(), yexception);
    }

    Y_UNIT_TEST(TestMakeCounterOrReuseExisting) {
        TStatsFixture fixture;
        auto& statsManager = fixture.Manager();

        TSharedCounter first = statsManager.MakeCounter("test_counter").AllowDuplicate().Build();
        TSharedCounter second = statsManager.MakeCounter("test_counter").AllowDuplicate().Build();

        fixture.FreezeAndInit();

        first.Set(1337);
        UNIT_ASSERT_VALUES_EQUAL(first.Value(), second.Value());

        second.Set(228);
        UNIT_ASSERT_VALUES_EQUAL(first.Value(), second.Value());
    }

    Y_UNIT_TEST(TestMakeHistogramTooManyIntervals) {
        TStatsFixture fixture;
        auto& statsManager = fixture.Manager();

        TVector<ui64> intervals;
        intervals.resize(100, 1);
        UNIT_ASSERT_EXCEPTION(statsManager.MakeHistogram("test", std::move(intervals)).Build(), yexception);
    }

    Y_UNIT_TEST(TestCountersInUnistatAnswer) {
        TStatsFixture fixture;
        auto& statsManager = fixture.Manager();

        // Also testing the combined name logic
        TSharedCounter first = statsManager.MakeCounter({"", "counter", "", "1", ""}).AllowDuplicate().Build();
        TSharedCounter second = statsManager.MakeCounter({"counter", "2"}).AllowDuplicate().Build();

        fixture.FreezeAndInit();

        first.Set(1337);
        second.Set(228);

        {
            auto stats = fixture.Unistat();
            UNIT_ASSERT_VALUES_EQUAL(stats["counter-1_summ"], 1337);
            UNIT_ASSERT_VALUES_EQUAL(stats["counter-2_summ"], 228);
        }

        first.Add(11);
        second.Sub(10);

        {
            auto stats = fixture.Unistat();
            UNIT_ASSERT_VALUES_EQUAL(stats["counter-1_summ"], 1337 + 11);
            UNIT_ASSERT_VALUES_EQUAL(stats["counter-2_summ"], 228 - 10);
        }

        first.Inc();
        second.Dec();

        {
            auto stats = fixture.Unistat();
            UNIT_ASSERT_VALUES_EQUAL(stats["counter-1_summ"], 1337 + 11 + 1);
            UNIT_ASSERT_VALUES_EQUAL(stats["counter-2_summ"], 228 - 10 - 1);
        }
    }

    Y_UNIT_TEST(TestHistogramsInUnistatAnswer) {
        TStatsFixture fixture;
        auto& statsManager = fixture.Manager();

        TVector<ui64> intervals = { 10, 25, 100 };

        auto first = statsManager.MakeHistogram("first", intervals).Build();
        auto second = statsManager.MakeHistogram("second", intervals).Scale(100).Build();

        fixture.FreezeAndInit();

        {
            auto stats = fixture.Unistat();
            const auto& hist1 = stats["first_hgram"].GetArray();
            const auto& hist2 = stats["second_hgram"].GetArray();
            for (size_t i = 0; i < intervals.size(); i++) {
                UNIT_ASSERT_VALUES_EQUAL(hist1[i][0], intervals[i]);
                UNIT_ASSERT_VALUES_EQUAL(hist2[i][0], 1.0 * intervals[i] / 100);
            }
        }

        first.AddValue(1); // ignored
        first.AddValue(25); // bucket #1
        first.AddValue(26); // bucket #1
        first.AddValue(99); // bucket #1
        first.AddValue(101); // bucket #2

        {
            auto stats = fixture.Unistat();
            const auto& hist = stats["first_hgram"].GetArray();
            UNIT_ASSERT_VALUES_EQUAL(hist[0][1], 0);
            UNIT_ASSERT_VALUES_EQUAL(hist[1][1], 3);
            UNIT_ASSERT_VALUES_EQUAL(hist[2][1], 1);
        }

        second.AddValue(10); // bucket #0
        second.AddValue(25); // bucket #1
        second.AddValue(100); // bucket #2

        {
            auto stats = fixture.Unistat();
            const auto& hist = stats["second_hgram"].GetArray();
            UNIT_ASSERT_VALUES_EQUAL(hist[0][1], 1);
            UNIT_ASSERT_VALUES_EQUAL(hist[1][1], 1);
            UNIT_ASSERT_VALUES_EQUAL(hist[2][1], 1);
        }
    }

    Y_UNIT_TEST(TestPostProcessing) {
        TStatsFixture fixture;
        auto& statsManager = fixture.Manager();

        TSharedCounter c(statsManager.MakeCounter("counter").PostProcess([](ui64 x) { return 2 * x; }).Build());

        fixture.FreezeAndInit();

        c.Inc();
        c.Inc();

        {
            auto stats = fixture.Unistat();
            UNIT_ASSERT_VALUES_EQUAL(stats["counter_summ"], 4);
        }
    }

#ifdef _linux_
    constexpr size_t WorkersCount = 5;

    template <typename Runner>
    static void RunProcesses(Runner&& runner) {
        TVector<pid_t> pids;
        for (size_t i = 0; i <= WorkersCount; i++) {
            pid_t pid = fork();
            if (pid == -1) {
                UNIT_ASSERT_C(false, "fork failed");
                for (pid_t child : pids) {
                    kill(child, SIGKILL);
                    waitpid(child, nullptr, 0);
                }
            } else if (pid == 0) {
                runner(i);
                _exit(0);
            } else {
                pids.push_back(pid);
            }
        }

        for (pid_t pid : pids) {
            int status = 0;
            UNIT_ASSERT(waitpid(pid, &status, 0) != -1);
            UNIT_ASSERT(!WIFSIGNALED(status) && WIFEXITED(status) && WEXITSTATUS(status) == 0);
        }
    }

    Y_UNIT_TEST(TestCountersInMultipleProcesses) {
        TStatsFixture fixture({.Workers = WorkersCount});
        auto& statsManager = fixture.Manager();

        TSharedCounter first = statsManager.MakeCounter("counter1").AllowDuplicate().Build();
        TSharedCounter second = statsManager.MakeCounter("counter2").AllowDuplicate().Build();

        fixture.Allocator().Freeze();

        RunProcesses([&] (size_t workerId) {
            if (workerId == 0) {
                return;
            }

            for (size_t i = 0; i < 10; i++) {
                first.Inc();
            }

            for (size_t i = 0; i < 10; i++) {
                second.Dec();
            }
        });


        {
            auto stats = fixture.Unistat();
            UNIT_ASSERT_VALUES_EQUAL(stats["counter1_summ"], WorkersCount * 10);
            UNIT_ASSERT_VALUES_EQUAL(stats["counter2_summ"], -WorkersCount * 10);
        }
    }

    Y_UNIT_TEST(TestHistogramsInMultipleProcesses) {
        TStatsFixture fixture({.Workers = WorkersCount});
        auto& statsManager = fixture.Manager();

        TVector<ui64> intervals = { 10, 25, 100 };
        auto first = statsManager.MakeHistogram("first", intervals).Build();
        auto second = statsManager.MakeHistogram("second", intervals).Build();

        fixture.Allocator().Freeze();

        RunProcesses([&] (size_t workerId) {
            if (workerId == 0) {
                return;
            }

            for (size_t i = 0; i < 100; i++) {
                first.AddValue(i);
            }

            for (size_t i = 0; i < 100; i++) {
                second.AddValue(10 + i);
            }
        });

        {
            auto stats = fixture.Unistat();
            const auto& hist1 = stats["first_hgram"].GetArray();
            UNIT_ASSERT_VALUES_EQUAL(hist1[0][1], WorkersCount * 15);
            UNIT_ASSERT_VALUES_EQUAL(hist1[1][1], WorkersCount * 75);
            UNIT_ASSERT_VALUES_EQUAL(hist1[2][1], 0);

            const auto& hist2 = stats["second_hgram"].GetArray();
            UNIT_ASSERT_VALUES_EQUAL(hist2[0][1], WorkersCount * 15);
            UNIT_ASSERT_VALUES_EQUAL(hist2[1][1], WorkersCount * 75);
            UNIT_ASSERT_VALUES_EQUAL(hist2[2][1], WorkersCount * 10);
        }

    }

    Y_UNIT_TEST(TestWrongHistogramIntervals) {
        TStatsFixture fixture({.Workers = WorkersCount});
        auto& statsManager = fixture.Manager();

        UNIT_ASSERT_EXCEPTION(statsManager.MakeHistogram("test", {1, 3, 2}).Build(), yexception);
        UNIT_ASSERT_EXCEPTION(statsManager.MakeHistogram("test", {1, 2, 2}).Build(), yexception);
    }

#endif

    Y_UNIT_TEST(TestHistogramAggregationInMultipleThreads) {
        constexpr size_t ThreadsNumber = 4;

        TStatsFixture fixture({.Workers = ThreadsNumber});
        auto& statsManager = fixture.Manager();

        TVector<ui64> intervals = { 10, 25, 100 };
        TSharedHistogram histogram = statsManager.MakeHistogram("first", intervals).Build();

        fixture.Allocator().Freeze();

        TThreadPool pool;
        pool.Start(ThreadsNumber);

        for (size_t workerId = 1; workerId <= ThreadsNumber; workerId++) {
            UNIT_ASSERT(pool.AddFunc([workerId, &histogram, &intervals]() {
                TSharedHistogram workerHistogram(histogram, workerId);
                for (size_t i = 0; i < intervals.back(); i++)
                    workerHistogram.AddValue(i);
            }));
        }

        pool.Stop();

        auto stats = fixture.Unistat();
        const auto& hist = stats["first_hgram"].GetArray();
        UNIT_ASSERT_VALUES_EQUAL(hist[0][1], ThreadsNumber * 15);
        UNIT_ASSERT_VALUES_EQUAL(hist[1][1], ThreadsNumber * 75);
        UNIT_ASSERT_VALUES_EQUAL(hist[2][1], 0);
    }

    Y_UNIT_TEST(TestGaugeAggregationTypesInMultipleThreads) {
        constexpr size_t ThreadsNumber = 4;

        TStatsFixture fixture({.Workers = ThreadsNumber});
        auto& statsManager = fixture.Manager();

        TSharedCounter minCounter = statsManager.MakeGauge("min").Aggregation(TWorkerAggregation::Min).Default(1000).Build();
        TSharedCounter maxCounter = statsManager.MakeGauge("max").Aggregation(TWorkerAggregation::Max).Default(0).Build();

        fixture.Allocator().Freeze();

        TThreadPool pool;
        pool.Start(ThreadsNumber);

        for (size_t workerId = 1; workerId <= ThreadsNumber; workerId++) {
            UNIT_ASSERT(pool.AddFunc([workerId, &minCounter, &maxCounter]() {
                TSharedCounter workerMinCounter(minCounter, workerId);
                workerMinCounter.Set(workerId);
                TSharedCounter workerMaxCounter(maxCounter, workerId);
                workerMaxCounter.Set(workerId);
            }));
        }

        pool.Stop();

        {
            auto stats = fixture.Unistat();
            UNIT_ASSERT_VALUES_EQUAL(stats["min_annn"], 1);
            UNIT_ASSERT_VALUES_EQUAL(stats["max_axxx"], 4);
        }
    }
}
