#include "client.h"

#include <saas/library/rtyt/lib/operation/factory.h>
#include <saas/library/rtyt/lib/io/ut_row.pb.h>
#include <library/cpp/testing/unittest/registar.h>
#include <mapreduce/yt/interface/operation.h>
#include <mapreduce/yt/interface/io.h>
#include <util/system/fs.h>

namespace {
    class TProtoMultiplyingMapperSaveLoad
        : public NYT::IMapper<NYT::TTableReader<NRTYT::NTesting::TRow>, NYT::TTableWriter<NRTYT::NTesting::TRow>>
    {
    public:
        TProtoMultiplyingMapperSaveLoad() = default;

        void Do(TReader* input, TWriter* output) override {
            NRTYT::NTesting::TRow row;
            for (; input->IsValid(); input->Next()) {
                row = input->GetRow();
                i64 result = row.GetInt32() * row.GetFixed64();
                row.SetFixed64(result);
                output->AddRow(row);
            }
        }
    };

    REGISTER_RTYT_MAPPER(TProtoMultiplyingMapperSaveLoad);

    TFsPath path = "./rtyt_test";

    void Cleanup() {
        TVector<TFsPath> children;
        if (path.Exists()) {
            path.List(children);
            for (const auto& child : children) {
                child.ForceDelete();
            }
        }
    }
}

Y_UNIT_TEST_SUITE(RTYT_SAVE_LOAD) {
    Y_UNIT_TEST(SIMPLE_SAVE_LOAD) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            NRTYT::TClientBase client(".");
            auto writer = client.CreateTableWriter("//test_table", *sample.GetDescriptor());
            TRow first;
            first.SetString("a");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);
            writer->Finish();
            client.Sync();
        }
        NRTYT::TClientBase client(".");
        client.Mount(/*storagePath = */".", /*mountPoint = */"//1");
        auto reader = client.CreateTableReader<TRow>("//1/test_table");
        TRow first = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(first.GetString(), "a");
        UNIT_ASSERT_VALUES_EQUAL(first.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(first.GetFixed64(), 5);
    }

    Y_UNIT_TEST(COMPLEX_SAVE_LOAD) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            NRTYT::TClientBase client(".");
            client.Create("//test", NYT::ENodeType::NT_MAP, NYT::TCreateOptions());
            auto writer = client.CreateTableWriter("//test/test_table", *sample.GetDescriptor());
            TRow first;
            first.SetString("a");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);
            writer->Finish();
            client.Copy("//test", "//test2", NYT::TCopyOptions());
            client.Sync();
        }
        NRTYT::TClientBase client;
        client.Mount(/*storagePath = */".", /*mountPoint = */"//1");
        auto reader = client.CreateTableReader<TRow>("//1/test2/test_table");
        TRow first = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(first.GetString(), "a");
        UNIT_ASSERT_VALUES_EQUAL(first.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(first.GetFixed64(), 5);
    }
    Y_UNIT_TEST(MOVE_DIR_SAVE_LOAD) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        TFsPath path("./test1");
        {
            path.MkDir();
            NFs::SetCurrentWorkingDirectory(path.GetPath());
            NRTYT::TClientBase client(".");
            auto writer = client.CreateTableWriter("//test_table", *sample.GetDescriptor());
            TRow first;
            first.SetString("a");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);
            writer->Finish();
            client.Sync();
            NFs::SetCurrentWorkingDirectory("..");
        }
        path.RenameTo("./test2");
        NRTYT::TClientBase client;
        client.Mount(/*storagePath = */"./test2", /*mountPoint = */"//1");
        auto reader = client.CreateTableReader<TRow>("//1/test_table");
        TRow first = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(first.GetString(), "a");
        UNIT_ASSERT_VALUES_EQUAL(first.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(first.GetFixed64(), 5);
    }

    Y_UNIT_TEST(MULTI_DIR_SAVE_LOAD) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            TFsPath path("./test1");
            path.MkDir();
            NFs::SetCurrentWorkingDirectory(path.GetPath());
            NRTYT::TClientBase client(".");
            auto writer = client.CreateTableWriter("//test_table", *sample.GetDescriptor());
            TRow first;
            first.SetString("a");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);
            writer->Finish();
            client.Sync();
            NFs::SetCurrentWorkingDirectory("..");
        }
        {
            TFsPath path("./test2");
            path.MkDir();
            NFs::SetCurrentWorkingDirectory(path.GetPath());
            NRTYT::TClientBase client(".");
            auto writer = client.CreateTableWriter("//test_table2", *sample.GetDescriptor());
            TRow first;
            first.SetString("b");
            first.SetInt32(1);
            first.SetFixed64(0);
            writer->AddRow(first, 0);
            writer->Finish();
            client.Sync();
            NFs::SetCurrentWorkingDirectory("..");
        }
        NRTYT::TClientBase client;
        client.Mount(/*storagePath = */"./test1", /*mountPoint = */"//1");
        client.Mount(/*storagePath = */"./test2", /*mountPoint = */"//2");
        {
            auto reader = client.CreateTableReader<TRow>("//1/test_table");
            TRow first = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(first.GetString(), "a");
            UNIT_ASSERT_VALUES_EQUAL(first.GetInt32(), 0);
            UNIT_ASSERT_VALUES_EQUAL(first.GetFixed64(), 5);
        }
        {
            auto reader = client.CreateTableReader<TRow>("//2/test_table2");
            TRow first = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(first.GetString(), "b");
            UNIT_ASSERT_VALUES_EQUAL(first.GetInt32(), 1);
            UNIT_ASSERT_VALUES_EQUAL(first.GetFixed64(), 0);
        }
    }
    Y_UNIT_TEST(SAVE_LOAD_UNMOUNT_READ) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            TFsPath path("./test1");
            path.MkDir();
            NFs::SetCurrentWorkingDirectory(path.GetPath());
            NRTYT::TClientBase client(".");
            auto writer = client.CreateTableWriter("//test_table", *sample.GetDescriptor());
            TRow first;
            first.SetString("a");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);
            writer->Finish();
            client.Sync();
            NFs::SetCurrentWorkingDirectory("..");
        }
        {
            TFsPath path("./test2");
            path.MkDir();
            NFs::SetCurrentWorkingDirectory(path.GetPath());
            NRTYT::TClientBase client(".");
            auto writer = client.CreateTableWriter("//test_table2", *sample.GetDescriptor());
            TRow first;
            first.SetString("b");
            first.SetInt32(1);
            first.SetFixed64(0);
            writer->AddRow(first, 0);
            writer->Finish();
            client.Sync();
            NFs::SetCurrentWorkingDirectory("..");
        }
        NRTYT::TClientBase client;
        client.Mount(/*storagePath = */"./test1", /*mountPoint = */"//1");
        {
            auto reader = client.CreateTableReader<TRow>("//1/test_table");
            TRow first = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(first.GetString(), "a");
            UNIT_ASSERT_VALUES_EQUAL(first.GetInt32(), 0);
            UNIT_ASSERT_VALUES_EQUAL(first.GetFixed64(), 5);
        }
        client.Umount(client.GetNodeId("//1"));

        client.Mount(/*storagePath = */"./test2", /*mountPoint = */"//1");
        {
            auto reader = client.CreateTableReader<TRow>("//1/test_table2");
            TRow first = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(first.GetString(), "b");
            UNIT_ASSERT_VALUES_EQUAL(first.GetInt32(), 1);
            UNIT_ASSERT_VALUES_EQUAL(first.GetFixed64(), 0);
        }
    }
}


