#ifndef __CDS_PUB_SUB_PLUGINS_ZLIB_HH__
#define __CDS_PUB_SUB_PLUGINS_ZLIB_HH__

#include <stdexcept>

#include <cds-pubsub/common.hh>
#include <zstd.h>
#include <boost/endian/arithmetic.hpp>

namespace pub_sub
{
    namespace detail
    {
        /*!
         * @brief Simple wrapper to take a decompress a compressed
         * pub_sub::Message
         */
        class InflateStream
        {
        public:
            InflateStream( )
            {
            }

            /*!
             * @brief inflate/decompress a message
             * @param input the compressed message
             * @return an uncompressed copy of the message
             * @note raises a runtime_error on error
             */
            pub_sub::Message
            operator( )( const pub_sub::Message& input )
            {
                if ( input.length < sizeof( boost::endian::little_uint32_t ) )
                {
                    throw std::runtime_error( "ZStd message is too small" );
                }

                auto size_ptr =
                    reinterpret_cast< const boost::endian::little_uint32_t* >(
                        input.data( ) );
                auto uncompressed_size =
                    static_cast< std::size_t >( *size_ptr );

                auto data =
                    pub_sub::DataPtr( new unsigned char[ uncompressed_size ] );

                auto rc = ZSTD_decompress(
                    data.get( ),
                    uncompressed_size,
                    reinterpret_cast< unsigned char* >( input.data( ) ) +
                        sizeof( std::uint32_t ),
                    input.length - sizeof( std::uint32_t ) );
                if ( ZSTD_isError( rc ) )
                {
                    throw std::runtime_error( "Unable to decompress input" );
                }

                return pub_sub::Message( std::move( data ), rc );
            }

            ~InflateStream( ) = default;

        private:
        };

        /*!
         * @brief Simple wrapper to compress a pub_sub::Message
         */
        class DeflateStream
        {
        public:
            DeflateStream( int compression_level = ZSTD_CLEVEL_DEFAULT )
                : compression_level_{ compression_level }
            {
            }

            ~DeflateStream( ) = default;

            /*!
             * @brief deflate/compress a message
             * @param input the uncompressed message
             * @return an compressed copy of the message
             * @note raises a runtime_error on error
             */
            pub_sub::Message
            operator( )( const pub_sub::Message& input )
            {
                auto req_size = ZSTD_compressBound( input.length );
                auto data = pub_sub::DataPtr(
                    new unsigned char[ req_size + sizeof( std::uint32_t ) ] );

                auto* data_start = data.get( ) + sizeof( std::uint32_t );
                auto  rc = ZSTD_compress( data_start,
                                         req_size,
                                         input.data( ),
                                         input.length,
                                         compression_level_ );

                if ( ZSTD_isError( rc ) )
                {
                    throw std::runtime_error( "Unable to compress input" );
                }
                *reinterpret_cast< boost::endian::little_uint32_t* >(
                    data.get( ) ) = input.length;
                auto final_size = rc + sizeof( std::uint32_t );

                return pub_sub::Message( std::move( data ), final_size );
            }

        private:
            int compression_level_;
        };

    } // namespace detail
} // namespace pub_sub
#endif // __CDS_PUB_SUB_PLUGINS_ZLIB_HH__