#pragma once

#include <util/generic/vector.h>
#include <util/generic/algorithm.h>
#include <util/system/yassert.h>

namespace NDoom {


/**
 *
 * Optimized TOffroadWadBuffer which takes hits pre-sorted by DocId
 */
//TODO: (yrum) Remove copypaste - the class should be inherited from TOffroadWadBuffer and moved to kernel/doom

template<class Hit>
class TOffroadWadDocBuffer {
    using THit = Hit;
public:
    TOffroadWadDocBuffer() {}

    void Reset() {
        DocId_ = 0;
        TotalDocCount_ = 0;
        HitsByDocId_.clear();
    }

    void Write(const THit& hit) {
        ui32 docId = DocId_;
        if (docId >= HitsByDocId_.size())
            HitsByDocId_.resize(docId + 1);

        //hit.GetDocId() contains TermId, push it as is
        //The caller is responsible for it to match sequential TermIds generated by KeyBuffer_.WriteKey()
        HitsByDocId_[docId].push_back(hit);
    }

    void Seek(ui32 docId) {
        Y_VERIFY_DEBUG(docId >= HitsByDocId_.size() || HitsByDocId_[docId].empty(), "docis must be unique");
        DocId_ = docId;
        ++TotalDocCount_;
    }

    size_t Size() const {
        return HitsByDocId_.size();
    }

    const TVector<THit>& Hits(ui32 docId) {
        return HitsByDocId_[docId];
    }

    void RemapTermIds(const TVector<ui32>& positions) {
        Y_ENSURE(HitsByDocId_.size() == TotalDocCount_);
        for (size_t docId = 0, docIdEnd = HitsByDocId_.size(); docId < docIdEnd; ++docId) {
            for (THit& hit : HitsByDocId_[docId]) {
                hit.SetDocId(positions[hit.DocId()]);
            }
            Sort(HitsByDocId_[docId].begin(), HitsByDocId_[docId].end());
        }
    }

private:
    ui32 DocId_;
    ui32 TotalDocCount_;
    TVector<TVector<THit>> HitsByDocId_;
};


} // namespace NDoom

