#pragma once

#include <util/generic/buffer.h>
#include <util/generic/store_policy.h>
#include <util/generic/string.h>
#include <util/stream/output.h>

class TCappedBufferedOutputBase : public IOutputStream {
public:

public:
    TCappedBufferedOutputBase(IOutputStream* slave, size_t writeCap)
        : WriteCap(writeCap)
        , Slave(slave)
        , Buffer(WriteCap)
    {}
    ~TCappedBufferedOutputBase() {
        DoFlush();
    }

protected:
    void DoWrite(const void* buf, size_t len) override;
    void DoFlush() override;
    void DoFinish() override;

private:
    void DumpBufferToSlave();
    void DumpDataToSlave(const char*& buf, size_t& len);
    void PushDataIntoBuffer(const char*& buf, size_t& len);

    size_t GetBufferAvail() const {
        // Cannot rely on Buffer.Avail() because TBuffer may allocate more memory than requested
        return WriteCap - Buffer.Size();
    }

private:
    const size_t WriteCap;
    IOutputStream* const Slave;
    TBuffer Buffer;
};

template <typename TSlave>
class TCappedBufferedOutput : private TEmbedPolicy<TSlave>, public TCappedBufferedOutputBase {
    using TSlaveBase = TEmbedPolicy<TSlave>;
public:
    template <typename... Args>
    TCappedBufferedOutput(size_t writeCap, Args&&... args)
        : TSlaveBase(std::forward<Args>(args)...)
        , TCappedBufferedOutputBase(TSlaveBase::Ptr(), writeCap)
    {
    }
};