Y_UNIT_TEST_SUITE(SAVE_LOAD_OPERATIONS) {
    Y_UNIT_TEST(PROTO_MAPPER_TEST) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        
        TFsPath path("./saveload/test1");
        {
            path.MkDirs();
            NFs::SetCurrentWorkingDirectory(path.GetPath());

            NRTYT::TClientBase client(".");
            client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
            client.Create("//test1/test2", NYT::NT_TABLE);
            client.Create("//test1/output", NYT::NT_TABLE);
            TRow sample;
            {
                auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
                TRow first;
                first.SetString("abcabc");
                first.SetInt32(0);
                first.SetFixed64(5);
                writer->AddRow(first, 0);

                TRow second;
                second.SetString("");
                second.SetInt32(2);
                second.SetFixed64(2);
                writer->AddRow(second, 0);
            }
            client.Sync();
            NFs::SetCurrentWorkingDirectory("../..");
        }
        path.RenameTo("./saveload/test2");
        {
            NRTYT::TClientBase client;
            client.Mount("./saveload/test2", "//in");
            client.Mount("./saveload/test3", "//out");

            NYT::TMapOperationSpec spec;
            spec.AddInput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//in/test1/test2")).AddOutput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//out/output"));
            client.Map(spec, new TProtoMultiplyingMapperSaveLoad(), NYT::TOperationOptions());
            client.Sync(client.GetNodeId("//out"));
        }

        NRTYT::TClientBase client;
        client.Mount("./saveload/test3", "//result");
        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//result/output"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 4);
    }
}
