#pragma once

#include <library/cpp/offroad/flat/flat_writer.h>
#include <library/cpp/offroad/flat/flat_searcher.h>
#include <library/cpp/offroad/custom/ui64_vectorizer.h>

#include <util/generic/maybe.h>
#include <bitset> // for typed helpers

namespace NRTYSerf {
    using TColMask = ui64;
    struct TRowfatKey {
        ui32 RowId;
        ui32 OffsetDw;
    };

    // Level 1: key = TRowfatSeek, keyData = key.Position, data = TRowfatSubkey
    // Level 2: key = TRowfatSubkey, keyData = key.OffsetDw, data = TColMask
    struct TRowfatSeek {
        ui32 RowId;
        ui32 Position; // in subfat

        static constexpr ui32 SuffixBits = 4;

        inline operator bool() const {
            return RowId != Max<ui32>();
        }
    };

    struct TRowfatSubkey {
        ui32 RowId;
        ui32 OffsetDw;
    };

    // We do not use TTupleInputBuffer, so our "Subtractor" handles the "tuple" without indirection
    struct TRowfatSubkeyOps {
        using T = TRowfatSubkey;
        Y_FORCE_INLINE static void Integrate(T value, T delta, T& next) {
            Y_ASSERT(!IsStopMark(delta));
            next.RowId = value.RowId + delta.RowId - 1;
            next.OffsetDw = value.OffsetDw + delta.OffsetDw;
        }

        Y_FORCE_INLINE static void Differentiate(T value, T& delta, T next) {
            Y_ASSERT(next.RowId >= value.RowId);
            delta.RowId = next.RowId - value.RowId + 1;
            delta.OffsetDw = next.OffsetDw - value.OffsetDw;
        }

        Y_FORCE_INLINE static bool IsStopMark(T delta) {
            return delta.RowId == 0;
        }

        inline static constexpr T StopMark() {
            return TRowfatSubkey{0, 0};
        }
    };

    struct TRowfatSeekVectorizer {
        enum {
            TupleSize = 2
        };

        template <class Key, class Storage>
        Y_FORCE_INLINE static void Scatter(Key&& key, Storage&& data) {
            data[0] = key.RowId;
            data[1] = key.Position; // keydata
        }

        template <class Key, class Storage>
        Y_FORCE_INLINE static void Gather(Storage&& data, Key&& key) {
            key->RowId = data[0];
            key->Position = data[1];
        }
    };

    struct TRowfatSubkeyVectorizer {
        enum {
            TupleSize = 2
        };

        template <class Key, class Storage>
        Y_FORCE_INLINE static void Scatter(Key&& key, Storage&& data) {
            data[0] = key.RowId;
            data[1] = key.OffsetDw; // keydata
        }

        template <class Key, class Storage>
        Y_FORCE_INLINE static void Gather(Storage&& data, Key&& key) {
            key->RowId = data[0];
            key->OffsetDw = data[1];
        }
    };

    struct TRowfatWriterIo {
    public:
        using TKey = TRowfatKey;
        using TSeek = TRowfatSeek;
        using TSubkey = TRowfatSubkey;
        using TData = TColMask;
        using TKeyVectorizer = TRowfatSeekVectorizer;
        using TSubkeyVectorizer = TRowfatSubkeyVectorizer;
        using TSubkeyOps = TRowfatSubkeyOps;
    };

    class TRowfatWriter : TRowfatWriterIo {
    public:
        using TIo = TRowfatWriterIo;

        using TFatWriter = NOffroad::TFlatWriter<TSeek, TSubkey, TKeyVectorizer, TSubkeyVectorizer>;
        using TSubfatWriter = NOffroad::TFlatWriter<TSubkey, TData, TSubkeyVectorizer, NOffroad::TUi64Vectorizer>;

    public:
        TRowfatWriter(IOutputStream* fat, IOutputStream* subfat) {
            Reset(fat, subfat);
        }

        Y_FORCE_INLINE static constexpr ui32 AsPrefix(ui32 rowId) {
            return rowId & ~((1 << TRowfatSeek::SuffixBits) - 1);
        }

