#include "node_table_reader.h"
#include "node_table_writer.h"
#include "proto_table_reader.h"
#include "proto_table_writer.h"
#include "multi_table_reader.h"
#include "range_comparators.h"
#include <saas/library/rtyt/lib/io/ut_row.pb.h>

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

Y_UNIT_TEST_SUITE(RTYT_IO) {
    Y_UNIT_TEST(SIMPLE_WRITE_READ_TEST) {
        NRTYT::TNodeWriter writer("test_path");
        writer.AddRow(NYT::TNode()("a", 1)("b", "b"), 0);
        writer.AddRow(NYT::TNode()("a", 2)("b", "b"), 0);

        NRTYT::TNodeReader reader("test_path");
        auto row = reader.GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 1);
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsString("b"), "b");
        reader.Next();
        row = reader.GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsString("b"), "b");
    }

    Y_UNIT_TEST(PROTO_WRITE_READ_TEST) {
        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            NRTYT::TProtoWriter writer("proto_path", 
                sample.GetDescriptor());

            TRow first;
            first.SetString("abcabc");
            first.SetInt32(-1);
            first.SetFixed64(5);
            writer.AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(0);
            second.SetFixed64(0);
            writer.AddRow(second, 0);
        }

        NRTYT::TProtoReader reader("proto_path", sample.GetDescriptor());
        TRow row;
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), -1);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
        reader.Next();
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
    }
}

Y_UNIT_TEST_SUITE(RTYT_IO_RANGES) {
    Y_UNIT_TEST(PROTO_TEST) {
        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            NRTYT::TProtoWriter writer("proto_ranges_path", 
                sample.GetDescriptor());

            TRow first;
            first.SetString("abcabc");
            first.SetInt32(-1);
            first.SetFixed64(5);
            writer.AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(0);
            second.SetFixed64(0);
            writer.AddRow(second, 0);

            second.SetInt32(1);
            writer.AddRow(second, 0);
        }

        NRTYT::TProtoReader reader(
            "proto_ranges_path", 
            sample.GetDescriptor(),
            {"String"}
        );
        TRow row;
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), -1);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
        reader.Next();
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "first range, 1 row");
        reader.NextKey();

        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader.Next();
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 1);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader.Next();
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "second range, 2 rows");
    }

    Y_UNIT_TEST(PROTO_MULTIREADER_TEST) {
        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            NRTYT::TProtoWriter writer("proto_ranges_path_first",
                sample.GetDescriptor());

            TRow first;
            first.SetString("abcabc");
            first.SetInt32(0);
            first.SetFixed64(0);
            writer.AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(1);
            second.SetFixed64(1);
            writer.AddRow(second, 0);

            second.SetInt32(2);
            writer.AddRow(second, 0);
        }

        {
            NRTYT::TProtoWriter writer("proto_ranges_path_second",
                sample.GetDescriptor());

            TRow first;
            first.SetString("abcabc");
            first.SetInt32(3);
            first.SetFixed64(0);
            writer.AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(4);
            second.SetFixed64(0);
            writer.AddRow(second, 0);

            second.SetInt32(5);
            second.SetFixed64(1);
            writer.AddRow(second, 0);
        }

        NRTYT::TMultiProtoReader reader(
            TVector<TString>({"proto_ranges_path_first", "proto_ranges_path_second"}),
            TVector<const ::google::protobuf::Descriptor*>({sample.GetDescriptor(), sample.GetDescriptor()}),
            {"Fixed64"},
            {"Fixed64"}
        );

        TVector<bool> readRows(6, false);
        TRow row;
        for (int i = 0; i < 3; i++) {
            reader.ReadRow(&row);
            UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
            UNIT_ASSERT(!readRows[row.GetInt32()]);
            readRows[row.GetInt32()] = true;
            reader.Next();
        }
        UNIT_ASSERT(readRows[0] && readRows[3] && readRows[4]); // Rows wich have 0 in Fixed64 field
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "first range, 3 rows");
        reader.NextKey();

        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
        reader.Next();
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
        reader.Next();
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
        reader.Next();
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "second range, 3 rows");
        reader.NextKey();
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "reading after the end");
    }

    Y_UNIT_TEST(PROTO_MULTIWRITER_TEST) {
        using NRTYT::NTesting::TRow;
        TRow sample;
        {
            NRTYT::TProtoWriter writer(
                {"proto_ranges_writer_first", "proto_ranges_writer_second"},
                {sample.GetDescriptor(), sample.GetDescriptor()});

            TRow first;
            first.SetString("abcabc");
            first.SetInt32(0);
            first.SetFixed64(0);
            writer.AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(1);
            second.SetFixed64(1);
            writer.AddRow(second, 0);

            second.SetInt32(2);
            writer.AddRow(second, 0);

            first.SetString("abcabc");
            first.SetInt32(3);
            first.SetFixed64(0);
            writer.AddRow(first, 1);

            second.SetString("");
            second.SetInt32(4);
            second.SetFixed64(0);
            writer.AddRow(second, 1);

            second.SetInt32(5);
            second.SetFixed64(1);
            writer.AddRow(second, 1);
        }

        NRTYT::TMultiProtoReader reader(
            TVector<TString>({"proto_ranges_writer_first", "proto_ranges_writer_second"}),
            TVector<const ::google::protobuf::Descriptor*>({sample.GetDescriptor(), sample.GetDescriptor()}),
            {"Fixed64"},
            {"Fixed64"}
        );

        TVector<bool> readRows(6, false);
        TRow row;
        for (int i = 0; i < 3; i++) {
            reader.ReadRow(&row);
            UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
            UNIT_ASSERT(!readRows[row.GetInt32()]);
            readRows[row.GetInt32()] = true;
            reader.Next();
        }
        UNIT_ASSERT(readRows[0] && readRows[3] && readRows[4]); // Rows wich have 0 in Fixed64 field
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "first range, 3 rows");
        reader.NextKey();

        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
        reader.Next();
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
        reader.Next();
        reader.ReadRow(&row);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
        reader.Next();
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "second range, 3 rows");
        reader.NextKey();
        UNIT_CHECK_GENERATED_EXCEPTION_C(reader.ReadRow(&row), yexception, "reading after the end");
    }
}

