#include <balancer/kernel/memory/shared.h>
#include <library/cpp/testing/unittest/registar.h>
#include <util/generic/utility.h>
#include <util/memory/pool.h>
#include <util/system/thread.h>

#ifdef __unix__
#   include <unistd.h>
#   include <sys/wait.h>
#endif

namespace {
    using NSrvKernel::TSharedAllocator;

    class TTestObject {
    public:
        TTestObject()
            : TestDouble(0.314)
            , TestUi64(1337)
        {
        }

        double TestDouble = 0;
        ui64 TestUi64 = 0;
    };

    class alignas(4096) TCustomAlignasObject : public TTestObject {
    };

    static size_t PageSize = NSystemInfo::GetPageSize();
}

Y_UNIT_TEST_SUITE(SharedAllocatorTest) {
    Y_UNIT_TEST(TestSimple) {
        TSharedAllocator allocator;
        TTestObject* obj = allocator.Allocate<TTestObject>();
        UNIT_ASSERT_EQUAL(obj->TestDouble, 0.314);
        UNIT_ASSERT_EQUAL(obj->TestUi64, 1337);
    }

    Y_UNIT_TEST(TestAllocateMultipleObjects) {
        TSharedAllocator allocator;
        TVector<TTestObject*> objects;
        for (size_t i = 0; i < 4096; i++) {
            TTestObject* obj = allocator.Allocate<TTestObject>();
            obj->TestDouble = 0.99 * i;
            obj->TestUi64 = 1337 + i * 228;
            objects.push_back(obj);
        }
        for (size_t i = 0; i < 4096; i++) {
            const TTestObject* obj = objects[i];
            UNIT_ASSERT_EQUAL(obj->TestDouble, 0.99 * i);
            UNIT_ASSERT_EQUAL(obj->TestUi64, 1337 + i * 228);
        }
    }

    Y_UNIT_TEST(TestAlignmentCompliance) {
        TSharedAllocator allocator;
        for (size_t i = 0; i < 1000; i++) {
            double* d = allocator.Allocate<double>(0.314);
            UNIT_ASSERT_EQUAL(*d, 0.314);
            UNIT_ASSERT(AlignDown(d, alignof(double)) == d);
            UNIT_ASSERT(AlignUp(d, alignof(double)) == d);
        }

        for (size_t i = 0; i < 1000; i++) {
            ui64* n = allocator.Allocate<ui64>(1337);
            UNIT_ASSERT_EQUAL(*n, 1337);
            UNIT_ASSERT(AlignDown(n, alignof(ui64)) == n);
            UNIT_ASSERT(AlignUp(n, alignof(ui64)) == n);
        }

        for (size_t i = 0; i < 1000; i++) {
            TTestObject* obj = allocator.Allocate<TTestObject>();
            UNIT_ASSERT_EQUAL(obj->TestDouble, 0.314);
            UNIT_ASSERT_EQUAL(obj->TestUi64, 1337);
            UNIT_ASSERT(AlignDown(obj, alignof(TTestObject)) == obj);
            UNIT_ASSERT(AlignUp(obj, alignof(TTestObject)) == obj);
        }

        for (size_t i = 0; i < 1000; i++) {
            TCustomAlignasObject* obj = allocator.Allocate<TCustomAlignasObject>();
            UNIT_ASSERT(AlignDown(obj, 4096) == obj);
            UNIT_ASSERT(AlignUp(obj, 4096) == obj);
        }
    }

    Y_UNIT_TEST(TestFrozenState) {
        TSharedAllocator allocator;
        allocator.Freeze();
        UNIT_CHECK_GENERATED_EXCEPTION(allocator.Allocate<double>(1), yexception);
    }
}
