#include "dir_digest.h"

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

#include <util/generic/algorithm.h>

class TTestFilesSizeInfo : public TFilesSizeInfo {
public:
    TTestFilesSizeInfo()
        : TFilesSizeInfo()
    {
    }

    TTestFilesSizeInfo(const std::initializer_list<std::pair<TString, ui32>>& entries)
        : TFilesSizeInfo()
    {
        for (const auto& [key, value]: entries) {
            AddInfoUnlocked(key, value);
        }
    }

    TVector<std::pair<TString, ui32>> GetReportDump() const {
        NJson::TJsonValue value = GetReport(/* isHumanReadable = */ false);
        UNIT_ASSERT(value.IsMap());
        TVector<std::pair<TString, ui32>> res;
        res.reserve(value.GetMap().size());
        for (const auto& entry: value.GetMap()) {
            res.emplace_back(entry.first, entry.second.GetUInteger());
        }
        Sort(res);
        return res;
    }
};

static inline void AssertSameVec(const TVector<std::pair<TString, ui32>>& res, const TVector<std::pair<TString, ui32>>& etalon) {
    UNIT_ASSERT_EQUAL_C(etalon.size(), res.size(), ToString(etalon.size()) << " != " << ToString(res.size()));
    for (ui32 i = 0; i < res.size(); ++i) {
        UNIT_ASSERT_EQUAL_C(etalon[i].first, res[i].first, ToString(i) << " " << etalon[i].first << " != " << res[i].first);
        UNIT_ASSERT_EQUAL_C(etalon[i].second, res[i].second, ToString(i) << " " << etalon[i].second << " != " << res[i].second);
    }
}

