#include "iterator.h"
#include "key.h"

#include <library/cpp/containers/comptrie/comptrie.h>
#include <library/cpp/logger/global/global.h>

class RtyAnnUi32 {
public:
    ui32 Value;
    RtyAnnUi32(ui32 value=0) : Value(value) {};
    inline RtyAnnUi32& operator <<= (ui32 shift) { Value <<= shift; return *this;}
    inline RtyAnnUi32& operator |= (ui32 shift) { Value |= shift; return *this; }
    inline operator ui32() { return Value; }
};

template<>
struct TCompactTrieKeySelector<RtyAnnUi32> {
    typedef NRTYAnn::TFixedSizeKey<3> TKey;
    typedef NRTYAnn::TFixedSizeKey<3> TKeyBuf;
};

template <>
void Out<NRTYAnn::TFixedSizeKey<3>>(IOutputStream& output, TTypeTraits<NRTYAnn::TFixedSizeKey<3>>::TFuncParam type) {
    output << type.Quote();
}

namespace NRTYAnn {
    class TStreamIteratorImpl {
        using TTrie = TCompactTrie<RtyAnnUi32, ui32>;
    public:
        TStreamIteratorImpl(const IDocBlobAccessor* docAccessor, TAtomicSharedPtr<IDataBlockAccessor> dataAccessor, const IStreamIdRemapper* remapper = nullptr)
            : DocAccessor(docAccessor)
            , DataAccessor(dataAccessor)
            , Remapper(remapper)
        {
        }

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

        bool Valid() const {
            return Valid_;
        }

        const NIndexAnn::THit* Next() {
            ++CurrentStream;
            if (CurrentStream == DataAccessor->GetDataCount()) {
                ++CurrentIter;
                CurrentStream = 0;
            }
            Valid_ = CurrentIter != EndIter;
            if (Valid_) {
                UpdateCurrent();
                return &Current_;
            }
            return nullptr;
        }

        void UpdateCurrent() {
            CHECK_WITH_LOG(Valid());
            DataAccessor->SetBlockBegining(StreamData + CurrentIter.GetValue());
            auto data = DataAccessor->Get(CurrentStream);
            ui32 value = NIndexAnn::DataRegionToUi32(data.begin(), data.size());

            const auto& key = CurrentIter.GetKey();
            Current_ = NIndexAnn::THit(
                Mask.DocId(),
                key[0],
                key[1],
                Remapper ? Remapper->GetGlobalStreamIndex(CurrentStream) : CurrentStream,
                value
            );
            Valid_ = Mask.Matches(Current_);
        }

        void Restart(const NIndexAnn::THitMask& mask) {
            Y_ENSURE(mask.HasDoc(), "iterating over all docs is not supported yet");
            Mask = mask;
            CurrentStream = 0;
            if (Current_.DocId() != Mask.DocId()) {
                RawData = DocAccessor->Get(Mask.DocId());
                if (RawData.Empty()) {
                    Valid_ = false;
                    return;
                }
                TStreamDataHeader* header = (TStreamDataHeader*)RawData.Data();
                Trie.Init((char*)RawData.Data() + sizeof(TStreamDataHeader), header->Size);
                StreamData = (char*)RawData.Data() + header->Size + sizeof(TStreamDataHeader);
            }

            if (!Mask.HasBreak()) {
                CurrentIter = Trie.Begin();
                EndIter = Trie.End();
            } else {
                TVector<ui32> begin = {Mask.Break()};
                if (Mask.HasRegion()) {
                    begin.push_back(Mask.Region());
                }
                TVector<ui32> end = begin;
                ++end.back();
                CurrentIter = Trie.UpperBound(TTrie::TKeyBuf(begin.data(), begin.size()));
                EndIter = Trie.UpperBound(TTrie::TKeyBuf(end.data(), end.size()));
            }
            Valid_ = CurrentIter != EndIter;
            if (Valid_) {
                UpdateCurrent();
            }
        }
    private:
        const IDocBlobAccessor* DocAccessor;
        TAtomicSharedPtr<IDataBlockAccessor> DataAccessor;
        const IStreamIdRemapper* Remapper;

        TBlob RawData;
        TTrie Trie;
        char* StreamData;

        NIndexAnn::THit Current_ = NIndexAnn::THit(-1, -1, -1, -1, -1);
        NIndexAnn::THitMask Mask;
        bool Valid_ = false;

        ui32 CurrentStream = 0;
        TCompactTrie<RtyAnnUi32, ui32>::TConstIterator CurrentIter;
        TCompactTrie<RtyAnnUi32, ui32>::TConstIterator EndIter;
    };


    TStreamIterator::TStreamIterator(const IDocBlobAccessor* docAccessor, TAtomicSharedPtr<IDataBlockAccessor> dataAccessor, const IStreamIdRemapper* remapper)
        : Impl(new TStreamIteratorImpl(docAccessor, dataAccessor, remapper))
    {}

    TStreamIterator::~TStreamIterator()
    {}

    const NIndexAnn::THit& TStreamIterator::Current() const {
        return Impl->Current();
    }

    bool TStreamIterator::Valid() const {
        return Impl->Valid();
    }

    const NIndexAnn::THit* TStreamIterator::Next() {
        return Impl->Next();
    }

    void TStreamIterator::Restart(const NIndexAnn::THitMask& mask) {
        Impl->Restart(mask);
    }
}
