#pragma once

#include <util/generic/algorithm.h>
#include <util/generic/vector.h>
#include <util/generic/yexception.h>

namespace NFusion {
    template <typename TIterator, typename TComparer>
    class TIteratorHeap {
    public:
        using TIteratorPtr = TIterator*;
    private:
        using THeap = TVector<TIteratorPtr>;
        using THeapIter = typename THeap::iterator;
        THeap Heap;
        THeapIter HeapEnd;
    public:
        TIteratorHeap()
        {
            Heap.resize(16);
            HeapEnd = Heap.begin();
        }

        void Add(TIteratorPtr source) {
            Y_ENSURE(HeapEnd < Heap.end());
            if (!source->Valid())
                return;

            *HeapEnd = source;
            MakeHeap<THeapIter>(Heap.begin(), ++HeapEnd, TComparer());
        }

        TIteratorPtr GetHead() const {
            Y_ASSERT(HeapEnd > Heap.begin());
            return Heap.front();
        }

        bool Valid() {
            Y_ASSERT(!(HeapEnd > Heap.begin()) || GetHead()->Valid());
            return HeapEnd > Heap.begin();
        }

        void Next() {
            Y_ASSERT(HeapEnd >= Heap.begin());
            TIteratorPtr head = Heap.front();
            PopHeap<THeapIter>(Heap.begin(), HeapEnd, TComparer());

            head->Next();

            if (head->Valid()) {
                PushHeap<THeapIter>(Heap.begin(), HeapEnd, TComparer());
            } else {
                --HeapEnd;
                *HeapEnd = nullptr;
            }
        }
    };
}
