#include "entity.h"
#include "writer.h"

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

#include <kernel/indexann/interface/reader.h>

class TArrayAnnWriter : public NIndexAnn::IDocDataWriter {
public:
    TArrayAnnWriter(const TArchiveOwner::TPtr& storage)
        : Storage_(storage)
    {}

    void Add(ui32 breakId, ui32 regionId, ui32 streamId, TArrayRef<const char> data) override {
        Y_ASSERT(DocId_ != static_cast<ui32>(-1));

        TArrayAnnRecord record(
            breakId,
            regionId,
            0,
            streamId,
            NIndexAnn::DataRegionToUi32(data.data(), data.size())
        );
        Records_.push_back(record);
    }

    void StartDoc(ui32 docId) override {
        DocId_ = docId;
    }

    void FinishDoc() override {
        Y_ASSERT(DocId_ != static_cast<ui32>(-1));

        if (Records_.empty())
            return;

        std::sort(Records_.begin(), Records_.end(), TArrayAnnRecordLess());

        Storage_->PutDocument(TBlob::Copy(Records_.data(), Records_.size() * sizeof(TArrayAnnRecord)), DocId_);
        Records_.clear();
    }

private:
    TArchiveOwner::TPtr Storage_;
    ui32 DocId_ = -1;
    TVector<TArrayAnnRecord> Records_;
};

THolder<NIndexAnn::IDocDataWriter> NewArrayDocDataWriter(const TArchiveOwner::TPtr& storage) {
    return THolder(new TArrayAnnWriter(storage));
}
