#pragma once

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

namespace NSrvKernel {
    class TLengthLimitInput : public IIoInput {
    public:
        TLengthLimitInput(IIoInput* s, ui64 limit) noexcept
            : S_(s)
            , L_(limit)
        {}

        bool Depleted() const noexcept {
            return !L_;
        }

        void SetLimit(ui64 limit) noexcept {
            L_ = limit;
        }

        ui64 GetLimit() const noexcept {
            return L_;
        }

    private:
        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            if (L_) {
                Y_PROPAGATE_ERROR(S_->Recv(lst, deadline));
                size_t readed = lst.size();
                if (readed <= L_) {
                    L_ -= readed;
                } else {
                    TChunkList tmp = CutPrefix(L_, lst);
                    tmp.Swap(lst);
                    S_->UnRecv(std::move(tmp));
                    L_ = 0;
                }
            }
            return {};
        }

    private:
        IIoInput* const S_ = nullptr;
        ui64 L_ = 0;
    };

    class TLengthEqualInput : public IIoInput {
    public:
        TLengthEqualInput(IIoInput* s, ui64 limit) noexcept
            : S_(s, limit)
        {}

    private:
        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            Y_PROPAGATE_ERROR(S_.Recv(lst, deadline));

            Y_REQUIRE(!lst.Empty() || S_.Depleted(),
                      (THttpError{400, "body length mismatch"}));
            return {};
        }

    private:
        TLengthLimitInput S_;
    };
}
