#include "entity.h"
#include "iterator.h"

#include <saas/rtyserver/components/ann/storage/interfaces.h>

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

class TArrayDocDataIterator: public NIndexAnn::IDocDataIterator {
public:
    TArrayDocDataIterator(const NRTYAnn::IDocBlobAccessor* docAccessor, const NRTYAnn::IStreamIdRemapper* remapper)
        : DocAccessor_(docAccessor)
        , Remapper_(remapper)
    {}

    void Restart(const NIndexAnn::THitMask& mask) override {
        Y_ASSERT(mask.HasDoc());

        if (Mask_.DocId() != mask.DocId() || !Mask_.HasDoc()) {
            Blob_ = DocAccessor_->Get(mask.DocId());
            Y_VERIFY(Blob_.Size() % sizeof(TArrayAnnRecord) == 0);
        }

        Mask_ = mask;
        Current_.SetDocId(mask.DocId());

        Ptr_ = std::lower_bound(
            reinterpret_cast<const TArrayAnnRecord*>(Blob_.Begin()),
            reinterpret_cast<const TArrayAnnRecord*>(Blob_.End()),
            TArrayAnnRecord(mask.Break(), mask.Region(), 0, mask.Stream(), 0),
            TArrayAnnRecordLess()
        );

        Advance();
    }

    bool Valid() const override {
        return Valid_;
    }

    const NIndexAnn::THit& Current() const override {
        return Current_;
    }

    const NIndexAnn::THit* Next() override {
        Advance();

        return Valid_ ? &Current_ : nullptr;
    }

private:
    void Advance() {
        const TArrayAnnRecord* end = reinterpret_cast<const TArrayAnnRecord*>(Blob_.End());
        if (Ptr_ >= end) {
            Valid_ = false;
            return;
        }

        Current_.SetBreak(Ptr_->Break);
        Current_.SetRegion(Ptr_->Region);
        Current_.SetStream(Remapper_ ? Remapper_->GetGlobalStreamIndex(Ptr_->Stream) : Ptr_->Stream);
        Current_.SetValue(Ptr_->Value);
        Valid_ = Mask_.Matches(Current_);
        Ptr_++;
    }

private:
    const NRTYAnn::IDocBlobAccessor* DocAccessor_;
    const NRTYAnn::IStreamIdRemapper* Remapper_;
    NIndexAnn::THitMask Mask_;
    bool Valid_ = false;
    NIndexAnn::THit Current_;
    TBlob Blob_;
    const TArrayAnnRecord* Ptr_ = nullptr;
};

THolder<NIndexAnn::IDocDataIterator> NewArrayDocDataIterator(const NRTYAnn::IDocBlobAccessor* docAccessor, const NRTYAnn::IStreamIdRemapper* remapper) {
    return THolder(new TArrayDocDataIterator(docAccessor, remapper));
}
