#pragma once

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

#include <util/generic/yexception.h>

#include <utility>


namespace NSrvKernel {
    //TODO - optimize some cases
    class TStreamIterator {
    public:
        // To handle events which look like EOFs in the
        enum EEmptyReadAction {
            Stop, Restart
        };

        TStreamIterator(IIoInput* slave) noexcept
            : S_(slave)
        {}

        explicit TStreamIterator(TChunkList&& data) noexcept
            : L_(std::move(data))
        {
        }

        TStreamIterator(TStreamIterator&& rhs) noexcept
            : S_(rhs.S_)
            , L_(std::move(rhs.L_))
        {
            rhs.S_ = nullptr;
        }

        TStreamIterator& operator=(TStreamIterator&& rhs) noexcept {
            TStreamIterator tmp{ std::move(rhs) };
            DoSwap(*this, tmp);
            return *this;
        }

        virtual ~TStreamIterator() noexcept {
            if (S_) {
                S_->UnRecv(std::move(L_));
            }
        }

        TErrorOr<size_t> ReadSome(
            void* buf, size_t len, TInstant deadline = TInstant::Max()) noexcept
        {
            TChunkList tmp;

            Y_PROPAGATE_ERROR(Next(len, tmp, deadline));

            return tmp.CopyDataTo(buf, len);
        }

        TError Read(
            void* bufin, size_t len, TInstant deadline = TInstant::Max()) noexcept
        {
            const auto totalLen = len;
            char* buf = reinterpret_cast<char*>(bufin);

            while (len) {
                size_t readed;
                Y_PROPAGATE_ERROR(ReadSome(buf, len, deadline).AssignTo(readed));
                if (readed) {
                    buf += readed;
                    len -= readed;
                } else {
                    Y_REQUIRE(EEmptyReadAction::Stop != OnEmptyRead(),
                              TStreamError(len, totalLen));
                }
            }
            return {};
        }

        TError Read(
            TChunkList& lst, size_t len, TInstant deadline = TInstant::Max()) noexcept
        {
            const auto totalLen = len;

            while (len) {
                TChunkList tmp;

                Y_PROPAGATE_ERROR(Next(len, tmp, deadline));

                if (size_t readed = tmp.size()) {
                    len -= readed;
                    lst.Append(std::move(tmp));
                } else {
                    Y_REQUIRE(EEmptyReadAction::Stop != OnEmptyRead(),
                              TStreamError(len, totalLen));
                }
            }
            return {};
        }

        TErrorOr<bool> Next(char& ch, TInstant deadline = TInstant::Max()) noexcept {
            size_t readed;
            Y_PROPAGATE_ERROR(ReadSome(&ch, 1, deadline).AssignTo(readed));
            return readed > 0;
        }

        TErrorOr<char> Next(TInstant deadline = TInstant::Max()) noexcept {
            char ch;

            bool got;
            Y_PROPAGATE_ERROR(Next(ch, deadline).AssignTo(got));
            if (!got) {
                Y_REQUIRE(EEmptyReadAction::Stop != OnEmptyRead(),
                          TStreamError(1, 1));
            }

            return ch;
        }

        TError Next(
            size_t len, TChunkList& lst, TInstant deadline = TInstant::Max()) noexcept
        {
            bool got;
            Y_PROPAGATE_ERROR(Peek(deadline).AssignTo(got));
            if (got) {
                lst = CutPrefix(len, L_);
            }
            return {};
        }

        TError Next(TChunkList& lst, TInstant deadline = TInstant::Max()) noexcept {
            bool got;
            Y_PROPAGATE_ERROR(Peek(deadline).AssignTo(got));
            if (got) {
                lst.Append(std::move(L_));
            }
            return {};
        }

        TErrorOr<bool> Peek(TInstant deadline = TInstant::Max()) noexcept {
            if (S_ && L_.Empty()) {
                do {
                    Y_PROPAGATE_ERROR(S_->Recv(L_, deadline));
                } while (L_.Empty() && EEmptyReadAction::Stop != OnEmptyRead());
            }

            return !L_.Empty();
        }

        void UnGet(TChunkList& data) noexcept {
            L_.Prepend(std::move(data));
        }

        bool HasBufferedData() const noexcept {
            return !L_.Empty() || (S_ && S_->HasBufferedData());
        }

        size_t BytesBuffered() const noexcept {
            return L_.size() + (S_ ? S_->BytesBuffered() : 0);
        }

    private:
        virtual EEmptyReadAction OnEmptyRead() noexcept {
            return EEmptyReadAction::Stop;
        }

    private:
        IIoInput* S_ = nullptr;
        TChunkList L_;
    };
}
