//
// Copyright (c) 2013-2017 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)
//

#include <beast/core.hpp>
#include <boost/asio.hpp>
#include <cstddef>
#include <memory>
#include <utility>

//[core_sample_echo_op_1

// Read a line and echo it back
//
template<class AsyncStream, class CompletionToken>
beast::async_return_type<CompletionToken, void(beast::error_code)>
async_echo(AsyncStream& stream, CompletionToken&& token);

//]

//[core_sample_echo_op_3

// This composed operation reads a line of input and echoes it back.
//
template<class AsyncStream, class Handler>
class echo_op
{
    // This holds all of the state information required by the operation.
    struct state
    {
        // The stream to read and write to
        AsyncStream& stream;

        // Indicates what step in the operation's state machine
        // to perform next, starting from zero.
        int step = 0;

        // The buffer used to hold the input and output data.
        // Note that we use a custom allocator for performance,
        // this allows the implementation of the io_service to
        // make efficient re-use of memory allocated by composed
        // operations during continuations.
        //
        boost::asio::basic_streambuf<beast::handler_alloc<char, Handler>> buffer;

        // handler_ptr requires that the first parameter to the
        // contained object constructor is a reference to the
        // managed final completion handler.
        //
        explicit state(Handler& handler, AsyncStream& stream_)
            : stream(stream_)
            , buffer((std::numeric_limits<std::size_t>::max)(),
                beast::handler_alloc<char, Handler>{handler})
        {
        }
    };

    // This smart pointer container allocates our state using the
    // memory allocation hooks associated with the final completion
    // handler, manages the lifetime of that handler for us, and
    // enforces the destroy-before-invocation requirement on memory
    // allocated using the hooks.
    //
    beast::handler_ptr<state, Handler> p_;

public:
    // Boost.Asio requires that handlers are CopyConstructible.
    // In some cases, it takes advantage of handlers that are
    // MoveConstructible. This operation supports both.
    //
    echo_op(echo_op&&) = default;
    echo_op(echo_op const&) = default;

    // The constructor simply creates our state variables in
    // the smart pointer container.
    //
    template<class DeducedHandler, class... Args>
    echo_op(AsyncStream& stream, DeducedHandler&& handler)
        : p_(std::forward<DeducedHandler>(handler), stream)
    {
    }

    // Determines if the next asynchronous operation represents a
    // continuation of the asynchronous flow of control associated
    // with the final handler. If we are past step one, it means
    // we have performed an asynchronous operation therefore any
    // subsequent operation would represent a continuation.
    // Otherwise, we propagate the handler's associated value of
    // is_continuation. Getting this right means the implementation
    // may schedule the invokation of the invoked functions more
    // efficiently.
    //
    friend bool asio_handler_is_continuation(echo_op* op)
    {
        // This next call is structured to permit argument
        // dependent lookup to take effect.
        using boost::asio::asio_handler_is_continuation;

        // Always use std::addressof to pass the pointer to the handler,
        // otherwise an unwanted overload of operator& may be called instead.
        return op->p_->step > 1 ||
            asio_handler_is_continuation(std::addressof(op->p_.handler()));
    }

    // Handler hook forwarding. These free functions invoke the hooks
    // associated with the final completion handler. In effect, they
    // make the Asio implementation treat our composed operation the
    // same way it would treat the final completion handler for the
    // purpose of memory allocation and invocation.
    //
    // Our implementation just passes through the call to the hook
    // associated with the final handler.

    friend void* asio_handler_allocate(std::size_t size, echo_op* op)
    {
        using boost::asio::asio_handler_allocate;
        return asio_handler_allocate(size, std::addressof(op->p_.handler()));
    }

    friend void asio_handler_deallocate(void* p, std::size_t size, echo_op* op)
    {
        using boost::asio::asio_handler_deallocate;
        return asio_handler_deallocate(p, size, std::addressof(op->p_.handler()));
    }

    template<class Function>
    friend void asio_handler_invoke(Function&& f, echo_op* op)
    {
        using boost::asio::asio_handler_invoke;
        return asio_handler_invoke(f, std::addressof(op->p_.handler()));
    }

    // Our main entry point. This will get called as our
    // intermediate operations complete. Definition below.
    //
    void operator()(beast::error_code ec, std::size_t bytes_transferred);
};

//]

//[core_sample_echo_op_4

template<class AsyncStream, class Handler>
void echo_op<AsyncStream, Handler>::
operator()(beast::error_code ec, std::size_t bytes_transferred)
{
    // Store a reference to our state. The address of the state won't
    // change, and this solves the problem where dereferencing the
    // data member is undefined after a move.
    auto& p = *p_;

    // Now perform the next step in the state machine
    switch(ec ? 2 : p.step)
    {
        // initial entry
        case 0:
            // read up to the first newline
            p.step = 1;
            return boost::asio::async_read_until(p.stream, p.buffer, "\n", std::move(*this));

        case 1:
            // write everything back
            p.step = 2;
            // async_read_until could have read past the newline,
            // use buffer_prefix to make sure we only send one line
            return boost::asio::async_write(p.stream,
                beast::buffer_prefix(bytes_transferred, p.buffer.data()), std::move(*this));

        case 2:
            p.buffer.consume(bytes_transferred);
            break;
    }

    // Invoke the final handler. If we wanted to pass any arguments
    // which come from our state, they would have to be moved to the
    // stack first, since the `handler_ptr` guarantees that the state
    // is destroyed before the handler is invoked.
    //
    p_.invoke(ec);
    return;
}

//]

//[core_sample_echo_op_2

template<class AsyncStream, class Handler>
class echo_op;

// Read a line and echo it back
//
template<class AsyncStream, class CompletionToken>
beast::async_return_type<CompletionToken, void(beast::error_code)>
async_echo(AsyncStream& stream, CompletionToken&& token)
{
    // Make sure stream meets the requirements. We use static_assert
    // to cause a friendly message instead of an error novel.
    //
    static_assert(beast::is_async_stream<AsyncStream>::value,
        "AsyncStream requirements not met");

    // This helper manages some of the handler's lifetime and
    // uses the result and handler specializations associated with
    // the completion token to help customize the return value.
    //
    beast::async_completion<CompletionToken, void(beast::error_code)> init{token};

    // Create the composed operation and launch it. This is a constructor
    // call followed by invocation of operator(). We use handler_type
    // to convert the completion token into the correct handler type,
    // allowing user defined specializations of the async result template
    // to take effect.
    //
    echo_op<AsyncStream, beast::handler_type<CompletionToken, void(beast::error_code)>>{
        stream, init.completion_handler}(beast::error_code{}, 0);

    // This hook lets the caller see a return value when appropriate.
    // For example this might return std::future<error_code> if
    // CompletionToken is boost::asio::use_future, or this might
    // return an error code if CompletionToken specifies a coroutine.
    //
    return init.result.get();
}

//]

int main()
{
    using address_type = boost::asio::ip::address;
    using socket_type = boost::asio::ip::tcp::socket;
    using endpoint_type = boost::asio::ip::tcp::endpoint;

    // Create a listening socket, accept a connection, perform
    // the echo, and then shut everything down and exit.
    boost::asio::io_service ios;
    socket_type sock{ios};
    boost::asio::ip::tcp::acceptor acceptor{ios};
    endpoint_type ep{address_type::from_string("0.0.0.0"), 0};
    acceptor.open(ep.protocol());
    acceptor.bind(ep);
    acceptor.listen();
    acceptor.accept(sock);
    async_echo(sock,
        [&](beast::error_code ec)
        {
        });
    ios.run();
    return 0;
}