#pragma once

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

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

#include <array>

namespace NSrvKernel {
    class TDoubleInputMaster : TMoveOnly {
        friend class TDoubleInputSlave;
    public:
        explicit TDoubleInputMaster(IIoInput& slave, size_t bufferLimit)
            : UnderlyingInput_(slave)
            , BufferLimit_(bufferLimit)
        {
            StoredLength_.fill(0);
            Exceeded_.fill(false);
        }

        TError RecvMaster(TInstant deadline, size_t exclusiveIndex) {
            // Slave with exclusive index is allowed to exceed the limit
            // as it will immediately recv all stored data

            TChunkList lst;

            Y_PROPAGATE_ERROR(UnderlyingInput_.Recv(lst, deadline));

            if (lst.Empty()) {
                return {};
            }

            size_t currentLength = lst.size();

            for (size_t index = 0; index < 2; ++index) {
                if (!Exceeded_[index]) {
                    // If slave has ever exceeded the limit, he has skipped a part of input
                    // and thus shouldn't receive input anymore
                    if (StoredLength_[index] <= BufferLimit_ || index == exclusiveIndex) {
                        Stored_[index].Append(lst.Copy());
                        StoredLength_[index] += currentLength;
                    } else {
                        Exceeded_[index] = true;
                    }
                }
            }

            return {};
        }

        TError RecvSlave(TChunkList& lst, TInstant deadline, size_t index) {
            Y_PROPAGATE_ERROR(RecvMaster(deadline, index));

            lst = std::move(Stored_[index]);
            StoredLength_[index] = 0;

            return {};
        }

        void ResetRewind() noexcept {
            StoredLength_.fill(0);
            for (TChunkList& list : Stored_) {
                list.Clear();
            }
            BufferLimit_ = 0;
        }

    private:
        IIoInput& UnderlyingInput_;
        std::array<TChunkList, 2> Stored_;
        std::array<size_t, 2> StoredLength_;
        std::array<bool, 2> Exceeded_;
        size_t BufferLimit_ = 0;
    };

    class TDoubleInputSlave : public IIoInput {
    public:
        TDoubleInputSlave(TDoubleInputMaster& master, size_t index)
            : Master_(master)
            , Index_(index)
        {}

        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            return Master_.RecvSlave(lst, deadline, Index_);
        }

    private:
        TDoubleInputMaster& Master_;
        size_t Index_ = 0;
    };
}
