#include "yt_chunked_input.h"

#include <mapreduce/yt/interface/io.h>

#include <util/generic/size_literals.h>
#include <util/generic/string.h>
#include <util/generic/queue.h>
#include <util/stream/file.h>
#include <util/stream/null.h>
#include <util/system/tempfile.h>

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

namespace {
    using TMockProtoTable = TQueue<NSaas::TYTBlobBase>;

    class TMockProtoTableReader : public NYT::IProtoReaderImpl
    {
    public:
        TMockProtoTableReader(TMockProtoTable& table)
            : Table(table)
        {
        }
        virtual ~TMockProtoTableReader() = default;

        void ReadRow(NYT::Message* message) override {
            message->CopyFrom(Table.front());
        }
        bool IsValid() const override {
            return !Table.empty();
        }
        void Next() override {
            Table.pop();
        }
        ui32 GetTableIndex() const override {
            return 0;
        }
        ui32 GetRangeIndex() const override {
            return 0;
        }
        ui64 GetRowIndex() const override {
            return 0;
        }
        void NextKey() override {
        }

    private:
        TMockProtoTable& Table;
    };

    TYTBlob PrepareFileChunk(const TString& fileName, const TString& fileContents, ui32 shardId = 0, ui32 chunkId = 0) {
        TYTBlob b;
        b.SetShardId(shardId);
        b.GetBlobRef().SetName(fileName);
        b.GetBlobRef().SetData(fileContents);
        b.GetBlobRef().SetDataLength(fileContents.size());
        b.GetBlobRef().SetPartIndex(chunkId);
        return b;
    }

    void CheckFile(TFsPath filePath, const TString& fileContents) {
        UNIT_ASSERT(filePath.Exists());
        UNIT_ASSERT_VALUES_EQUAL(fileContents, TUnbufferedFileInput(filePath).ReadAll());
    }

    struct TTempFiles : public TVector<TTempFile> {
        using TBase = TVector<TTempFile>;
        TTempFiles(const TVector<TString>& names)
            : TBase(names.begin(), names.end())
        {
        }
    };

    struct TWrapperThrottled : public TYTChunkedInputToThrottledFiles {
        TWrapperThrottled(TTableReader* tableReader, const TFsPath& path)
            : TYTChunkedInputToThrottledFiles(tableReader, 10_MB, path, EYTBlobFormat::Old)
        {
        }
    };

    struct TFailedConstructorOutput : public TNullOutput {
        struct TEx : public yexception {
        };
        TFailedConstructorOutput() {
            ythrow TEx() << "bad constructor";
        }
    };

    struct TFailedWriteOutput : public TNullOutput {
        struct TEx : public yexception {
        };
        void DoWrite(const void*, size_t) override {
            ythrow TEx() << "bad write";
        }
    };

    struct TFailedFlushOutput : public TNullOutput {
        struct TEx : public yexception {
        };
        void DoFlush() override {
            ythrow TEx() << "bad flush";
        }
    };

    struct TFailedFinishOutput : public TNullOutput {
        struct TEx : public yexception {
        };
        void DoFinish() override {
            ythrow TEx() << "bad finish";
        }
    };

    /*
    This can't be tested because ~IOutputStream is noexcept, program gets terminated.

    struct TFailedDestructorOutput : public TNullOutput {
        ~TFailedDestructorOutput() noexcept(false) override {
            ythrow yexception() << "bad destructor";
        }
    };
    */

    template <typename TStreamType>
    class TFailedChunkedInput : public IYTChunkedInputProcessor {
    public:
        TFailedChunkedInput(TTableReader* tableReader, TFsPath = TFsPath(), EYTBlobFormat format = EYTBlobFormat::Old)
            : IYTChunkedInputProcessor(tableReader, format)
    {
    }

    protected:
        THolder<IOutputStream> CreateOutputStream(TString) override {
            return THolder(new TStreamType);
        }
    };

    // Creates empty files. Needed to verify that our test harness actually checks file contents.
    class TBrokenChunkedInput : public IYTChunkedInputProcessor {
    public:
        TBrokenChunkedInput(TTableReader* tableReader, TFsPath root = TFsPath(), EYTBlobFormat format = EYTBlobFormat::Old)
            : IYTChunkedInputProcessor(tableReader, format)
            , Root(root)
    {
    }

    protected:
        THolder<IOutputStream> CreateOutputStream(TString name) override {
            (Root / name).Touch();
            return THolder(new TNullOutput);
        }
        TFsPath Root;
    };


}

