#include "merger.h"

#include <kernel/doc_url_index/doc_url_index.h>

#include <library/cpp/on_disk/2d_array/array2d_writer.h>

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

using namespace NRTYArchive;

Y_UNIT_TEST_SUITE(TDocUrlMergerSuite) {
    void InitLog() {
        if (!GlobalLogInitialized())
            DoInitGlobalLog("console", 7, false, false);
    }

    void Clear(const TFsPath& archive) {
        TVector<TFsPath> children;
        archive.Parent().List(children);
        for (auto& f : children)
            if (f.GetName().StartsWith(archive.GetName()))
                f.ForceDelete();

    }

    class TTestDecoderSplit : public IRTYMergerDocIdDecoder {
    public:
        TTestDecoderSplit(ui32 size)
            : HalfSize(size / 2)
        {}

        TRTYMergerAddress Decode(ui32 /*clusterId*/, ui32 docId) const override {
            ui32 rDocid = 2 * HalfSize - docId - 1;
            return TRTYMergerAddress( rDocid / HalfSize, rDocid % HalfSize);
        }

        bool Check(ui32 /*clusterId*/, ui32 /*docId*/) const override {
            return true;
        }
        ui32 GetSizeOfCluster(ui32 clusterId) const override {
            return clusterId ? 0 : 2 * HalfSize;
        }

        ui32 GetSize() const override {
            return 0;
        }
        bool IsValidClusterId(ui32 /*clusterId*/) const override {
            return true;
        }
        ui32 GetNewDocsCount(ui32 /*clusterId*/) const override {
            return HalfSize;
        }

        bool NewToOld(ui32 clusterId, ui32 docId, TRTYMergerAddress& addr) const override {
            if (clusterId >= 2 || docId >= HalfSize)
                return false;
            addr.ClusterId = 0;
            ui32 rDocid = clusterId * HalfSize + docId;
            addr.DocId = 2 * HalfSize - rDocid - 1;
            return true;
        }

        void PatchDestMap(ui32 /*clusterId*/, const TVector<ui32>& /*remap*/) override {
        }
    private:
        ui32 HalfSize;
    };

    class TTestDecoderFusion : public IRTYMergerDocIdDecoder {
    public:
        TTestDecoderFusion(ui32 size)
            : HalfSize(size / 2) {
        }

        TRTYMergerAddress Decode(ui32 clusterId, ui32 docId) const override {
            ui32 rDocid = clusterId * HalfSize + docId;
            ui32 docid = 2 * HalfSize - rDocid - 1;
            return TRTYMergerAddress(0, docid);
        }
        bool Check(ui32 /*clusterId*/, ui32 /*docId*/) const override {
            return true;
        }
        ui32 GetSizeOfCluster(ui32 clusterId) const override {
            return clusterId < 2 ? 2 * HalfSize : 0;
        }

        ui32 GetSize() const override {
            return 0;
        }
        bool IsValidClusterId(ui32 /*clusterId*/) const override {
            return true;
        }
        ui32 GetNewDocsCount(ui32 clusterId) const override {
            if (!clusterId || clusterId == Max<ui32>())
                    return 2 * HalfSize;
            return 0;
        }
        bool NewToOld(ui32 clusterId, ui32 docId, TRTYMergerAddress& addr) const override {
            if (clusterId || docId >= 2 * HalfSize)
                return false;
            ui32 rDocid = 2 * HalfSize - docId - 1;
            addr.ClusterId = rDocid / HalfSize;
            addr.DocId = rDocid % HalfSize;
            return true;
        }
        void PatchDestMap(ui32 /*clusterId*/, const TVector<ui32>& /*remap*/) override {
        }
    private:
        ui32 HalfSize;
    };

    Y_UNIT_TEST(TestMerge) {
        InitLog();
        Clear("src");
        Clear("dst1");
        Clear("dst2");
        Clear("dst");
        TFsPath("src").MkDirs();
        TFsPath("dst1").MkDirs();
        TFsPath("dst2").MkDirs();
        TFsPath("dst").MkDirs();

        TFile2DArrayWriter<ui32, char> writer("src/index.docurl");
        bool isFirst = true;
        for (ui32 i = 0; i < 12; ++i) {
            if (!isFirst) {
                writer.NewLine();
            } else {
                isFirst = false;
            }
            const TString url = ToString(i);
            for (const char c: url) {
                writer.Write(c);
            }
        }
        writer.Finish();
        TVector<TString> srcs(1, "src");
        TVector<TString> dsts;
        dsts.push_back("dst1");
        dsts.push_back("dst2");
        TTestDecoderSplit decoderS(12);
        TTestDecoderFusion decoderF(12);
        TRTYMerger merger(nullptr, TRTYMerger::otDocUrl);

        TRTYMerger::TContext ctxS(srcs, dsts, "", &decoderS, nullptr);
        merger.MergeIndicies(ctxS);
        for (ui32 d = 0; d < dsts.size(); ++d) {
            TDocUrlIndexReader reader(dsts[d] + "/index.docurl");
            for (ui32 docid = 0; docid < 6; ++docid) {
                TString url(reader.Get(docid));
                TRTYMergerAddress oldAddr;
                UNIT_ASSERT(decoderS.NewToOld(d, docid, oldAddr));
                UNIT_ASSERT_EQUAL(url, ToString(oldAddr.DocId));
            }
        }
        TRTYMerger::TContext ctxF(dsts, TVector<TString>(1, "dst"), "", &decoderF, nullptr);
        merger.MergeIndicies(ctxF);
        TDocUrlIndexReader reader("dst/index.docurl");
        for (ui32 docid = 0; docid < 12; ++docid) {
            TString url(reader.Get(docid));
            TRTYMergerAddress oldAddr;
            UNIT_ASSERT(decoderF.NewToOld(0, docid, oldAddr));
            UNIT_ASSERT(decoderS.NewToOld(oldAddr.ClusterId, oldAddr.DocId, oldAddr));
            UNIT_ASSERT_EQUAL(oldAddr.DocId, docid);
            UNIT_ASSERT_EQUAL(url, ToString(docid));
        }
    }
}

