//
// Copyright (c) 2016-2019 Vinnie Falco (vinnie dot falco at gmail dot com)
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
// Official repository: https://github.com/boostorg/beast
//

#ifndef BOOST_BEAST_WEBSOCKET_IMPL_WRITE_HPP
#define BOOST_BEAST_WEBSOCKET_IMPL_WRITE_HPP

#include <boost/beast/websocket/detail/mask.hpp>
#include <boost/beast/core/async_base.hpp>
#include <boost/beast/core/bind_handler.hpp>
#include <boost/beast/core/buffer_traits.hpp>
#include <boost/beast/core/buffers_cat.hpp>
#include <boost/beast/core/buffers_prefix.hpp>
#include <boost/beast/core/buffers_range.hpp>
#include <boost/beast/core/buffers_suffix.hpp>
#include <boost/beast/core/flat_static_buffer.hpp>
#include <boost/beast/core/stream_traits.hpp>
#include <boost/beast/core/detail/bind_continuation.hpp>
#include <boost/beast/core/detail/clamp.hpp>
#include <boost/beast/core/detail/config.hpp>
#include <boost/beast/websocket/detail/frame.hpp>
#include <boost/beast/websocket/impl/stream_impl.hpp>
#include <boost/asio/coroutine.hpp>
#include <boost/assert.hpp>
#include <boost/config.hpp>
#include <boost/throw_exception.hpp>
#include <algorithm>
#include <memory>

