#include "wad_merge_host_adapter.h"
#include <kernel/groupattrs/attrmap.h>
#include <kernel/groupattrs/docsattrswriter.h>
#include <kernel/groupattrs/mutdocattrs.h>

#include <robot/jupiter/library/rtdoc/merger/groupattrsmap.h>
#include <robot/jupiter/library/rtdoc/file/docidmap_io.h>
#include <library/cpp/testing/unittest/registar.h>
#include <util/folder/path.h>
#include <util/stream/file.h>
#include <util/stream/str.h>

namespace NRtDoc {
    class TWadHostAdapterUnitTest: public NUnitTest::TTestBase {
        UNIT_TEST_SUITE(TWadHostAdapterUnitTest)
            UNIT_TEST(TestAaCreationForTest);
            UNIT_TEST(TestAdapter);
        UNIT_TEST_SUITE_END();

        static void WriteIndexAa(const TString& indexAaPrefix, const TVector<ui32> hostIdByDocId) {
            NMemoryAttrMap::TMemoryConfigWrapper memConfig;
            NGroupingAttrs::TConfig::Type idType = NGroupingAttrs::TConfig::I32;
            ui32 domainAttrNum, hostAttrNum;
            memConfig.AddMetainfo("d", domainAttrNum, idType);
            memConfig.AddMetainfo("h", hostAttrNum, idType);

            NGroupingAttrs::TConfig config = memConfig.Config();

            const TFsPath tmpDir = TFsPath(indexAaPrefix).Dirname();

            TFileOutput out(indexAaPrefix + "aa");
            out.SetFinishPropagateMode(true);

            NGroupingAttrs::TDocsAttrsWriter wr(
                    config,
                    NGroupingAttrs::TDocsAttrsWriter::DefaultWritableVersion,
                    NGroupingAttrs::TConfig::Search,
                    (TString)tmpDir,
                    out,
                    ""
                );

            NGroupingAttrs::TMutableDocAttrs docAttrs(config);

            for (ui32 docId = 0; docId < hostIdByDocId.size(); ++docId) {
                docAttrs.Clear();
                docAttrs.SetAttr(hostAttrNum, hostIdByDocId[docId]);
                docAttrs.SetAttr(domainAttrNum, hostIdByDocId[docId]); //for the test, we set domainId == hostId. It does not really matter
                wr.Write(docAttrs);
            }

            wr.Close();
            out.Finish();
        }

        void Cleanup() {
            TFsPath("./segments").ForceDelete();
        }

        void SetUp() override {
            Cleanup();
            TFsPath("./segments").MkDirs();
        }

        void TearDown() override {
            Cleanup();
        }

        void TestAaCreationForTest() {
            //
            // Tests that the WriteIndexAa() helper method is implemented correctly
            //
            const TString indexaaPrefix = "./segments/index";
            TVector<ui32> expectedHostIds{2, 1, 2, 1};
            WriteIndexAa(indexaaPrefix, expectedHostIds);
            TCategMap indexaa(indexaaPrefix);
            UNIT_ASSERT_EQUAL(4, indexaa.Size());

            TVector<ui32> actualHostIds;
            for(ui32 docId = 0; docId < indexaa.Size(); ++docId) {
                actualHostIds.push_back(indexaa[docId].GetHostId());
            }
            UNIT_ASSERT_EQUAL(expectedHostIds, actualHostIds);
        }

        class TMockSlaveMerger : public IWadMerger {
        public:
            struct TMockChunk {
                const TString WadPath;
                const TString DocIdMapStr;
                const size_t DocIdMapSize;
                const TString DisplayName;
            };

        public:
            void Init(const TString& outputFileName) override {
                OutputFileName = outputFileName;
            }

            void Add(const TString& wadPath, THolder<TDocIdMap> docIdMap, IWadPatcher::TPtr wadPatcher, const TString& displayName, bool /*isDelta*/) override {
                UNIT_ASSERT(!wadPatcher);
                UNIT_ASSERT(docIdMap);
                TString docIdMapStr;
                TStringOutput ss(docIdMapStr);
                for (ui32 docId = 0; docId < docIdMap->Size(); ++docId) {
                    if (docId > 0) {
                        ss << " ";
                    }
                    const ui32 newId = docIdMap->Map(docId);
                    if (newId != TDocIdMap::DeletedDocument()) {
                        ss << docIdMap->Map(docId);
                    } else {
                        ss << "X";
                    }
                }

                Chunks.push_back(TMockChunk{wadPath, docIdMapStr, docIdMap->Size(), displayName});
            }