Y_UNIT_TEST_SUITE(TRtyUtilFileSizesAggr) {
    Y_UNIT_TEST(TestFileSizesAggregation) {
        const std::initializer_list<std::pair<TString, ui32>> entries = {
            {"files_info", 1103},
            {"indexarc.fat", 1600000},
            {"indexarc.part.0", 76425033},
            {"indexarc.part.0.hdr", 58628},
            {"indexarc.part.0.meta", 5},
            {"indexarc.part.1", 76425033},
            {"indexarc.part.1.hdr", 58628},
            {"indexarc.part.1.meta", 5},
            {"indexddk.rty", 7600000},
            {"indexddk.rty.hdr", 22},
            {"index.docurl", 21763360},
            {"indexerf.rty", 44000000},
            {"indexerf.rty.hdr", 127},
            {"indexfrq", 400000},
            {"indexfullarc.base.fat", 1600000},
            {"indexfullarc.base.meta", 6},
            {"indexfullarc.base.part.0", 183205},
            {"indexfullarc.base.part.0.hdr", 11868},
            {"indexfullarc.base.part.0.meta", 5},
            {"indexfullarc.base.part.1", 896083},
            {"indexfullarc.base.part.1.hdr", 50973},
            {"indexfullarc.base.part.1.meta", 5},
            {"indexfullarc.base.part.2", 648652 },
            {"indexfullarc.base.part.2.hdr", 36397},
            {"indexfullarc.base.part.2.meta", 5},
            {"indexann.fat", 19781960},
            {"indexann.inv", 37910666},
            {"indexann.key", 2089987},
            {"indexann.part.0", 19597384},
            {"indexann.part.0.hdr", 545668},
            {"indexann.part.0.meta", 52},
            {"indexann.part.1", 0},
            {"indexann.part.1.hdr", 0},
            {"indexann.part.1.meta", 13},
            {"indexfactorann.fat", 19781960},
            {"indexfactorann.inv", 116729492},
            {"indexfactorann.key", 7629561},
            {"indexfactorann.part.0", 147687772},
            {"indexfactorann.part.0.hdr", 1630108},
            {"indexfactorann.part.0.meta", 52},
            {"indexfactorann.part.1", 0},
            {"indexfactorann.part.1.hdr", 0},
            {"indexfactorann.part.1.meta", 13},
            {"indexfactorann.sent", 23218696},
            {"indexfactorann.sent.title", 9890984}
        };

        TTestFilesSizeInfo fileSizeInfo(entries);
        const TVector<std::pair<TString, ui32>> res = fileSizeInfo.GetReportDump();

        const TVector<std::pair<TString, ui32>> res_entries = {
            {"__COUNT", 27},
            {"__SUM", 638253511},
            {"files_info", 1103},
            {"index.docurl", 21763360},
            {"indexann.fat", 19781960},
            {"indexann.inv", 37910666},
            {"indexann.key", 2089987},
            {"indexann.part", 19597384},
            {"indexann.part.hdr", 545668},
            {"indexann.part.hdr_count", 2},
            {"indexann.part.meta", 65},
            {"indexann.part.meta_count", 2},
            {"indexann.part_count", 2},
            {"indexarc.fat", 1600000},
            {"indexarc.part", 152850066},
            {"indexarc.part.hdr", 117256},
            {"indexarc.part.hdr_count", 2},
            {"indexarc.part.meta",10},
            {"indexarc.part.meta_count",2},
            {"indexarc.part_count", 2},
            {"indexddk.rty", 7600000},
            {"indexddk.rty.hdr", 22},
            {"indexerf.rty", 44000000},
            {"indexerf.rty.hdr", 127},
            {"indexfactorann.fat", 19781960},
            {"indexfactorann.inv", 116729492},
            {"indexfactorann.key", 7629561},
            {"indexfactorann.part", 147687772},
            {"indexfactorann.part.hdr", 1630108},
            {"indexfactorann.part.hdr_count", 2},
            {"indexfactorann.part.meta", 65},
            {"indexfactorann.part.meta_count", 2},
            {"indexfactorann.part_count", 2},
            {"indexfactorann.sent", 23218696},
            {"indexfactorann.sent.title", 9890984},
            {"indexfrq", 400000},
            {"indexfullarc.base.fat", 1600000},
            {"indexfullarc.base.meta", 6},
            {"indexfullarc.base.part", 1727940},
            {"indexfullarc.base.part.hdr", 99238},
            {"indexfullarc.base.part.hdr_count", 3},
            {"indexfullarc.base.part.meta", 15},
            {"indexfullarc.base.part.meta_count", 3},
            {"indexfullarc.base.part_count", 3},
        };
        AssertSameVec(res, res_entries);
        UNIT_ASSERT_EQUAL(638253511, fileSizeInfo.GetTotalSize());
    }

    Y_UNIT_TEST(TestPartFileNamesProcessing) {
        const std::initializer_list<std::pair<TString, ui32>> entries = {
            {"files_info", 1},
            {"indexarc1.part.", 1},
            {"indexarc2.part.0", 2},
            {"indexarc2.part.1", 2},
            {"indexarc3.part.0.hdr", 3},
            {"indexarc3.part.1.hdr", 3},
        };

        TTestFilesSizeInfo fileSizeInfo(entries);
        const TVector<std::pair<TString, ui32>> res = fileSizeInfo.GetReportDump();

        const TVector<std::pair<TString, ui32>> res_entries = {
            {"__COUNT", 4},
            {"__SUM", 12},
            {"files_info", 1},
            {"indexarc1.part.", 1},
            {"indexarc2.part", 4},
            {"indexarc2.part_count", 2},
            {"indexarc3.part.hdr", 6},
            {"indexarc3.part.hdr_count", 2},
        };
        AssertSameVec(res, res_entries);
        UNIT_ASSERT_EQUAL(12, fileSizeInfo.GetTotalSize());
    }

        Y_UNIT_TEST(TestMerge) {
        const std::initializer_list<std::pair<TString, ui32>> entries1 = {
            {"files_info", 1},
            {"indexarc1.part.", 1},
            {"indexarc2.part.0", 2},
            {"indexarc2.part.1", 2},
            {"indexarc3.part.0.hdr", 3},
            {"indexarc3.part.1.hdr", 3},
            {"indexarc6", 3},
        };

        const std::initializer_list<std::pair<TString, ui32>> entries2 = {
            {"files_info", 1},
            {"indexarc1.part.", 1},
            {"indexarc2.part.0", 2},
            {"indexarc2.part.1", 2},
            {"indexarc3.part.0.hdr", 3},
            {"indexarc3.part.1.hdr", 3},
            {"indexarc7", 3},
        };

        TTestFilesSizeInfo fileSizeInfo1(entries1);
        TTestFilesSizeInfo fileSizeInfo2(entries2);
        TTestFilesSizeInfo merged;
        merged.Merge(fileSizeInfo1);
        merged.Merge(fileSizeInfo2);
        const TVector<std::pair<TString, ui32>> res = merged.GetReportDump();

        const TVector<std::pair<TString, ui32>> res_entries = {
            {"__COUNT", 12},
            {"__SUM", 30},
            {"files_info", 2},
            {"files_info_count", 2},
            {"indexarc1.part.", 2},
            {"indexarc1.part._count", 2},
            {"indexarc2.part", 8},
            {"indexarc2.part_count", 4},
            {"indexarc3.part.hdr", 12},
            {"indexarc3.part.hdr_count", 4},
            {"indexarc6", 3},
            {"indexarc7", 3},
        };
        AssertSameVec(res, res_entries);
        UNIT_ASSERT_EQUAL(30, merged.GetTotalSize());
    }
}