namespace boost {
namespace beast {
namespace websocket {

template<class NextLayer, bool deflateSupported>
template<class Handler, class Buffers>
class stream<NextLayer, deflateSupported>::write_some_op
    : public beast::async_base<
        Handler, beast::executor_type<stream>>
    , public asio::coroutine
{
    enum
    {
        do_nomask_nofrag,
        do_nomask_frag,
        do_mask_nofrag,
        do_mask_frag,
        do_deflate
    };

    boost::weak_ptr<impl_type> wp_;
    buffers_suffix<Buffers> cb_;
    detail::frame_header fh_;
    detail::prepared_key key_;
    std::size_t bytes_transferred_ = 0;
    std::size_t remain_;
    std::size_t in_;
    int how_;
    bool fin_;
    bool more_ = false; // for ubsan
    bool cont_ = false;

public:
    static constexpr int id = 2; // for soft_mutex

    template<class Handler_>
    write_some_op(
        Handler_&& h,
        boost::shared_ptr<impl_type> const& sp,
        bool fin,
        Buffers const& bs)
        : beast::async_base<Handler,
            beast::executor_type<stream>>(
                std::forward<Handler_>(h),
                    sp->stream().get_executor())
        , wp_(sp)
        , cb_(bs)
        , fin_(fin)
    {
        auto& impl = *sp;

        // Set up the outgoing frame header
        if(! impl.wr_cont)
        {
            impl.begin_msg();
            fh_.rsv1 = impl.wr_compress;
        }
        else
        {
            fh_.rsv1 = false;
        }
        fh_.rsv2 = false;
        fh_.rsv3 = false;
        fh_.op = impl.wr_cont ?
            detail::opcode::cont : impl.wr_opcode;
        fh_.mask =
            impl.role == role_type::client;

        // Choose a write algorithm
        if(impl.wr_compress)
        {
            how_ = do_deflate;
        }
        else if(! fh_.mask)
        {
            if(! impl.wr_frag)
            {
                how_ = do_nomask_nofrag;
            }
            else
            {
                BOOST_ASSERT(impl.wr_buf_size != 0);
                remain_ = buffer_bytes(cb_);
                if(remain_ > impl.wr_buf_size)
                    how_ = do_nomask_frag;
                else
                    how_ = do_nomask_nofrag;
            }
        }
        else
        {
            if(! impl.wr_frag)
            {
                how_ = do_mask_nofrag;
            }
            else
            {
                BOOST_ASSERT(impl.wr_buf_size != 0);
                remain_ = buffer_bytes(cb_);
                if(remain_ > impl.wr_buf_size)
                    how_ = do_mask_frag;
                else
                    how_ = do_mask_nofrag;
            }
        }
        (*this)({}, 0, false);
    }

    void operator()(
        error_code ec = {},
        std::size_t bytes_transferred = 0,
        bool cont = true);
};

template<class NextLayer, bool deflateSupported>
template<class Buffers, class Handler>
void
stream<NextLayer, deflateSupported>::
write_some_op<Buffers, Handler>::
operator()(
    error_code ec,
    std::size_t bytes_transferred,
    bool cont)
{
    using beast::detail::clamp;
    std::size_t n;
    net::mutable_buffer b;
    auto sp = wp_.lock();
    if(! sp)
    {
        ec = net::error::operation_aborted;
        bytes_transferred_ = 0;
        return this->complete(cont, ec, bytes_transferred_);
    }
    auto& impl = *sp;
    BOOST_ASIO_CORO_REENTER(*this)
    {
        // Acquire the write lock
        if(! impl.wr_block.try_lock(this))
        {
        do_suspend:
            BOOST_ASIO_CORO_YIELD
            impl.op_wr.emplace(std::move(*this));
            impl.wr_block.lock(this);
            BOOST_ASIO_CORO_YIELD
            net::post(std::move(*this));
            BOOST_ASSERT(impl.wr_block.is_locked(this));
        }
        if(impl.check_stop_now(ec))
            goto upcall;

        //------------------------------------------------------------------

        if(how_ == do_nomask_nofrag)
        {
            // send a single frame
            fh_.fin = fin_;
            fh_.len = buffer_bytes(cb_);
            impl.wr_fb.clear();
            detail::write<flat_static_buffer_base>(
                impl.wr_fb, fh_);
            impl.wr_cont = ! fin_;
            BOOST_ASIO_CORO_YIELD
            net::async_write(impl.stream(),
                buffers_cat(impl.wr_fb.data(), cb_),
                    beast::detail::bind_continuation(std::move(*this)));
            bytes_transferred_ += clamp(fh_.len);
            if(impl.check_stop_now(ec))
                goto upcall;
            goto upcall;
        }

        //------------------------------------------------------------------

        if(how_ == do_nomask_frag)
        {
            // send multiple frames
            for(;;)
            {
                n = clamp(remain_, impl.wr_buf_size);
                fh_.len = n;
                remain_ -= n;
                fh_.fin = fin_ ? remain_ == 0 : false;
                impl.wr_fb.clear();
                detail::write<flat_static_buffer_base>(
                    impl.wr_fb, fh_);
                impl.wr_cont = ! fin_;
                // Send frame
                BOOST_ASIO_CORO_YIELD
                net::async_write(impl.stream(), buffers_cat(
                    impl.wr_fb.data(),
                    buffers_prefix(clamp(fh_.len), cb_)),
                        beast::detail::bind_continuation(std::move(*this)));
                n = clamp(fh_.len); // restore `n` on yield
                bytes_transferred_ += n;
                if(impl.check_stop_now(ec))
                    goto upcall;
                if(remain_ == 0)
                    break;
                cb_.consume(n);
                fh_.op = detail::opcode::cont;

                // Give up the write lock in between each frame
                // so that outgoing control frames might be sent.
                impl.wr_block.unlock(this);
                if( impl.op_close.maybe_invoke()
                    || impl.op_idle_ping.maybe_invoke()
                    || impl.op_rd.maybe_invoke()
                    || impl.op_ping.maybe_invoke())
                {
                    BOOST_ASSERT(impl.wr_block.is_locked());
                    goto do_suspend;
                }
                impl.wr_block.lock(this);
            }
            goto upcall;
        }

        //------------------------------------------------------------------

        if(how_ == do_mask_nofrag)
        {
            // send a single frame using multiple writes
            remain_ = beast::buffer_bytes(cb_);
            fh_.fin = fin_;
            fh_.len = remain_;
            fh_.key = impl.create_mask();
            detail::prepare_key(key_, fh_.key);
            impl.wr_fb.clear();
            detail::write<flat_static_buffer_base>(
                impl.wr_fb, fh_);
            n = clamp(remain_, impl.wr_buf_size);
            net::buffer_copy(net::buffer(
                impl.wr_buf.get(), n), cb_);
            detail::mask_inplace(net::buffer(
                impl.wr_buf.get(), n), key_);
            remain_ -= n;
            impl.wr_cont = ! fin_;
            // write frame header and some payload
            BOOST_ASIO_CORO_YIELD
            net::async_write(impl.stream(), buffers_cat(
                impl.wr_fb.data(),
                net::buffer(impl.wr_buf.get(), n)),
                    beast::detail::bind_continuation(std::move(*this)));
            // VFALCO What about consuming the buffer on error?
            bytes_transferred_ +=
                bytes_transferred - impl.wr_fb.size();
            if(impl.check_stop_now(ec))
                goto upcall;
            while(remain_ > 0)
            {
                cb_.consume(impl.wr_buf_size);
                n = clamp(remain_, impl.wr_buf_size);
                net::buffer_copy(net::buffer(
                    impl.wr_buf.get(), n), cb_);
                detail::mask_inplace(net::buffer(
                    impl.wr_buf.get(), n), key_);
                remain_ -= n;
                // write more payload
                BOOST_ASIO_CORO_YIELD
                net::async_write(impl.stream(),
                    net::buffer(impl.wr_buf.get(), n),
                        beast::detail::bind_continuation(std::move(*this)));
                bytes_transferred_ += bytes_transferred;
                if(impl.check_stop_now(ec))
                    goto upcall;
            }
            goto upcall;
        }

        //------------------------------------------------------------------

        if(how_ == do_mask_frag)
        {
            // send multiple frames
            for(;;)
            {
                n = clamp(remain_, impl.wr_buf_size);
                remain_ -= n;
                fh_.len = n;
                fh_.key = impl.create_mask();
                fh_.fin = fin_ ? remain_ == 0 : false;
                detail::prepare_key(key_, fh_.key);
                net::buffer_copy(net::buffer(
                    impl.wr_buf.get(), n), cb_);
                detail::mask_inplace(net::buffer(
                    impl.wr_buf.get(), n), key_);
                impl.wr_fb.clear();
                detail::write<flat_static_buffer_base>(
                    impl.wr_fb, fh_);
                impl.wr_cont = ! fin_;
                // Send frame
                BOOST_ASIO_CORO_YIELD
                net::async_write(impl.stream(), buffers_cat(
                    impl.wr_fb.data(),
                    net::buffer(impl.wr_buf.get(), n)),
                        beast::detail::bind_continuation(std::move(*this)));
                n = bytes_transferred - impl.wr_fb.size();
                bytes_transferred_ += n;
                if(impl.check_stop_now(ec))
                    goto upcall;
                if(remain_ == 0)
                    break;
                cb_.consume(n);
                fh_.op = detail::opcode::cont;
                // Give up the write lock in between each frame
                // so that outgoing control frames might be sent.
                impl.wr_block.unlock(this);
                if( impl.op_close.maybe_invoke()
                    || impl.op_idle_ping.maybe_invoke()
                    || impl.op_rd.maybe_invoke()
                    || impl.op_ping.maybe_invoke())
                {
                    BOOST_ASSERT(impl.wr_block.is_locked());
                    goto do_suspend;
                }
                impl.wr_block.lock(this);
            }
            goto upcall;
        }

        //------------------------------------------------------------------

        if(how_ == do_deflate)
        {
            // send compressed frames
            for(;;)
            {
                b = net::buffer(impl.wr_buf.get(),
                    impl.wr_buf_size);
                more_ = impl.deflate(b, cb_, fin_, in_, ec);
                if(impl.check_stop_now(ec))
                    goto upcall;
                n = buffer_bytes(b);
                if(n == 0)
                {
                    // The input was consumed, but there is
                    // no output due to compression latency.
                    BOOST_ASSERT(! fin_);
                    BOOST_ASSERT(buffer_bytes(cb_) == 0);
                    goto upcall;
                }
                if(fh_.mask)
                {
                    fh_.key = impl.create_mask();
                    detail::prepared_key key;
                    detail::prepare_key(key, fh_.key);
                    detail::mask_inplace(b, key);
                }
                fh_.fin = ! more_;
                fh_.len = n;
                impl.wr_fb.clear();
                detail::write<
                    flat_static_buffer_base>(impl.wr_fb, fh_);
                impl.wr_cont = ! fin_;
                // Send frame
                BOOST_ASIO_CORO_YIELD
                net::async_write(impl.stream(), buffers_cat(
                    impl.wr_fb.data(), b),
                        beast::detail::bind_continuation(std::move(*this)));
                bytes_transferred_ += in_;
                if(impl.check_stop_now(ec))
                    goto upcall;
                if(more_)
                {
                    fh_.op = detail::opcode::cont;
                    fh_.rsv1 = false;
                    // Give up the write lock in between each frame
                    // so that outgoing control frames might be sent.
                    impl.wr_block.unlock(this);
                    if( impl.op_close.maybe_invoke()
                        || impl.op_idle_ping.maybe_invoke()
                        || impl.op_rd.maybe_invoke()
                        || impl.op_ping.maybe_invoke())
                    {
                        BOOST_ASSERT(impl.wr_block.is_locked());
                        goto do_suspend;
                    }
                    impl.wr_block.lock(this);
                }
                else
                {
                    if(fh_.fin)
                        impl.do_context_takeover_write(impl.role);
                    goto upcall;
                }
            }
        }

    //--------------------------------------------------------------------------

    upcall:
        impl.wr_block.unlock(this);
        impl.op_close.maybe_invoke()
            || impl.op_idle_ping.maybe_invoke()
            || impl.op_rd.maybe_invoke()
            || impl.op_ping.maybe_invoke();
        this->complete(cont, ec, bytes_transferred_);
    }
}

template<class NextLayer, bool deflateSupported>
struct stream<NextLayer, deflateSupported>::
    run_write_some_op
{
    template<
        class WriteHandler,
        class ConstBufferSequence>
    void
    operator()(
        WriteHandler&& h,
        boost::shared_ptr<impl_type> const& sp,
        bool fin,
        ConstBufferSequence const& b)
    {
        // If you get an error on the following line it means
        // that your handler does not meet the documented type
        // requirements for the handler.

        static_assert(
            beast::detail::is_invocable<WriteHandler,
                void(error_code, std::size_t)>::value,
            "WriteHandler type requirements not met");

        write_some_op<
            typename std::decay<WriteHandler>::type,
            ConstBufferSequence>(
                std::forward<WriteHandler>(h),
                sp,
                fin,
                b);
    }
};

//------------------------------------------------------------------------------

template<class NextLayer, bool deflateSupported>
template<class ConstBufferSequence>
std::size_t
stream<NextLayer, deflateSupported>::
write_some(bool fin, ConstBufferSequence const& buffers)
{
    static_assert(is_sync_stream<next_layer_type>::value,
        "SyncStream type requirements not met");
    static_assert(net::is_const_buffer_sequence<
        ConstBufferSequence>::value,
            "ConstBufferSequence type requirements not met");
    error_code ec;
    auto const bytes_transferred =
        write_some(fin, buffers, ec);
    if(ec)
        BOOST_THROW_EXCEPTION(system_error{ec});
    return bytes_transferred;
}

template<class NextLayer, bool deflateSupported>
template<class ConstBufferSequence>
std::size_t
stream<NextLayer, deflateSupported>::
write_some(bool fin,
    ConstBufferSequence const& buffers, error_code& ec)
{
    static_assert(is_sync_stream<next_layer_type>::value,
        "SyncStream type requirements not met");
    static_assert(net::is_const_buffer_sequence<
        ConstBufferSequence>::value,
            "ConstBufferSequence type requirements not met");
    using beast::detail::clamp;
    auto& impl = *impl_;
    std::size_t bytes_transferred = 0;
    ec = {};
    if(impl.check_stop_now(ec))
        return bytes_transferred;
    detail::frame_header fh;
    if(! impl.wr_cont)
    {
        impl.begin_msg();
        fh.rsv1 = impl.wr_compress;
    }
    else
    {
        fh.rsv1 = false;
    }
    fh.rsv2 = false;
    fh.rsv3 = false;
    fh.op = impl.wr_cont ?
        detail::opcode::cont : impl.wr_opcode;
    fh.mask = impl.role == role_type::client;
    auto remain = buffer_bytes(buffers);
    if(impl.wr_compress)
    {

        buffers_suffix<
            ConstBufferSequence> cb(buffers);
        for(;;)
        {
            auto b = net::buffer(
                impl.wr_buf.get(), impl.wr_buf_size);
            auto const more = impl.deflate(
                b, cb, fin, bytes_transferred, ec);
            if(impl.check_stop_now(ec))
                return bytes_transferred;
            auto const n = buffer_bytes(b);
            if(n == 0)
            {
                // The input was consumed, but there
                // is no output due to compression
                // latency.
                BOOST_ASSERT(! fin);
                BOOST_ASSERT(buffer_bytes(cb) == 0);
                fh.fin = false;
                break;
            }
            if(fh.mask)
            {
                fh.key = this->impl_->create_mask();
                detail::prepared_key key;
                detail::prepare_key(key, fh.key);
                detail::mask_inplace(b, key);
            }
            fh.fin = ! more;
            fh.len = n;
            detail::fh_buffer fh_buf;
            detail::write<
                flat_static_buffer_base>(fh_buf, fh);
            impl.wr_cont = ! fin;
            net::write(impl.stream(),
                buffers_cat(fh_buf.data(), b), ec);
            if(impl.check_stop_now(ec))
                return bytes_transferred;
            if(! more)
                break;
            fh.op = detail::opcode::cont;
            fh.rsv1 = false;
        }
        if(fh.fin)
            impl.do_context_takeover_write(impl.role);
    }
    else if(! fh.mask)
    {
        if(! impl.wr_frag)
        {
            // no mask, no autofrag
            fh.fin = fin;
            fh.len = remain;
            detail::fh_buffer fh_buf;
            detail::write<
                flat_static_buffer_base>(fh_buf, fh);
            impl.wr_cont = ! fin;
            net::write(impl.stream(),
                buffers_cat(fh_buf.data(), buffers), ec);
            if(impl.check_stop_now(ec))
                return bytes_transferred;
            bytes_transferred += remain;
        }
        else
        {
            // no mask, autofrag
            BOOST_ASSERT(impl.wr_buf_size != 0);
            buffers_suffix<
                ConstBufferSequence> cb{buffers};
            for(;;)
            {
                auto const n = clamp(remain, impl.wr_buf_size);
                remain -= n;
                fh.len = n;
                fh.fin = fin ? remain == 0 : false;
                detail::fh_buffer fh_buf;
                detail::write<
                    flat_static_buffer_base>(fh_buf, fh);
                impl.wr_cont = ! fin;
                net::write(impl.stream(),
                    beast::buffers_cat(fh_buf.data(),
                        beast::buffers_prefix(n, cb)), ec);
                bytes_transferred += n;
                if(impl.check_stop_now(ec))
                    return bytes_transferred;
                if(remain == 0)
                    break;
                fh.op = detail::opcode::cont;
                cb.consume(n);
            }
        }
    }
    else if(! impl.wr_frag)
    {
        // mask, no autofrag
        fh.fin = fin;
        fh.len = remain;
        fh.key = this->impl_->create_mask();
        detail::prepared_key key;
        detail::prepare_key(key, fh.key);
        detail::fh_buffer fh_buf;
        detail::write<
            flat_static_buffer_base>(fh_buf, fh);
        buffers_suffix<
            ConstBufferSequence> cb{buffers};
        {
            auto const n =
                clamp(remain, impl.wr_buf_size);
            auto const b =
                net::buffer(impl.wr_buf.get(), n);
            net::buffer_copy(b, cb);
            cb.consume(n);
            remain -= n;
            detail::mask_inplace(b, key);
            impl.wr_cont = ! fin;
            net::write(impl.stream(),
                buffers_cat(fh_buf.data(), b), ec);
            bytes_transferred += n;
            if(impl.check_stop_now(ec))
                return bytes_transferred;
        }
        while(remain > 0)
        {
            auto const n =
                clamp(remain, impl.wr_buf_size);
            auto const b =
                net::buffer(impl.wr_buf.get(), n);
            net::buffer_copy(b, cb);
            cb.consume(n);
            remain -= n;
            detail::mask_inplace(b, key);
            net::write(impl.stream(), b, ec);
            bytes_transferred += n;
            if(impl.check_stop_now(ec))
                return bytes_transferred;
        }
    }
    else
    {
        // mask, autofrag
        BOOST_ASSERT(impl.wr_buf_size != 0);
        buffers_suffix<
            ConstBufferSequence> cb(buffers);
        for(;;)
        {
            fh.key = this->impl_->create_mask();
            detail::prepared_key key;
            detail::prepare_key(key, fh.key);
            auto const n =
                clamp(remain, impl.wr_buf_size);
            auto const b =
                net::buffer(impl.wr_buf.get(), n);
            net::buffer_copy(b, cb);
            detail::mask_inplace(b, key);
            fh.len = n;
            remain -= n;
            fh.fin = fin ? remain == 0 : false;
            impl.wr_cont = ! fh.fin;
            detail::fh_buffer fh_buf;
            detail::write<
                flat_static_buffer_base>(fh_buf, fh);
            net::write(impl.stream(),
                buffers_cat(fh_buf.data(), b), ec);
            bytes_transferred += n;
            if(impl.check_stop_now(ec))
                return bytes_transferred;
            if(remain == 0)
                break;
            fh.op = detail::opcode::cont;
            cb.consume(n);
        }
    }
    return bytes_transferred;
}

template<class NextLayer, bool deflateSupported>
template<class ConstBufferSequence, class WriteHandler>
BOOST_BEAST_ASYNC_RESULT2(WriteHandler)
stream<NextLayer, deflateSupported>::
async_write_some(bool fin,
    ConstBufferSequence const& bs, WriteHandler&& handler)
{
    static_assert(is_async_stream<next_layer_type>::value,
        "AsyncStream type requirements not met");
    static_assert(net::is_const_buffer_sequence<
        ConstBufferSequence>::value,
            "ConstBufferSequence type requirements not met");
    return net::async_initiate<
        WriteHandler,
        void(error_code, std::size_t)>(
            run_write_some_op{},
            handler,
            impl_,
            fin,
            bs);
}

//------------------------------------------------------------------------------

template<class NextLayer, bool deflateSupported>
template<class ConstBufferSequence>
std::size_t
stream<NextLayer, deflateSupported>::
write(ConstBufferSequence const& buffers)
{
    static_assert(is_sync_stream<next_layer_type>::value,
        "SyncStream type requirements not met");
    static_assert(net::is_const_buffer_sequence<
        ConstBufferSequence>::value,
            "ConstBufferSequence type requirements not met");
    error_code ec;
    auto const bytes_transferred = write(buffers, ec);
    if(ec)
        BOOST_THROW_EXCEPTION(system_error{ec});
    return bytes_transferred;
}

template<class NextLayer, bool deflateSupported>
template<class ConstBufferSequence>
std::size_t
stream<NextLayer, deflateSupported>::
write(ConstBufferSequence const& buffers, error_code& ec)
{
    static_assert(is_sync_stream<next_layer_type>::value,
        "SyncStream type requirements not met");
    static_assert(net::is_const_buffer_sequence<
        ConstBufferSequence>::value,
            "ConstBufferSequence type requirements not met");
    return write_some(true, buffers, ec);
}

template<class NextLayer, bool deflateSupported>
template<class ConstBufferSequence, class WriteHandler>
BOOST_BEAST_ASYNC_RESULT2(WriteHandler)
stream<NextLayer, deflateSupported>::
async_write(
    ConstBufferSequence const& bs, WriteHandler&& handler)
{
    static_assert(is_async_stream<next_layer_type>::value,
        "AsyncStream type requirements not met");
    static_assert(net::is_const_buffer_sequence<
        ConstBufferSequence>::value,
            "ConstBufferSequence type requirements not met");
    return net::async_initiate<
        WriteHandler,
        void(error_code, std::size_t)>(
            run_write_some_op{},
            handler,
            impl_,
            true,
            bs);
}

} // websocket
} // beast
} // boost

#endif