#pragma once

#include <balancer/serval/contrib/cone/cone.hh>

#include <util/generic/maybe.h>
#include <util/generic/strbuf.h>

#include <memory>

namespace NSv {
    class TBuffer : TNonCopyable {
    public:
        TBuffer(size_t size = 2048) noexcept
            : Size(size)
        {
        }

        // Wait until at least one byte is available or EOF is reached, return the current contents.
        TMaybe<TStringBuf> Read() noexcept {
            while (!Used && State == Open) {
                if (!More.wait() MUN_RETHROW) {
                    return {};
                }
            }
            if ((State == Failed || State == Discarding) && mun_error(ECONNRESET, "buffer aborted")) {
                return {};
            }
            return TStringBuf(Data.get(), Used);
        }

        // Skip `n` previously read bytes from the buffer.
        bool Consume(size_t n) noexcept {
            if (n) {
                Used -= n;
                memmove(Data.get(), Data.get() + n, Used);
                Less.wake();
            }
            return true;
        }

        // Copy data into the buffer. If not enough space is available, copy what fits,
        // sleep until `Read` is called with nonzero `n`, then repeat.
        bool Write(TStringBuf v, bool corked = false) noexcept {
            while (v) {
                while (Used >= Size && State == Open) {
                    if (!Less.wait() MUN_RETHROW) {
                        return false;
                    }
                }
                if (State != Open) {
                    return State == Discarding || !mun_error(EPIPE, "buffer closed");
                }
                if (!Data) {
                    Data.reset(new char[Size]);
                }
                auto part = v.NextTokAt(Size - Used);
                memmove(Data.get() + Used, part.data(), part.size());
                Used += part.size();
                if (!corked || Used >= Size) {
                    More.wake();
                }
            }
            return true;
        }

        // If there is data, wake the reader. Must be used after a corked `Write`
        // if it is known that there will be no more data for a while.
        void Uncork() noexcept {
            if (Used) {
                More.wake();
            }
        }

        // Whether `Read`s is guaranteed to only return empty buffers. (It may still fail.)
        bool AtEnd() const noexcept {
            return !Used && State == Closed;
        }

        // Whether `Write`s could succeed.
        bool IsOpen() const noexcept {
            return State == Open;
        }

        // Make all concurrent and future `Write`s return EPIPE; allow `Read`s to return 0 bytes.
        void Close() noexcept {
            SetState(Closed);
        }

        // Same as `Close`, but `Read` fails with ECONNRESET.
        void Abort() noexcept {
            SetState(Failed);
        }

        // Make further `Write`s no-ops.
        void Discard() noexcept {
            SetState(Discarding);
        }

    private:
        std::unique_ptr<char[]> Data;
        size_t Used = 0;
        size_t Size = 0;
        cone::event More;
        cone::event Less;
        enum EState { Open, Discarding, Closed, Failed } State = Open;

    private:
        void SetState(EState state) noexcept {
            if (State == Closed || State == Failed || State == state)
                return;
            State = state;
            More.wake();
            Less.wake();
            if (state != Closed) {
                Used = 0;
                Data.reset();
            }
        }
    };

    // Concatenate the first n bytes read from a buffer-like object into a string.
    template <typename T>
    static inline TMaybe<TString> ReadFrom(T& stream, size_t n = (size_t)-1) {
        TString data;
        while (auto p = stream.Read()) {
            data += p->Trunc(n - data.size());
            if (p->size() && !stream.Consume(p->size())) {
                return {};
            }
            if (!*p || data.size() == n) {
                return data;
            }
        }
        return {};
    }
}
