#pragma once

#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/memory/chunks.h>

#include <util/datetime/base.h>
#include <util/generic/strbuf.h>

namespace NSrvKernel {
    class TLimitedRewindableInput : public IIoInput {
    public:
        TLimitedRewindableInput(IIoInput& slave, size_t limit) noexcept
            : Slave_(&slave)
            , Limit_(limit)
        {}

        void SetPrefix(TChunkList chunkList) noexcept {
            Prefix_ = std::move(chunkList);
        }

        void Rewind() noexcept {
            if (!Prefix_.Empty()) {
                Chunks_ = Prefix_.Copy();
                Chunks_.Append(Stored_.Copy());
            } else {
                Chunks_ = Stored_.Copy();
            }
        }

        [[nodiscard]] size_t StoredSize() const noexcept {
            return StoredSize_;
        }

        [[nodiscard]] const TChunkList& Stored() const noexcept {
            return Stored_;
        }

        [[nodiscard]] bool IsRewindable() const noexcept {
            return IsRewindable_;
        }

        void ResetRewind() noexcept {
            if (IsRewindable_) {
                IsRewindable_ = false;
                StoredSize_ = 0;
                Stored_.Clear();
            }
        }

        size_t Limit() const { return Limit_; }

    protected:
        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            if (Ended_) {
                return {};
            }

            if (IsRewindable_) {
                if (auto err = Slave_->Recv(lst, deadline)) {
                    if (!lst.Empty()) {
                        // Should never happen
                        ResetRewind();
                    }

                    return err;
                }

                if (lst.Empty()) {
                    Ended_ = true;
                    return {};
                }

                StoredSize_ += lst.size();

                if (StoredSize_ > Limit_) {
                    ResetRewind();
                    return {};
                }

                Stored_.Append(lst.Copy());

                return {};
            } else {
                return Slave_->Recv(lst, deadline);
            }
        }

    private:
        TChunkList Stored_;
        TChunkList Prefix_;
        IIoInput* const Slave_ = nullptr;
        const size_t Limit_;
        size_t StoredSize_ = 0;
        bool IsRewindable_ = true;
        bool Ended_ = false;
    };
}
