#include "wad_decoded_merge.h"
#include <robot/jupiter/library/rtdoc/merger/reg_erf.h>
#include <kernel/doom/offroad_reg_herf_wad/reg_herf_io.h>
#include <kernel/doom/standard_models_storage/standard_models_storage.h>
#include <library/cpp/testing/unittest/registar.h>
#include <ysite/yandex/erf_format/erf_format.h>
#include <util/folder/path.h>
#include <util/stream/file.h>
#include <util/stream/str.h>

namespace NRtDoc {
    class TWadDecodedMergeUnitTest: public NUnitTest::TTestBase {
        UNIT_TEST_SUITE(TWadDecodedMergeUnitTest)
            UNIT_TEST(TestTest);
            UNIT_TEST(TestWadMerge);
        UNIT_TEST_SUITE_END();

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

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

        void TearDown() override {
            Cleanup();
        }

        struct TMockErfData {
            ui32 HostId;
            ui32 Tag;
            ui32 RegionsNum;

            TString DebugString() const {
                TString res;
                TStringStream ss(res);
                ss << "(" << HostId << " " << Tag << " " << RegionsNum << ")";
                ss.Finish();
                return ss.Str();
            }
        };

        static TString DumpMockData(const TVector<TMockErfData>& records) {
            TString res;
            TStringStream ss(res);
            for (const auto& record : records) {
                if (!ss.Empty())
                    ss << "  ";
                ss << record.DebugString();
            }
            ss.Finish();
            return ss.Str();
        }

        static void WriteRegHerf(const TString& fileName, const TVector<TMockErfData>& records) {
            using TIo = NDoom::TRegHostErfIo;
            TIo::TWriter::TModel model = TStandardIoModelsStorage::Model<TIo::TWriter::TModel>(TIo::DefaultModel);
            TIo::TWriter wr(model, fileName);

            for (const TMockErfData& record : records) {
                const ui32 hostId = record.HostId;
                for (ui32 region = 0; region < record.RegionsNum; ++region) {
                    TIo::TKey key(hostId, region);
                    TRegHostErfInfo regherf;
                    memset(&regherf, 0, sizeof(regherf));
                    regherf.RegionalHostRank = record.Tag;
                    wr.Write(key, &regherf);
                }
            }
            wr.Finish();
        }

        static void ReadRegHerf(const TString& fileName, TVector<TMockErfData>& records) {
            using TIo = NDoom::TRegHostErfIo;
            TIo::TReader rd(fileName);

            records.clear();

            ui32 oldHostId = Max<ui32>();
            TMockErfData* record = nullptr;
            while(true) {
                TIo::TKey key;
                const TIo::TData* data;
                if (!rd.Next(&key, &data)) {
                    break;
                }

                const ui32 hostId = key.DocId();
                if (hostId != oldHostId) {
                    records.emplace_back();
                    record = &records.back();
                    record->HostId = hostId;
                    record->RegionsNum = 1;
                    oldHostId = hostId;
                } else {
                    Y_ASSERT(record);
                    record->RegionsNum++;
                }
                record->Tag = (int) data->RegionalHostRank;
            }
        }

        void TestTest() {
            TVector<TMockErfData> expected{
                {0, 5, 2},
                {1, 4, 2},
                {3, 3, 3}
            };

            WriteRegHerf("./testregherf.wad", expected);
            TVector<TMockErfData> actual;
            ReadRegHerf("./testregherf.wad", actual);

            UNIT_ASSERT_EQUAL(expected.size(), actual.size());

            const TString expectedStr = DumpMockData(expected);
            const TString actualStr = DumpMockData(actual);
            UNIT_ASSERT_EQUAL(expectedStr, actualStr);
        }

        void TestWadMerge() {
            //
            // Checks that TWadDecodedMerger handles regherf as expected
            //
            using TIo = NDoom::TRegHostErfIo;
            using TMerger = TWadDecodedMerger<TIo, NRtDoc::TKeyAccessor<TIo::TKey>>;
            TIntrusivePtr<TMerger> merger = MakeIntrusive<TMerger>();
            const TString wadName = "indexregherf.wad";
            TVector<TFsPath> segs{TFsPath("./segments/index_1"), TFsPath("./segments/index_2"), TFsPath("./segments/index_3")};
            TVector<NRtDoc::TDocIdMap> hostIdMaps;
            const TFsPath trgPath = "./segments/index_4_merge";
            TFsPath segPath;

            {
                const TVector<TMockErfData> mockData{
                    {0, 11, 2},
                    {1, 12, 1},
                    {3, 14, 3}
                };
                segPath = segs[0];
                segPath.MkDirs();
                hostIdMaps.emplace_back();
                NRtDoc::TDocIdMap& srcHostMap = hostIdMaps.back();
                *srcHostMap.MutableData() = TVector<ui32>{1, TDocIdMap::DeletedDocument(), TDocIdMap::DeletedDocument(), 4};
                WriteRegHerf(segPath / wadName, mockData);
            }
            {
                // empty segment (no docs)
                segPath = segs[1];
                segPath.MkDirs();
                const TVector<TMockErfData> mockData;
                hostIdMaps.emplace_back();
                NRtDoc::TDocIdMap& srcHostMap = hostIdMaps.back();
                *srcHostMap.MutableData() = TVector<ui32>{TDocIdMap::DeletedDocument(), TDocIdMap::DeletedDocument()};
                WriteRegHerf(segPath / wadName, mockData);
            }
            {
                // replace hostid=2(tag=12) with something with 5 regions, insert two new hosts (with tag=10 and tag=13)
                const TVector<TMockErfData> mockData{
                    {0, 10, 2},
                    {1, 12, 5},
                    {2, 13, 2}
                };
                segPath = segs[2];
                segPath.MkDirs();
                hostIdMaps.emplace_back();
                NRtDoc::TDocIdMap& srcHostMap = hostIdMaps.back();
                *srcHostMap.MutableData() = TVector<ui32>{0, 2, 3};
                WriteRegHerf(segPath / wadName, mockData);
            }
            {
                segPath = trgPath;
                segPath.MkDirs();
            }

            merger->Init(trgPath / wadName);
            for (ui32 segNo = 0; segNo < segs.size(); ++segNo) {
                THolder<TDocIdMap> hostIdMap = MakeHolder<TDocIdMap>(std::move(hostIdMaps[segNo]));
                merger->Add(segs[segNo] / wadName, std::move(hostIdMap), IWadPatcher::TPtr(), segs[segNo], false);
            }
            merger->Finish();

            const TVector<TMockErfData> expectedData{
                {0, 10, 2},
                {1, 11, 2},
                {2, 12, 5},
                {3, 13, 2},
                {4, 14, 3}
            };
            TVector<TMockErfData> actualData;
            ReadRegHerf(trgPath / wadName, actualData);
            const TString expectedStr = DumpMockData(expectedData);
            const TString actualStr = DumpMockData(actualData);
            UNIT_ASSERT_EQUAL(expectedStr, actualStr);

            merger.Drop();
        }
    };

}

UNIT_TEST_SUITE_REGISTRATION(NRtDoc::TWadDecodedMergeUnitTest);