        void Reset(IOutputStream* fat, IOutputStream* subfat) {
            Y_ASSERT(fat && subfat);
            LastFatKey_ = TKey();
            BaseSubkey_ = TSubkey();
            LastRowIdPrefix_ = 0;
            Fat_ = fat;
            SubFat_ = subfat;
            FatWriter_.Reset(Fat_);
            SubfatWriter_.Reset(SubFat_);
        }

        void WriteHit(const TData& data) {
            TSubkey delta;
            TSubkeyOps::Differentiate(BaseSubkey_, delta, LastSubkey_);
            SubfatWriter_.Write(delta, data);
            BaseSubkey_ = LastSubkey_;
        }

        void WriteKey(const TKey& aKey) {
            LastSubkey_ = TSubkey{aKey.RowId, aKey.OffsetDw};
            ui32 rowIdPrefix = AsPrefix(aKey.RowId);
            if (rowIdPrefix != LastRowIdPrefix_) {
                WriteSeek(aKey, LastSubkey_);
                LastRowIdPrefix_ = rowIdPrefix;
            }
        }

        void Finish() {
            if (SubfatWriter_.Size())
                SubfatWriter_.Write(TSubkeyOps::StopMark(), TData());

            SubfatWriter_.Finish();
            FatWriter_.Finish();
        }

    private:
        void WriteSeek(const TKey& key, const TSubkey& subkey) {
            SubfatWriter_.Write(TSubkeyOps::StopMark(), TData());

            LastRowIdPrefix_ = AsPrefix(key.RowId);
            TSeek s{LastRowIdPrefix_, (ui32)SubfatWriter_.Size()};
            BaseSubkey_ = subkey;
            FatWriter_.Write(s, BaseSubkey_);
        }

    private:
        TKey LastFatKey_;
        TSubkey BaseSubkey_;
        TSubkey LastSubkey_;
        ui32 LastRowIdPrefix_;

        IOutputStream* Fat_;
        IOutputStream* SubFat_;
        TFatWriter FatWriter_;
        TSubfatWriter SubfatWriter_;
    };

    class TRowfatSearcher: public TRowfatWriterIo {
    public:
        using TIo = TRowfatWriterIo;

    private:
        using TFatSearcher = NOffroad::TFlatSearcher<TSeek, TSubkey, TKeyVectorizer, TSubkeyVectorizer>;
        using TSubfatSearcher = NOffroad::TFlatSearcher<TSubkey, TData, TSubkeyVectorizer, NOffroad::TUi64Vectorizer>;

    public:
        // rarely used - typed representation of TReadResult
        struct TSimpleRow {
            ui32 RowId = 0;
            std::bitset<64> FactorMask;
            size_t ByteOffset = Max<size_t>();

            TSimpleRow() = default;

            TSimpleRow(ui32 rowId, ui32 offsetDw, ui64 factorMask)
                : RowId(rowId)
                , FactorMask(factorMask)
                , ByteOffset(offsetDw * sizeof(ui32))
            {
            }
        };

    public:
        TRowfatSearcher(const TBlob& fat, const TBlob& subfat) {
            Reset(fat, subfat);
        }

        void Reset(const TBlob& fat, const TBlob& subfat) {
            Fat_.Reset(fat);
            Subfat_.Reset(subfat);
            Validate();
        }

        // simplified ReadRow() for convenience
        bool ReadRow(ui32 rowId, TSimpleRow& simple) const {
            simple.ByteOffset = Max<size_t>();
            ReadRow(rowId, [&simple](ui32 rowId, ui32 offsetDw, ui64 factorMask) {
                Y_ENSURE(offsetDw != Max<ui32>());
                simple = TSimpleRow(rowId, offsetDw, factorMask);
            });
            return simple.ByteOffset != Max<size_t>();
        }

