#pragma once

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

#include <library/cpp/coroutine/engine/impl.h>

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

namespace NSrvKernel {
    class TChunkQueue {
    public:
        TChunkQueue(TContExecutor* e) noexcept
            : Executor_(e)
        {}

        TContExecutor* Executor() const noexcept {
            return Executor_;
        }

        void Push(TChunkList lst) noexcept {
            if (!lst.Empty() && !AllDone_) {
                Chunks_.Append(std::move(lst));
                Wake(PushWaiter_);
            }
        }

        [[nodiscard]] int Pop(TChunkList& lst, TInstant deadline) noexcept {
            int status = 0;
            if (!AllDone_ && Chunks_.Empty()) {
                status = Wait(&PushWaiter_, deadline);
            }

            lst.Append(std::move(Chunks_));
            Wake(PopWaiter_);
            return status;
        }

        [[nodiscard]] int WaitPop(TInstant deadline) noexcept {
            if (!AllDone_) {
                return Wait(&PopWaiter_, deadline);
            }
            return 0;
        }

        void Stop() noexcept {
            if (!AllDone_) {
                AllDone_ = true;
                Wake(PushWaiter_);
                Wake(PopWaiter_);
            }
        }

        bool Stopped() noexcept {
            return AllDone_;
        }

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

    private:
        int Wait(TCont** cont, TInstant deadline) noexcept {
            Y_ASSERT(!*cont);
            *cont = Executor_->Running();
            const int retval = (*cont)->SleepD(deadline);
            *cont = nullptr;
            return retval;
        }

        void Wake(TCont* cont) const noexcept {
            if (cont) {
                cont->ReSchedule();
                cont = nullptr;
            }
        }

    private:
        TContExecutor* const Executor_ = nullptr;
        TChunkList Chunks_;
        TCont* PushWaiter_ = nullptr;
        TCont* PopWaiter_ = nullptr;
        bool AllDone_ = false;
    };

    class TLimitedChunkQueue {
    public:
        TLimitedChunkQueue(TContExecutor* e, size_t limit) noexcept
            : Queue_(e), Limit_(limit)
        {}

    public:
        void Push(TChunkList lst) noexcept {
            size_t chunkSize = lst.size();
            if (CurrentSize_ + chunkSize > Limit_) {
                Dropped_ += chunkSize;
                return;
            }

            CurrentSize_ += chunkSize;
            Queue_.Push(std::move(lst));
        }

        int Pop(TChunkList& lst, TInstant deadline) noexcept {
            int status = Queue_.Pop(lst, deadline);
            CurrentSize_ = 0;
            return status;
        }

        size_t InQueue() const noexcept {
            return CurrentSize_;
        }

        size_t Dropped() const noexcept {
            return Dropped_;
        }

    private:
        TChunkQueue Queue_;
        size_t CurrentSize_ = 0;
        size_t Dropped_ = 0;
        const size_t Limit_;
    };

    class TQueueIo final : public IIoInput, public IIoOutput {
    public:
        TQueueIo(TContExecutor* e) noexcept
            : Queue_(e)
            , Limit_(0)
        {}

        TQueueIo(TContExecutor* e, size_t l) noexcept
            : Queue_(e)
            , Limit_(l)
        {}

        TQueueIo(TContExecutor* e, TDuration recvOpTimeout) noexcept
            : Queue_(e)
            , RecvOpTimeout_(recvOpTimeout)
            , Limit_(0)
        {}

        void Stop() noexcept {
            Queue_.Stop();
        }

    private:
        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            TInstant effectiveDeadline = deadline;
            if (RecvOpTimeout_ != TDuration::Zero()) {
                effectiveDeadline = Min(deadline, RecvOpTimeout_.ToDeadLine());
            }

            int status = Queue_.Pop(lst, effectiveDeadline);
            if (status == ETIMEDOUT) {
                Queue_.Executor()->Running()->Cancel();
                return Y_MAKE_ERROR(TRecvTimeoutError{} << "io error");
            }

            Y_REQUIRE(!Queue_.Executor()->Running()->Cancelled(),
                      TSystemError{ECANCELED} << "io error: input recv timed out");
            return {};
        }

        TError DoSend(TChunkList lst, TInstant deadline) noexcept override {
            if (lst.Empty()) {
                Queue_.Stop();
            } else if (!Limit_) {
                Queue_.Push(std::move(lst));
            } else {
                while (true) {
                    const size_t left = Limit_ - Min<size_t>(Limit_, Queue_.InQueue());

                    if (Queue_.Stopped() || lst.size() <= left) {
                        Queue_.Push(std::move(lst));
                        break;
                    }

                    TChunkList tmp = CutPrefix(left, lst);
                    Queue_.Push(std::move(tmp));
                    int status = Queue_.WaitPop(deadline);
                    if (status == ETIMEDOUT) {
                        Queue_.Executor()->Running()->Cancel();
                        return Y_MAKE_ERROR(TSystemError{ETIMEDOUT} << "io error");
                    }

                    Y_REQUIRE(!Queue_.Executor()->Running()->Cancelled(),
                              TSystemError{ECANCELED} << "io error");
                }
            }
            return {};
        }

    private:
        TChunkQueue Queue_;
        TDuration RecvOpTimeout_;
        const size_t Limit_;
    };
}
