#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/stream/output.h>
#include <util/stream/walk.h>

#include <utility>

namespace NSrvKernel {
    /* Dummy stream interfaces that is ensures
     * exceptionless write and read in its descendants.
     */
    class IExceptionlessInputStream : public IInputStream {
        size_t DoRead(void* buf, size_t len) noexcept override = 0;
    };

    class IExceptionlessOutputStream : public IOutputStream {
        void DoWrite(const void* data, size_t len) noexcept override = 0;
    };


    class TChunkInputStream : public IExceptionlessInputStream {
    public:
        explicit TChunkInputStream(TChunkList chunkList) noexcept
            : ChunkList_(std::move(chunkList))
        {}

    private:
        size_t DoRead(void* buf, size_t len) noexcept override {
            auto res = DoCutPrefix(len, ChunkList_);
            res.first.CopyDataTo(buf, len);
            return res.second;
        }

    private:
        TChunkList ChunkList_;
    };


    class TChunksOutputStream : public IExceptionlessOutputStream {
    public:
        explicit TChunksOutputStream(size_t chunkSize = 8192) noexcept
            : ChunkSize_(chunkSize)
        {}

        TChunkList& Chunks() noexcept {
            NextChunk(0);
            return Chunks_;
        }

    private:
        void DoWrite(const void* data, size_t len) noexcept override {
            if (Current_.Avail() < len) {
                NextChunk(std::max(len, ChunkSize_));
            }

            memcpy(Current_.Data() + Current_.Size(), data, len);
            Current_.Advance(len);
        }

        void DoFlush() noexcept override {
            NewChunkReserve(ChunkSize_);
        }

        void NextChunk(size_t len) {
            if (Current_) {
                Chunks_.Push(std::move(Current_));
            }
            Current_.Reserve(len);
        }

    private:
        TChunkList Chunks_;
        TBuffer Current_;
        const size_t ChunkSize_;
    };


    class TYieldingOutput : public IExceptionlessOutputStream {
    public:
        TYieldingOutput(IExceptionlessOutputStream& slave, TCont& currentCont, unsigned yieldFrequencyPower,
                        unsigned yieldLimit) noexcept
            : Slave_(slave)
            , Cont_(currentCont)
            , Frequency_((1 << yieldFrequencyPower) - 1)
            , Limit_(yieldLimit)
        {}

    private:
        void DoWrite(const void* data, size_t len) noexcept override {
            if (Limit_ != 0) {
                if ((Current_++ & Frequency_) == 0) {
                    Cont_.Yield();
                    --Limit_;
                }
            }
            return Slave_.Write(data, len);
        }

    private:
        IExceptionlessOutputStream& Slave_;
        TCont& Cont_;
        unsigned Frequency_ = 127;
        unsigned Limit_ = 10;
        unsigned Current_ = 0;
    };
}