template <typename TChunkedInput, bool ExpectException = false, typename TException = yexception>
void DoTestOneFile() {
    const TString fileName = "file.out";
    const TString fileContents = "test";
    const ui32 shardId = 321;
    TTempFile rm(fileName);
    TMockProtoTable t;
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileName, fileContents, shardId)));
    NYT::TTableReader<NSaas::TYTBlobBase> r(new TMockProtoTableReader(t));
    NYT::TTableReaderPtr<NYT::Message> abstractReader = NYT::CreateGenericProtobufReader(&r);
    TChunkedInput itf(abstractReader.Get(), TFsPath::Cwd());
    if (ExpectException) {
        UNIT_ASSERT_EXCEPTION(itf.ProcessAll(), TException);
    } else {
        itf.ProcessAll();
        UNIT_ASSERT(itf.GetShardId() == shardId);
        CheckFile(fileName, fileContents);
    }
}

template <typename TChunkedInput, bool ExpectException = false, typename TException = yexception>
void DoTestMultipleFiles() {
    const TVector<TString> fileNames{"file1.out", "file2.out", "file3.out"};
    const TString contents[] = {"test1", "test2\n", "xxxxxxxxxxxxxxxxxxxxxxxxxxx"};
    TTempFiles rm(fileNames);
    TMockProtoTable t;
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[0], contents[0])));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[1], contents[1])));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[2], contents[2])));
    NYT::TTableReader<NSaas::TYTBlobBase> r(new TMockProtoTableReader(t));
    NYT::TTableReaderPtr<NYT::Message> abstractReader = NYT::CreateGenericProtobufReader(&r);
    TChunkedInput itf(abstractReader.Get(), TFsPath::Cwd());
    if (ExpectException) {
        UNIT_ASSERT_EXCEPTION(itf.ProcessAll(), TException);
    } else {
        itf.ProcessAll();
        CheckFile(fileNames[0], contents[0]);
        CheckFile(fileNames[1], contents[1]);
        CheckFile(fileNames[2], contents[2]);
    }
}

template <typename TChunkedInput, bool ExpectException = false, typename TException = yexception>
void DoTestMultipleFileMixed() {
    const TVector<TString> fileNames{"file1.out", "file2.out", "file3.out"};
    const TString lines[] = {"line1\n", "line2\n"};
    const TString fileContents = lines[0] + lines[1];
    TTempFiles rm(fileNames);
    TMockProtoTable t;
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[0], lines[0], 0, 0)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[1], lines[0], 0, 0)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[2], lines[0], 0, 0)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[0], lines[1], 0, 1)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[1], lines[1], 0, 1)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[2], lines[1], 0, 1)));
    NYT::TTableReader<NSaas::TYTBlobBase> r(new TMockProtoTableReader(t));
    NYT::TTableReaderPtr<NYT::Message> abstractReader = NYT::CreateGenericProtobufReader(&r);
    TChunkedInput itf(abstractReader.Get(), TFsPath::Cwd());
    if (ExpectException) {
        UNIT_ASSERT_EXCEPTION(itf.ProcessAll(), TException);
    } else {
        itf.ProcessAll();
        CheckFile(fileNames[0], fileContents);
        CheckFile(fileNames[1], fileContents);
        CheckFile(fileNames[2], fileContents);
    }
}

template <typename TChunkedInput>
void DoTestWrongChunkOrderException() {
    const TString fileName = "file.out";
    const TString fileContents = "test";
    TTempFile rm(fileName);
    TMockProtoTable t;
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileName, fileContents, 0, 15)));
    NYT::TTableReader<NSaas::TYTBlobBase> r(new TMockProtoTableReader(t));
    NYT::TTableReaderPtr<NYT::Message> abstractReader = NYT::CreateGenericProtobufReader(&r);
    TChunkedInput itf(abstractReader.Get(), TFsPath::Cwd());
    UNIT_ASSERT_EXCEPTION(itf.ProcessAll(), yexception);
}

template <typename TChunkedInput>
void DoTestWrongChunkOrderException2() {
    const TString fileName = "file.out";
    const TString fileContents = "test";
    TTempFile rm(fileName);
    TMockProtoTable t;
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileName, fileContents, 0, 0)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileName, fileContents, 0, 1)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileName, fileContents, 0, 1)));
    NYT::TTableReader<NSaas::TYTBlobBase> r(new TMockProtoTableReader(t));
    NYT::TTableReaderPtr<NYT::Message> abstractReader = NYT::CreateGenericProtobufReader(&r);
    TChunkedInput itf(abstractReader.Get(), TFsPath::Cwd());
    UNIT_ASSERT_EXCEPTION(itf.ProcessAll(), yexception);
}

