#pragma once

#include "socket_adapter.h"
#include <google/protobuf/messagext.h>
#include <google/protobuf/io/coded_stream.h>
#include <util/generic/ptr.h>
#include <util/generic/yexception.h>

class TSocketAdapterZeroCopyStream: public ISocketAdapter {
public:
    TSocketAdapterZeroCopyStream(IInputStream& input, IOutputStream& output)
        : ISocketAdapter(input, output) {
    }

    virtual bool ReadMessage(::google::protobuf::Message& message) const override {
        if (!InputAdapter) {
            InputAdapter = new ::google::protobuf::io::TCopyingInputStreamAdaptor(&Input);
            Decoder = new ::google::protobuf::io::CodedInputStream(InputAdapter.Get());
        }
        return ParseFromCodedStreamSeq(&message, Decoder.Get());
    }

    using ISocketAdapter::WriteMessage;

    virtual void WriteMessage(const ::google::protobuf::Message& message, const TString& request="") const override {
        Y_UNUSED(request);
        ::google::protobuf::io::TCopyingOutputStreamAdaptor OutputAdapter(&Output);
        if (!SerializeToZeroCopyStreamSeq(&message, &OutputAdapter)) {
            ythrow yexception() << "Failed to serialize message";
        }
        if (!OutputAdapter.Flush()) {
            ythrow yexception() << "Failed to send message";
        }
    }

    virtual ~TSocketAdapterZeroCopyStream() override {}

    inline bool HasError() const {
        return !!InputAdapter && InputAdapter->HasError();
    }

protected:
private:
    mutable TAutoPtr< ::google::protobuf::io::TCopyingInputStreamAdaptor > InputAdapter;
    mutable TAutoPtr< ::google::protobuf::io::CodedInputStream > Decoder;
};