        template <typename TFoo>
        void ReadRow(ui32 rowId, TFoo&& rowConsumer) const {
            size_t fatSeek = SeekFat(rowId);
            ReadFrom(fatSeek, [&, this, rowId](ui32 subfatPos, TIo::TSubkey value) mutable {
                if (value.RowId < rowId)
                    return true;
                if (value.RowId == rowId) {
                    rowConsumer(rowId, value.OffsetDw, /*factorMask=*/Subfat_.ReadData(subfatPos));
                }
                return false;
            });
        }

        template <typename TFoo>
        void ReadAll(ui32 maxRowId, TFoo&& rowConsumer) const {
            ReadFrom(0ul, [&, this, lastRow = Max<ui32>()](ui32 subfatPos, TIo::TSubkey value) mutable {
                if (maxRowId < value.RowId)
                    return false;
                if (lastRow == value.RowId) {
                    Y_ASSERT(0 /* only one record per (docId,rowId), please */);
                    return true; // skip & continue
                }
                lastRow = value.RowId;
                return rowConsumer(lastRow, value.OffsetDw, /*factorMask=*/Subfat_.ReadData(subfatPos));
            });
        }

        bool IsEmpty() const {
            return Subfat_.Size() == 0;
        }

        size_t Size() const {
            if (!CachedRowsCount_.Defined()) {
                CachedRowsCount_ = GetRowsCount();
            }
            return *CachedRowsCount_;
        }

    private:
        size_t SeekFat(ui32 rowId) const {
            // Note: When Size() < 16, Fat_ is empty and SeekFat returns 0
            TSeek s{rowId + 1, 0};
            size_t lbFat = Fat_.LowerBound(s); // this is +1
            return lbFat;
        }

        std::pair<ui32, TIo::TSubkey> ReadFirstSubkey(size_t seekResult) const {
            Y_ASSERT(seekResult <= Fat_.Size());
            if (seekResult == 0)
                return {0, TIo::TSubkey()};
            --seekResult;
            ui32 subfatPos = Fat_.ReadKey(seekResult).Position;
            return {subfatPos, Fat_.ReadData(seekResult)};
        }

        TIo::TData ReadData(size_t subfatPos) const {
            return Subfat_.ReadData(subfatPos);
        }

        template <typename TFoo>
        void ReadFrom(size_t fatSeek, TFoo&& consumer) const {
            if (Y_UNLIKELY(IsEmpty()))
                return; // avoids assert in TFlatSearcher

            const size_t end = Fat_.Size() + 1;
            for (size_t seekResult = fatSeek; seekResult != end; ++seekResult) {
                ui32 subfatPos;
                TIo::TSubkey value;
                std::tie(subfatPos, value) = ReadFirstSubkey(seekResult);

                for (TIo::TSubkey delta = Subfat_.ReadKey(subfatPos); !TIo::TSubkeyOps::IsStopMark(delta);) {
                    TIo::TSubkeyOps::Integrate(value, delta, value);
                    bool shouldContinue = consumer(subfatPos, value);
                    if (!shouldContinue)
                        return;
                    delta = Subfat_.ReadKey(++subfatPos);
                }
            }
        }

    private:
        size_t GetRowsCount() const {
            size_t count = 0;
            if (Y_UNLIKELY(Subfat_.Size() == 0))
                return count; // avoids assert in TFlatSearcher

            for (size_t subfatPos = 0; subfatPos < Subfat_.Size(); ++subfatPos) {
                TIo::TSubkey delta = Subfat_.ReadKey(subfatPos);
                if (!TIo::TSubkeyOps::IsStopMark(delta))
                    count++;
            }
            return count;
        }

        void Validate() const {
            Y_ENSURE(Subfat_.Size() < ((1 << TRowfatSeek::SuffixBits) + 1) || Fat_.Size() != 0, "broken serf");
            Y_ENSURE(Subfat_.Size() == 0 || TIo::TSubkeyOps::IsStopMark(Subfat_.ReadKey(Subfat_.Size() - 1)), "broken serf");
        }

    private:
        TFatSearcher Fat_;
        TSubfatSearcher Subfat_;
        mutable TMaybe<size_t> CachedRowsCount_;
    };
}