            bool Empty() const override {
                return Chunks.empty();
            }

            void Finish() override {
                MockCalledFinish = true;
            }

        public:
            TVector<TMockChunk> Chunks;
            bool MockCalledFinish = false;

        private:
            TString OutputFileName;
        };

        void TestAdapter() {
            //
            // Checks that TWadHostAdapter generates correct HostId remappings
            //

            TIntrusivePtr<TMockSlaveMerger> slave = MakeIntrusive<TMockSlaveMerger>();
            TWadHostAdapter::TPtr merger = MakeIntrusive<TWadHostAdapter>(slave);

            const TString wadName = "indexregerf.wad"; //can be any name - the mock class does nothing
            TVector<TFsPath> segs{TFsPath("./segments/index_1"), TFsPath("./segments/index_2"), TFsPath("./segments/index_3")};
            const TFsPath trgPath = "./segments/index_4_merge";
            TFsPath segPath;

            {
                segPath = segs[0];
                segPath.MkDirs();
                NRtDoc::TDocIdMap srcDocIds;
                *srcDocIds.MutableData() = TVector<ui32>{0, 1, TDocIdMap::DeletedDocument(), TDocIdMap::DeletedDocument() /*drop hostId 2*/, 2};
                TVector<ui32> srcHostIds{1,3,1,2,0}; //mapped as is
                WriteIndexAa(segPath / "index", srcHostIds);
                TDocIdMapIo::Save(segPath / "map", &srcDocIds);
            }
            {
                // empty segment (no docs)
                segPath = segs[1];
                segPath.MkDirs();
                NRtDoc::TDocIdMap srcDocIds;
                *srcDocIds.MutableData() = TVector<ui32>{TDocIdMap::DeletedDocument(), TDocIdMap::DeletedDocument()};
                TVector<ui32> srcHostIds{}; // empty indexaa, why not
                WriteIndexAa(segPath / "index", srcHostIds);
                TDocIdMapIo::Save(segPath / "map", &srcDocIds);
            }
            {
                segPath = segs[2];
                segPath.MkDirs();
                NRtDoc::TDocIdMap srcDocIds;
                *srcDocIds.MutableData() = TVector<ui32>{3, 4, 5};
                TVector<ui32> srcHostIds{0, 0, 1}; //mapped to {0, 0, 3}
                WriteIndexAa(segPath / "index", srcHostIds);
                TDocIdMapIo::Save(segPath / "map", &srcDocIds);
            }
            {
                segPath = trgPath;
                segPath.MkDirs();
                TVector<ui32> trgHostIds{1, 2, 0, 0, 0, 3};
                WriteIndexAa(segPath / "index", trgHostIds);
            }

            merger->Init(trgPath / wadName);
            for (auto& seg : segs) {
                THolder<TDocIdMap> docIdMap = MakeHolder<TDocIdMap>();
                TDocIdMapIo::Load(docIdMap.Get(), seg / "map");
                merger->Add(seg / wadName, std::move(docIdMap), IWadPatcher::TPtr(), seg, false);
            }
            merger->Finish();

            UNIT_ASSERT_C(slave->MockCalledFinish, "Finish() not called");
            const auto& chunks = slave->Chunks;
            UNIT_ASSERT_EQUAL(segs.size(), chunks.size());
            for (ui32 segNo = 0; segNo < segs.size(); ++segNo) {
                UNIT_ASSERT_EQUAL(TFsPath(segs[segNo]) / wadName, TFsPath(chunks[segNo].WadPath));
                UNIT_ASSERT_EQUAL((TString)segs[segNo], chunks[segNo].DisplayName);
            }

            Y_ASSERT(3 == segs.size());
            UNIT_ASSERT_EQUAL("X 1 X 2", chunks.at(0).DocIdMapStr);
            UNIT_ASSERT_EQUAL("", chunks.at(1).DocIdMapStr);
            UNIT_ASSERT_EQUAL("0 3", chunks.at(2).DocIdMapStr);

            merger.Drop();
            slave.Drop();
        }
    };

}

UNIT_TEST_SUITE_REGISTRATION(NRtDoc::TWadHostAdapterUnitTest);
