#pragma once

#include <robot/jupiter/library/rtdoc/file/model/docidmap.h>
#include <robot/jupiter/library/rtdoc/file/docidmap_io.h>

namespace NRtDoc {
    struct TDocIdRemapper {
        TDocIdRemapper(TString mappingFileName, size_t totalSegments)
            : MappingFileName(std::move(mappingFileName))
            , Segments(totalSegments)
        {
        }

        void AddSegmentSize(size_t size) {
            Y_ENSURE(InitSegments < Segments.size());
            Y_ENSURE(ProcessedSegments == 0);
            Segments[InitSegments++] = size;
        }

        void InitFinalDocIds() {
            Y_ENSURE(InitSegments == Segments.size());
            Y_ENSURE(ProcessedSegments == 0);
            size_t docCount = std::accumulate(Segments.begin(), Segments.end(), 0);
            FinalDocIds.MutableData()->resize(docCount);
        }

        template <typename TCallback>
        void ProcessSegment(const NRtDoc::TDocIdMap& docIds, TCallback&& callback) {
            Y_ENSURE(InitSegments == Segments.size());
            Y_ENSURE(ProcessedSegments < Segments.size());

            const TVector<ui32>& docIdMap = *docIds.GetData();
            const size_t segmentSize = Segments[ProcessedSegments];
            const size_t mapSize = docIdMap.size();

            NRtDoc::TDocIdMap::TData& finalDocIds = *FinalDocIds.MutableData();
            if (NewDocId + mapSize >= finalDocIds.size()) {
                finalDocIds.resize(NewDocId + mapSize);
            }

            size_t docId = 0;
            for (; docId < mapSize; docId++) {
                if (docIdMap.at(docId) == TDocIdMap::DeletedDocument())
                    continue;

                const ui32 finalDocId = docIdMap[docId];
                finalDocIds[NewDocId] = finalDocId;

                if (docId >= segmentSize) {
                    ++NewDocId;
                    continue;
                }
                const ui32 outputDocId = UseFinalDocIds() ? finalDocId : NewDocId;

                callback(docId, outputDocId);

                ++NewDocId;
            }

            ++ProcessedSegments;
        }

        void Finish() {
            Y_ENSURE(ProcessedSegments == Segments.size());
            if (!UseFinalDocIds()) {
                FinalDocIds.MutableData()->resize(NewDocId);
                TDocIdMapIo::Save(TFsPath(MappingFileName), &FinalDocIds);
            }
        }

    private:
        bool UseFinalDocIds() const {
            return MappingFileName.empty();
        }

        const TString MappingFileName;
        TDocIdMap FinalDocIds;
        TVector<size_t> Segments;
        size_t InitSegments = 0;
        size_t ProcessedSegments = 0;
        size_t NewDocId = 0;
    };
}