template <typename TChunkedInput>
void DoTestUnexpectedShardIdException() {
    const TVector<TString> fileNames{"file1.out", "file2.out"};
    const TString fileContents = "test";
    TTempFiles rm(fileNames);
    TMockProtoTable t;
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[0], fileContents, 111)));
    t.push(static_cast<NSaas::TYTBlobBase>(PrepareFileChunk(fileNames[1], fileContents, 222)));
    NYT::TTableReader<NSaas::TYTBlobBase> r(new TMockProtoTableReader(t));
    NYT::TTableReaderPtr<NYT::Message> abstractReader = NYT::CreateGenericProtobufReader(&r);
    TChunkedInput itf(abstractReader.Get(), TFsPath::Cwd());
    UNIT_ASSERT_EXCEPTION(itf.ProcessAll(), yexception);
}

#define TEST(name, type) Y_UNIT_TEST(name) { Do ## name<type>() ;}
#define TEST_EX(name, type, exception_type) Y_UNIT_TEST(name) { Do ## name<type, true, exception_type>() ;}
#define TEST_SHOULD_FAIL(name, type) Y_UNIT_TEST(name) { UNIT_ASSERT_TEST_FAILS(Do ## name<type>()) ;}

Y_UNIT_TEST_SUITE(ChunkedInputSuite) {
    TEST(TestOneFile, TYTChunkedInputToFiles)
    TEST(TestMultipleFiles, TYTChunkedInputToFiles)
    TEST(TestMultipleFileMixed, TYTChunkedInputToFiles)
    TEST(TestWrongChunkOrderException, TYTChunkedInputToFiles)
    TEST(TestWrongChunkOrderException2, TYTChunkedInputToFiles)
    TEST(TestUnexpectedShardIdException, TYTChunkedInputToFiles)
}

Y_UNIT_TEST_SUITE(ChunkedInputThrottledSuite) {
    TEST(TestOneFile, TWrapperThrottled)
    TEST(TestMultipleFiles, TWrapperThrottled)
    TEST_EX(TestMultipleFileMixed, TWrapperThrottled, yexception) // Resume not supported yet
    TEST(TestWrongChunkOrderException, TWrapperThrottled)
    TEST(TestWrongChunkOrderException2, TWrapperThrottled)
    TEST(TestUnexpectedShardIdException, TWrapperThrottled)
}

Y_UNIT_TEST_SUITE(ChunkedInputFailedConstructorSuite) {
    TEST_EX(TestOneFile, TFailedChunkedInput<TFailedConstructorOutput>,  TFailedConstructorOutput::TEx)
    TEST_EX(TestMultipleFiles, TFailedChunkedInput<TFailedConstructorOutput>, TFailedConstructorOutput::TEx)
}

Y_UNIT_TEST_SUITE(ChunkedInputFailedWriteSuite) {
    TEST_EX(TestOneFile, TFailedChunkedInput<TFailedWriteOutput>, TFailedWriteOutput::TEx)
    TEST_EX(TestMultipleFiles, TFailedChunkedInput<TFailedWriteOutput>, TFailedWriteOutput::TEx)
}

Y_UNIT_TEST_SUITE(ChunkedInputFailedFlushSuite) {
    TEST_EX(TestOneFile, TFailedChunkedInput<TFailedFlushOutput>, TFailedFlushOutput::TEx)
    TEST_EX(TestMultipleFiles, TFailedChunkedInput<TFailedFlushOutput>, TFailedFlushOutput::TEx)
}

Y_UNIT_TEST_SUITE(ChunkedInputFailedFinishSuite) {
    TEST_EX(TestOneFile, TFailedChunkedInput<TFailedFinishOutput>, TFailedFinishOutput::TEx)
    TEST_EX(TestMultipleFiles, TFailedChunkedInput<TFailedFinishOutput>, TFailedFinishOutput::TEx)
}

Y_UNIT_TEST_SUITE(ChunkedInputBrokenSuite) {
    TEST_SHOULD_FAIL(TestOneFile, TBrokenChunkedInput)
    TEST_SHOULD_FAIL(TestMultipleFiles, TBrokenChunkedInput)
}
