#include <balancer/modules/hasher/lib/calc_hash.h>
#include <balancer/kernel/module/conn_descr.h>
#include <balancer/kernel/testing/conn_descr.h>
#include <balancer/kernel/testing/process_mock.h>
#include <balancer/kernel/helpers/default_instance.h>
#include <balancer/kernel/memory/chunks.h>

#include <quality/ab_testing/spu/metrics/special.h>

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

#include <util/generic/ptr.h>
#include <util/generic/singleton.h>
#include <util/random/easy.h>
#include <util/string/cast.h>

#include <cmath>

using namespace NSrvKernel;
using namespace NSrvKernel::NTesting;

const TString takeIpFrom = TString("X-Forwarded-For");

class TIpFsm : public TFsm, public TWithDefaultInstance<TIpFsm> {
public:
    TIpFsm()
        : TFsm(takeIpFrom, TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

class TIpProducer : public TWithDefaultInstance<TIpProducer> {
public:
    TIpProducer() = default;
    const TVector<TString>& Produce() {
        if (Ips_.empty()) {
            Generate();
        }
        return Ips_;
    }

    const TVector<TString>& ReProduce() {
        Generate();
        return Ips_;
    }

private:
    void Generate() {
        Ips_.resize(n);
        for (int i = 0; i < n; ++i) {
            bool ipv4 = Random();
            if (ipv4) {
                TVector<TString> ipPart;
                std::generate_n(std::back_inserter(ipPart), 4, []() {
                    unsigned num = Random();
                    return ToString(num % 256);
                });

                Ips_[i] = JoinStrings(ipPart, ".");
            } else {
                TVector<TString> ipPart;
                std::generate_n(std::back_inserter(ipPart), 8, []() {
                    unsigned num = Random();
                    return to_lower(IntToString<16>(num % (1 << 16)));
                });

                Ips_[i] = JoinStrings(ipPart, ":");
            }
        }
    }

private:
    TVector<TString> Ips_;
    int n = 2000;
};

Y_UNIT_TEST_SUITE(THasherUnitTest) {
    constexpr unsigned V4MASK = 16;
    constexpr unsigned V6MASK = 64;

    Y_UNIT_TEST(TestSubnetHeaderSelectFirstHeader) {
        TProcessMock process;
        TTestConnDescr descr1(process);
        descr1.ConnDescr().Request->Headers().Add("X-Forwarded-For-Y", "127.127.0.1");

        const auto& ips = Singleton<TIpProducer>()->Produce();

        for (size_t i = 0; i < 10; ++i) {
            descr1.ConnDescr().Request->Headers().Add(takeIpFrom, ips[i]);
        }
        ui64 hash1;
        auto err1 = CalcHashSubnet(descr1.ConnDescr(), &TIpFsm::Instance(), V4MASK, V6MASK).AssignTo(hash1);

        TTestConnDescr descr2(process);
        descr2.ConnDescr().Request->Headers().Add(takeIpFrom, ips.front());
        ui64 hash2;
        auto err2 = CalcHashSubnet(descr2.ConnDescr(), &TIpFsm::Instance(), V4MASK, V6MASK).AssignTo(hash2);

        UNIT_ASSERT(!err1);
        UNIT_ASSERT(!err2);
        UNIT_ASSERT_EQUAL(hash1, hash2);
    }

    double MakePearsonTest(const TVector<TString>& ips) {
        constexpr int MOD = 31;
        constexpr int BUCKETCOUNT = 25;
        TVector<int> ipCounter(MOD, 0);
        TProcessMock process;

        for (const auto& ip : ips) {
            TTestConnDescr descr(process);
            descr.ConnDescr().Request->Headers().Add(takeIpFrom, ip);

            ui64 hash;
            auto err = CalcHashSubnet(descr.ConnDescr(), &TIpFsm::Instance(), V4MASK, V6MASK).AssignTo(hash);
            UNIT_ASSERT(!err);

            ++ipCounter[hash % MOD];
        }

        int mean = ips.size() / MOD - 1;
        double var = 0;

        Sort(ipCounter);
        auto l = ipCounter.front();
        auto r = ipCounter.back();

        for (auto value : ipCounter) {
            var += (mean - value) * (mean - value);
        }

        var /= ipCounter.size() - 1;
        var = std::sqrt(var);

        double freq = 0.0;
        double bucketSize = (r - l) / BUCKETCOUNT;
        double prev = std::erf((l * 1.0 - mean) / (var * std::sqrt(2))) / std::sqrt(8);

        for (size_t i = 1, j = 0; i <= BUCKETCOUNT; ++i) {
            double cur = std::erf(((l + i * bucketSize)  * 1.0 - mean) / (var * std::sqrt(2))) / std::sqrt(8);
            double curFreq = (cur - prev) * 100;
            prev = cur;

            int tmp = 0;
            while (j < ipCounter.size() && ipCounter[j] <= l + i * bucketSize) {
                ++tmp;
                ++j;
            }
            freq += curFreq ? (curFreq - tmp) * (curFreq - tmp) / curFreq : 0;
        }
        return ChiSquaredDistribution(freq, BUCKETCOUNT - 1);
    }

    Y_UNIT_TEST(TestSubnetHeaderNormalDistributionBigMod) {
        const auto& ips = Singleton<TIpProducer>();
        auto pvalue = (MakePearsonTest(ips->Produce()) + MakePearsonTest(ips->ReProduce()) + MakePearsonTest(ips->ReProduce())) / 3;
        UNIT_ASSERT(pvalue > 0.05);
    }
}
