#pragma once

#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/coro/waked.h>

namespace NSrvKernel {
    class TChannelInput : public IIoInput {
    public:
        TChannelInput(TAtomicSharedPtr<TU2WChannel<TString>> channel, TContExecutor* executor)
            : Channel_(std::move(channel))
            , Executor_(executor)
            , Waker_(executor)
        {}

    private:
        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            if (Eof_) {
                return {};
            }
            TString message;
            EChannelStatus status = Channel_->Receive(message, deadline, Executor_->Running(), &Waker_);
            switch (status) {
                case EChannelStatus::Canceled:
                    return Y_MAKE_ERROR(TSystemError{ECANCELED});
                case EChannelStatus::TimedOut:
                    return Y_MAKE_ERROR(TSystemError{ETIMEDOUT});
                case EChannelStatus::Success:
                    if (message.empty()) {
                        Eof_ = true;
                    }
                    lst = TChunkList(std::move(message));
                    return {};
            }
        }

    private:
        TAtomicSharedPtr<TU2WChannel<TString>> Channel_;
        TContExecutor* Executor_ = nullptr;
        TEventWaker Waker_;
        bool Eof_ = false;
    };

    class TChannelOutput : public IIoOutput {
      public:
        explicit TChannelOutput(const TAtomicSharedPtr<TW2UChannel<TString>>& channel, TContExecutor* executor)
            : Channel_(channel)
            , Executor_(executor)
            , Waker_(executor)
        {}

      private:
        TError DoSend(TChunkList chunkList, TInstant deadline) noexcept override {
            do {
                TString chunk = chunkList.Empty() ? TString{} : TString{chunkList.PopFront()->AsStringBuf()};
                EChannelStatus status = Channel_->Send(std::move(chunk), deadline, Executor_->Running(), &Waker_);
                switch (status) {
                    case EChannelStatus::Canceled:
                        return Y_MAKE_ERROR(TSystemError{ECANCELED});
                    case EChannelStatus::TimedOut:
                        return Y_MAKE_ERROR(TSystemError{ETIMEDOUT});
                    case EChannelStatus::Success:
                        continue;
                }
            } while (!chunkList.Empty());
            return {};
        }

      private:
        TAtomicSharedPtr<TW2UChannel<TString>> Channel_;
        TContExecutor* Executor_ = nullptr;
        TEventWaker Waker_;
    };
}
