#pragma once

#include <balancer/kernel/memory/chunks.h>
#include <balancer/kernel/helpers/errors.h>

namespace NSrvKernel {

    enum class EStreamState {
        Open,
        Close,
        Error
    };

    class IIoInput : public TMoveOnly {
    public:
        virtual ~IIoInput() noexcept = default;

        TError Recv(TChunkList& lst, TInstant deadline) noexcept {
            if (Chunks_.Empty()) {
                lst.Clear();
                return DoRecv(lst, deadline);
            } else {
                lst = std::move(Chunks_);
            }
            return {};
        }

        bool HasBufferedData() const noexcept {
            return !Chunks_.Empty();
        }

        size_t BytesBuffered() const noexcept {
            return Chunks_.size();
        }

        [[nodiscard]] TChunkList RecvBuffered() noexcept {
            return std::move(Chunks_);
        }

        TErrorOr<bool> Peek(TInstant deadline) noexcept {
            TChunkList lst;

            Y_PROPAGATE_ERROR(Recv(lst, deadline));

            if (!lst.Empty()) {
                Chunks_.Append(std::move(lst));
                return true;
            }

            return false;
        }

        bool TryPeek(TInstant deadline) noexcept {
            bool ret = false;
            Y_TRY(TError, error) {
                return Peek(deadline).AssignTo(ret);
            }Y_CATCH {
                return false;
            }
            return ret;
        }

        virtual void UnRecv(TChunkList&& lst) noexcept {
            Chunks_.Prepend(std::move(lst));
        }

        TErrorOr<size_t> FillBuffer(
            size_t maxLen, TInstant deadline = TInstant::Max()) noexcept {
            auto currLen = Chunks_.size();
            while (currLen < maxLen) {
                TChunkList lst;
                Y_PROPAGATE_ERROR(DoRecv(lst, deadline));
                if (lst.Empty()) {
                    break;
                }
                currLen += lst.size();
                Chunks_.Append(std::move(lst));
            }
            return currLen;
        }

    private:
        virtual TError DoRecv(TChunkList& lst, TInstant deadline) noexcept = 0;

    protected:
        TChunkList Chunks_;
        EStreamState StreamState_ = EStreamState::Open;
    };


    class IIoOutput : public TMoveOnly {
    public:
        virtual ~IIoOutput() noexcept = default;

        TError Send(TChunkList lst, TInstant deadline) noexcept {
            if (lst.Empty()) {
                StreamState_ = EStreamState::Close;
            }
            auto error = DoSend(std::move(lst), deadline);
            if (error) {
                StreamState_ = EStreamState::Error;
            }
            return error;
        }

        TError SendEof(TInstant deadline) noexcept {
            return Send({}, deadline);
        }

        EStreamState StreamState() const noexcept {
            return StreamState_;
        }


//        TError Send(TChunkList& lst, TInstant deadline) {
//            TInstant old = SetDeadline(deadline);
//            Y_DEFER {
//                Y_UNUSED(SetDeadline(old));
//            };
//            return Send(lst);
//        }
//
    private:
        virtual TError DoSend(TChunkList lst, TInstant deadline) noexcept = 0;
//
//        // TODO(velavokr): need to implement passing in every delegating TIoOutput
//        [[nodiscard]] virtual TInstant SetDeadline(TInstant inst) noexcept {
//            return inst;
//        }

    private:
        EStreamState StreamState_ = EStreamState::Open;
    };

    class TResponse;
    class THeaders;

    class IHttpOutput : public IIoOutput {
    public:
        TError SendHead(TResponse&& response, const bool forceClose, TInstant deadline) {
            return DoSendHead(std::move(response), forceClose, deadline);
        }

        TError SendTrailers(THeaders&& headers, TInstant deadline) {
            return DoSendTrailers(std::move(headers), deadline);
        }

    private:
        virtual TError DoSendHead(TResponse&& response, const bool forceClose, TInstant deadline) = 0;
        virtual TError DoSendTrailers(THeaders&& headers, TInstant deadline) = 0;
    };


    template <class HeadOp, class BodyOp, class TrailersOp>
    inline constexpr bool IsHttpOutput =
        std::is_invocable_r_v<TError, HeadOp, TResponse&&, const bool, TInstant> &&
        std::is_invocable_r_v<TError, BodyOp, TChunkList, TInstant> &&
        std::is_invocable_r_v<TError, TrailersOp, THeaders&&, TInstant>;


    template <class HeadOp, class BodyOp, class TrailersOp>
    class TAnyHttpOutput final : public IHttpOutput {
    public:
        template <class H, class B, class T, class = std::enable_if_t<IsHttpOutput<HeadOp, BodyOp, TrailersOp>>>
        TAnyHttpOutput(H&& head, B&& body, T&& trailers)
            : HeadOp_(std::forward<H>(head))
            , BodyOp_(std::forward<B>(body))
            , TrailersOp_(std::forward<T>(trailers))
        {}

    private:
        TError DoSendHead(TResponse&& response, const bool forceClose, TInstant deadline) override {
            return HeadOp_(std::move(response), forceClose, deadline);
        }

        TError DoSend(TChunkList lst, TInstant deadline) noexcept override {
            return BodyOp_(std::move(lst), deadline);
        }

        TError DoSendTrailers (THeaders&& headers, TInstant deadline) override {
            return TrailersOp_(std::move(headers), deadline);
        }

      private:
        HeadOp HeadOp_;
        BodyOp BodyOp_;
        TrailersOp TrailersOp_;
    };

    template <class HeadOp, class BodyOp, class TrailersOp, class = std::enable_if_t<IsHttpOutput<HeadOp, BodyOp, TrailersOp>>>
    auto MakeHttpOutput(HeadOp&& head, BodyOp&& body, TrailersOp&& trailers) {
        return TAnyHttpOutput<std::decay_t<HeadOp>, std::decay_t<BodyOp>, std::decay_t<TrailersOp>>{
            std::forward<HeadOp>(head),
            std::forward<BodyOp>(body),
            std::forward<TrailersOp>(trailers)
        };
    }

    inline TError Transfer(IIoInput* in, IIoOutput* out, TInstant deadline) noexcept {
        bool finished = false;
        do {
            TChunkList lst;
            Y_PROPAGATE_ERROR(in->Recv(lst, deadline));

            finished = lst.Empty();
            Y_PROPAGATE_ERROR(out->Send(std::move(lst), deadline));
        } while (!finished);
        return {};
    }

    inline TError RecvAll(IIoInput* in, TChunkList& res, TInstant deadline) noexcept {
        while (true) {
            TChunkList lst;
            Y_PROPAGATE_ERROR(in->Recv(lst, deadline));

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

            res.Append(std::move(lst));
        }
        return {};
    }

    inline TError SkipAll(IIoInput* in, TInstant deadline) noexcept {
        while (true) {
            TChunkList lst;
            Y_PROPAGATE_ERROR(in->Recv(lst, deadline));

            if (lst.Empty()) {
                break;
            }
        }
        return {};
    }
}  // namespace NSrvKernel
