#include <balancer/kernel/fs/shared_files.h>

#include <balancer/kernel/fs/shared_file_rereader.h>

#include <balancer/kernel/testing/file_arbiter_mock.h>

#include <balancer/kernel/testing/simple_barrier.h>

#include <library/cpp/threading/task_scheduler/task_scheduler.h>

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

#include <thread>

using namespace NSrvKernel;
using namespace NSrvKernel::NTesting;

using ::testing::Return;
using ::testing::Throw;
using ::testing::InvokeWithoutArgs;
using ::testing::StrEq;
using ::testing::_;

Y_UNIT_TEST_SUITE(TBasicFileEntitiesTest) {
    constexpr char STRING_SAMPLE[] = "abc";
    constexpr char STRING_SAMPLE2[] = "def";
    constexpr char PATH_SAMPLE[] = "test.txt";

    Y_UNIT_TEST(FileDataTest) {
        // Init invariants
        // --------------------------------------------------------------------------------
        TSharedFileReReader::TData fileData;
        UNIT_ASSERT_EQUAL(fileData.Id(), 0);
        UNIT_ASSERT(fileData.Data().empty());
        UNIT_ASSERT(!fileData.Exists());
        // --------------------------------------------------------------------------------


        // Non-empty update
        // --------------------------------------------------------------------------------
        fileData.Update(STRING_SAMPLE);
        UNIT_ASSERT_EQUAL(fileData.Id(), 1);
        UNIT_ASSERT_EQUAL(fileData.Data(), STRING_SAMPLE);
        UNIT_ASSERT(fileData.Exists());
        // --------------------------------------------------------------------------------


        // Deleting update
        // --------------------------------------------------------------------------------
        fileData.Update({});
        UNIT_ASSERT_EQUAL(fileData.Id(), 2);
        UNIT_ASSERT(fileData.Data().empty());
        UNIT_ASSERT(!fileData.Exists());
        // --------------------------------------------------------------------------------


        // Empty update
        // --------------------------------------------------------------------------------
        fileData.Update("");
        UNIT_ASSERT_EQUAL(fileData.Id(), 3);
        UNIT_ASSERT(fileData.Data().empty());
        UNIT_ASSERT(fileData.Exists());
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(TFileUpdaterTest) {
        TFileManager::Instance().EnableTesting();

        auto fileArbiterMockPtr = MakeHolder<NSrvKernel::NTesting::TFileArbiterMock>();
        auto& fileArbiterMock = *fileArbiterMockPtr;

        TFileManager::TMockGuard guard{std::move(fileArbiterMockPtr)};

        Internal::TFileUpdater fup(PATH_SAMPLE);
        TFileStat retFileStat;

        // Init invariants
        // --------------------------------------------------------------------------------
        UNIT_ASSERT(!fup.Data().Exists());
        // --------------------------------------------------------------------------------


        // Situation when we read file first time
        // --------------------------------------------------------------------------------
        retFileStat.MTime = 1;
        retFileStat.CTime = 1;
        retFileStat.INode = 1;
        retFileStat.Size = 1;
        retFileStat.NLinks = 1;

        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Return(retFileStat));

        TString input = STRING_SAMPLE;
        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(1)
            .WillOnce(InvokeWithoutArgs([&] {
                return MakeHolder<TStringInput>(input);
            }));

        fup.Update();
        UNIT_ASSERT_EQUAL(fup.Data().Data(), STRING_SAMPLE);
        UNIT_ASSERT(fup.Data().Exists());
        // --------------------------------------------------------------------------------


        // Situation when file was deleted
        // --------------------------------------------------------------------------------
        retFileStat.Size = 0;
        retFileStat.NLinks = 0;

        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Return(retFileStat));

        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(0);

        fup.Update();
        UNIT_ASSERT(!fup.Data().Exists());
        // --------------------------------------------------------------------------------


        // Situation when file appeared and its last modification time was increased
        // --------------------------------------------------------------------------------
        retFileStat.MTime = 2;
        retFileStat.Size = 1;
        retFileStat.NLinks = 1;

        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Return(retFileStat));

        input = STRING_SAMPLE2;
        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(1)
            .WillOnce(InvokeWithoutArgs([&] {
                return MakeHolder<TStringInput>(input);
            }));

        fup.Update();
        UNIT_ASSERT_EQUAL(fup.Data().Data(), STRING_SAMPLE2);
        UNIT_ASSERT(fup.Data().Exists());
        // --------------------------------------------------------------------------------


        // Situation when file appeared and its last status change time was increased
        // --------------------------------------------------------------------------------
        retFileStat.CTime = 2;

        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Return(retFileStat));

        input = STRING_SAMPLE;
        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(1)
            .WillOnce(InvokeWithoutArgs([&] {
                return MakeHolder<TStringInput>(input);
            }));

        fup.Update();
        UNIT_ASSERT_EQUAL(fup.Data().Data(), STRING_SAMPLE);
        UNIT_ASSERT(fup.Data().Exists());
        // --------------------------------------------------------------------------------


        // Situation when file appeared and its inode count was changed
        // --------------------------------------------------------------------------------
        retFileStat.INode = 2;

        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Return(retFileStat));

        input = STRING_SAMPLE;
        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(1)
            .WillOnce(InvokeWithoutArgs([&] {
                return MakeHolder<TStringInput>(input);
            }));

        fup.Update();
        UNIT_ASSERT_EQUAL(fup.Data().Data(), STRING_SAMPLE);
        UNIT_ASSERT(fup.Data().Exists());
        // --------------------------------------------------------------------------------


        // Situation when file appeared and its size was changed
        // --------------------------------------------------------------------------------
        retFileStat.Size = 2;

        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Return(retFileStat));

        input = STRING_SAMPLE2;
        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(1)
            .WillOnce(InvokeWithoutArgs([&] {
                return MakeHolder<TStringInput>(input);
            }));

        fup.Update();
        UNIT_ASSERT_EQUAL(fup.Data().Data(), STRING_SAMPLE2);
        UNIT_ASSERT(fup.Data().Exists());
        // --------------------------------------------------------------------------------


        // Situation when file wasn't changed
        // --------------------------------------------------------------------------------
        auto lastData = fup.Data();

        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Return(retFileStat));

        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(0);

        fup.Update();
        UNIT_ASSERT_EQUAL(fup.Data().Id(), lastData.Id());
        UNIT_ASSERT_EQUAL(fup.Data().Data(), lastData.Data());
        // --------------------------------------------------------------------------------


        // Strong exception guarantee: GetStat failed with exception
        // --------------------------------------------------------------------------------
        EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
            .Times(1)
            .WillOnce(Throw(TSystemError{}));

        EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
            .Times(0);

        fup.Update();
        UNIT_ASSERT_EQUAL(fup.Data().Id(), lastData.Id());
        UNIT_ASSERT_EQUAL(fup.Data().Data(), lastData.Data());
        // --------------------------------------------------------------------------------
    }

    Y_UNIT_TEST(TSharedFileReReaderTest) {
        TFileManager::Instance().EnableTesting();

        auto fileArbiterMockPtr = MakeHolder<NSrvKernel::NTesting::TFileArbiterMock>();
        auto& fileArbiterMock = *fileArbiterMockPtr;

        TFileManager::TMockGuard guard{std::move(fileArbiterMockPtr)};

        TSharedFiles sharedFiles;
        auto fileReReader = sharedFiles.FileReReader(PATH_SAMPLE, TDuration::Seconds(1));

        TFileStat retFileStat;

        retFileStat.MTime = 1;
        retFileStat.CTime = 0;
        retFileStat.INode = 1;
        retFileStat.Size = 1;
        retFileStat.NLinks = 1;

        static constexpr size_t Iterations = 3;
        static constexpr TDuration TimeToReRead = TDuration::Seconds(2);

        for (size_t i = 0; i < Iterations; i++) {
            EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
                .WillRepeatedly(Return(retFileStat));

            TString input = STRING_SAMPLE + ToString(i);
            EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
                .WillRepeatedly(InvokeWithoutArgs([&] {
                    return MakeHolder<TStringInput>(input);
                }));

            if (i == 0) {
                // We can't do it outside the loop since calling a mocked function before setting a value is UB.
                sharedFiles.Start();
            }

            ++retFileStat.CTime;

            Sleep(TimeToReRead);

            UNIT_ASSERT(fileReReader.Data().Exists());
            UNIT_ASSERT_EQUAL(fileReReader.Data().Data(), input);
        }
    }

    Y_UNIT_TEST(TSharedFileReReaderDefaultConstructedTest) {
        TSharedFileReReader fileReReader;

        UNIT_ASSERT_EQUAL(fileReReader.Data().Id(), 0);
    }

    Y_UNIT_TEST(TSharedFileReReaderWorkersThreadSafetyTest) {
        static constexpr int32_t WorkerCount = 20;
        static const TString DataFirst = "data1";
        static const TString DataSecond = "data2";

        TFileManager::Instance().EnableTesting();

        auto fileArbiterMockPtr = MakeHolder<NSrvKernel::NTesting::TFileArbiterMock>();
        auto& fileArbiterMock = *fileArbiterMockPtr;

        TFileManager::TMockGuard guard{std::move(fileArbiterMockPtr)};

        TFileStat retFileStat;

        retFileStat.MTime = 1;
        retFileStat.CTime = 0;
        retFileStat.INode = 1;
        retFileStat.Size = 1;
        retFileStat.NLinks = 1;

        auto mockReturnData = [&fileArbiterMock](const TFileStat& stat, const TString& data) {
            EXPECT_CALL(fileArbiterMock, GetStat(StrEq(PATH_SAMPLE), false))
                .WillRepeatedly(Return(stat));

            EXPECT_CALL(fileArbiterMock, GetInput(TString(PATH_SAMPLE)))
                .WillRepeatedly(InvokeWithoutArgs([&] {
                    return MakeHolder<TStringInput>(data);
                }));
        };

        TSharedFiles sharedFiles;
        auto fileReReader = sharedFiles.FileReReader(PATH_SAMPLE, TDuration::Seconds(1));

        sharedFiles.Start();

        auto waitForStateChange = [&](const TSharedFileReReader& reader, ui64 expectedId, const TString& expectedData) {
            WaitForConditionOrDeadline([&]() { return reader.Data().Id() == expectedId; });

            UNIT_ASSERT_EQUAL(reader.Data().Id(), expectedId);
            UNIT_ASSERT_EQUAL(reader.Data().Data(), expectedData);
        };

        TSimpleBarrier barrier(WorkerCount);

        const TVector<std::pair<ui64, const TString&>> states = {
            {1, DataFirst},
            {2, DataSecond},
            {3, DataFirst},
            {4, DataSecond},
            {5, DataFirst},
            {6, DataSecond},
            {7, DataFirst},
        };

        auto worker = [&]() {
            Sleep(TDuration::Seconds(2));

            auto workerFileReReader = sharedFiles.FileReReader(PATH_SAMPLE, TDuration::Seconds(1));

            for (const auto& state : states) {
                waitForStateChange(workerFileReReader, /*expectedId =*/ state.first, state.second);
                barrier.Pass();
            }
        };

        TVector<std::thread> workers;
        for (size_t i = 0; i < WorkerCount; i++) {
            workers.emplace_back(worker);
        }

        for (const auto& state : states) {
            ++retFileStat.CTime;
            mockReturnData(retFileStat, state.second);
            waitForStateChange(fileReReader, /*expectedId =*/ state.first, state.second);
            barrier.WaitForAllWorkersToPass();
        }

        for (auto& worker : workers) {
            worker.join();
        }
    }
}
