//
// Websocket, implemented via CRTP for both plain and ssl websockets
//
#include "websocket.h"

namespace beast = boost::beast;         // from <boost/beast.hpp>
namespace http = beast::http;           // from <boost/beast/http.hpp>
namespace net = boost::asio;            // from <boost/asio.hpp>
namespace ssl = boost::asio::ssl;       // from <boost/asio/ssl.hpp>
namespace websocket = beast::websocket;
using tcp = boost::asio::ip::tcp;       // from <boost/asio/ip/tcp.hpp>
using namespace std::placeholders;

// Server session, asynchronous, proxying, implemented w/ CRTP for plain+ssl variants
template<class Derived>
class websocket_session
{
private:
 Derived& derived()
 {
  return static_cast<Derived&>(*this);
 }

 boost::asio::io_context& ioc_;
 boost::asio::ip::tcp::resolver resolver_;
 boost::beast::flat_buffer buffer_in_;
 boost::beast::websocket::stream<beast::tcp_stream> ws_app_;
 boost::beast::flat_buffer buffer_out_;
 std::string host_;
 std::string port_;
 std::string subprotocol_;
 std::string relative_target_;

public: 
 explicit websocket_session(boost::asio::io_context& ioc, std::string&& websocket_address):
  ioc_(ioc),
  resolver_(boost::asio::make_strand(ioc_)),
  ws_app_(boost::asio::make_strand(ioc_)),
  host_{},
  port_{},
  subprotocol_{},
  relative_target_{}
 {
  // Parse websocket address host:port :

  auto colon_pos{websocket_address.find_last_of(':')};

  if (colon_pos == std::string::npos) {
   std::cerr << "Warning: Bad websocket address (colon missing): " << websocket_address << std::endl;
   return;
  }
  
  auto slash_pos{websocket_address.find('/')};
  if (slash_pos == std::string::npos) {
   std::cerr << "Warning: Bad websocket address (slash missing): " << websocket_address << std::endl;
   return;
  }
  if (slash_pos <= colon_pos) {
   std::cerr << "Warning: Bad websocket address: " << websocket_address << std::endl;
   return;
  }

  host_ = websocket_address.substr(0, colon_pos);
  port_ = websocket_address.substr(colon_pos + 1, slash_pos - (colon_pos + 1));
  relative_target_ = websocket_address.substr(slash_pos);
 }

 void fail(boost::beast::error_code ec, char const* what)
 {
  boost::beast::error_code ec2;
  
  if (auto& ws_in{derived().ws_in()}; ws_in.is_open())
   ws_in.close(beast::websocket::close_code::going_away, ec2);
  if (auto& ws_app{ws_app_}; ws_app.is_open())
   ws_app.close(beast::websocket::close_code::going_away, ec2);

  if (ec == websocket::error::closed)
   return;

  std::cerr << what << ": " << ec.message() << "\n";
 }

 //
 // The initial setup path
 //
 
 // Start the asynchronous accept operation
 void do_accept_in(request_type req)
 {
  // Set suggested timeout settings for the websocket
  derived().ws_in().set_option(
      websocket::stream_base::timeout::suggested(
          beast::role_type::server));

  // Set a decorator to change the Server of the handshake
  derived().ws_in().set_option(websocket::stream_base::decorator(
      [](websocket::response_type& res)
      {
          res.set(http::field::server,
              std::string{"Reichwein.IT Webserver"});
      }));

  // Forward subprotocol from request to target websocket
  subprotocol_ = std::string{req[http::field::sec_websocket_protocol]};

  // Accept the websocket handshake
  derived().ws_in().async_accept(
      req,
      beast::bind_front_handler(
          &websocket_session::on_accept_in,
          derived().shared_from_this()));
 }

private:
 void on_accept_in(beast::error_code ec)
 {
  if (ec)
   return fail(ec, "accept in");

  resolver_.async_resolve(host_, port_,
                          beast::bind_front_handler(&websocket_session::on_resolve_app, derived().shared_from_this()));
 }
 
 void on_resolve_app(beast::error_code ec, tcp::resolver::results_type results)
 {
  if (ec)
   return fail(ec, "resolve app");

  beast::get_lowest_layer(ws_app_).async_connect(results,
                          beast::bind_front_handler(&websocket_session::on_connect_app, derived().shared_from_this()));
 }

 void on_connect_app(beast::error_code ec, tcp::resolver::results_type::endpoint_type endpoint)
 {
  if (ec)
   return fail(ec, "connect app");

  beast::get_lowest_layer(ws_app_).expires_never();

  host_ += ':' + std::to_string(endpoint.port());

  // Set suggested timeout settings for the websocket
  ws_app_.set_option(
      websocket::stream_base::timeout::suggested(
	  beast::role_type::client));

  ws_app_.set_option(boost::beast::websocket::stream_base::decorator(
   [](boost::beast::websocket::request_type& req)
   {
    req.set(boost::beast::http::field::user_agent, "Reichwein.IT Webserver Websocket client");
   }));
 
  ws_app_.set_option(boost::beast::websocket::stream_base::decorator(
   [this](boost::beast::websocket::request_type& req)
   {
    req.set(boost::beast::http::field::sec_websocket_protocol, subprotocol_);
   }));

  ws_app_.async_handshake(host_, relative_target_,
                          beast::bind_front_handler(&websocket_session::on_handshake_app, derived().shared_from_this()));
 }

 void on_handshake_app(beast::error_code ec)
 {
  if (ec)
   return fail(ec, "handshake app");

  // Start reading messages from both sides, asynchronously
  do_read_in();
  do_read_app();
 }

 //
 // The input path (client,ws_in_ -> app,ws_app_) via
 //
 
 void
 do_read_in()
 {
  // Read a message into our buffer
  derived().ws_in().async_read(
      buffer_in_,
      beast::bind_front_handler(
          &websocket_session::on_read_in,
          derived().shared_from_this()));
 }

 void
 on_read_in(
     beast::error_code ec,
     std::size_t bytes_transferred)
 {
  boost::ignore_unused(bytes_transferred);

  if (ec) {
   // not error::closed as above, but no other strategy known
   return fail(ec, "read in");
  }

  ws_app_.text(derived().ws_in().got_text());
  
  do_write_app();
 }
 
 void do_write_app()
 {
  ws_app_.async_write(buffer_in_.data(),
      beast::bind_front_handler(
          &websocket_session::on_write_app,
          derived().shared_from_this()));
 }

 void on_write_app(beast::error_code ec, std::size_t bytes_transferred)
 {
  boost::ignore_unused(bytes_transferred);

  if (ec)
   return fail(ec, "write app");

  buffer_in_.consume(buffer_in_.size());

  // Do another read
  do_read_in();
 }

 //
 // The output path (app,ws_app_ -> client,ws_in_)
 //
 
 void do_read_app()
 {
  // Read a message into our buffer
  ws_app_.async_read(
      buffer_out_,
      beast::bind_front_handler(
          &websocket_session::on_read_app,
          derived().shared_from_this()));
 }

 void on_read_app(beast::error_code ec, std::size_t bytes_transferred)
 {
  boost::ignore_unused(bytes_transferred);

  if (ec == websocket::error::closed)
   return;

  if (ec) {
   // not error::closed as above, but no other strategy known
   return fail(ec, "read app");
  }

  do_write_out();
 }

 void do_write_out()
 {
  derived().ws_in().async_write(buffer_out_.data(),
      beast::bind_front_handler(
          &websocket_session::on_write_out,
          derived().shared_from_this()));
 }

 void on_write_out(
     beast::error_code ec,
     std::size_t bytes_transferred)
 {
  boost::ignore_unused(bytes_transferred);

  if (ec)
   return fail(ec, "write out");

  // Clear the buffer
  buffer_out_.consume(buffer_out_.size());

  // Do another read
  do_read_app();
 }

}; // class

class plain_websocket_session:
 public websocket_session<plain_websocket_session>,
 public std::enable_shared_from_this<plain_websocket_session>
{
 boost::beast::websocket::stream<beast::tcp_stream> ws_in_;

public:

 explicit plain_websocket_session(boost::asio::io_context& ioc, beast::tcp_stream&& stream, std::string&& websocket_address):
  websocket_session(ioc, std::move(websocket_address)),
  ws_in_(std::move(stream))
 {
 }

 boost::beast::websocket::stream<beast::tcp_stream>& ws_in()
 {
  return ws_in_;
 }
}; // class

class ssl_websocket_session:
 public websocket_session<ssl_websocket_session>,
 public std::enable_shared_from_this<ssl_websocket_session>
{
 boost::beast::websocket::stream<beast::ssl_stream<beast::tcp_stream>> ws_in_;

public:

 explicit ssl_websocket_session(boost::asio::io_context& ioc, beast::ssl_stream<beast::tcp_stream>&& stream, std::string&& websocket_address):
  websocket_session(ioc, std::move(websocket_address)),
  ws_in_(std::move(stream))
 {
 }

 boost::beast::websocket::stream<beast::ssl_stream<beast::tcp_stream>>& ws_in()
 {
  return ws_in_;
 }
}; // class

void make_websocket_session(boost::asio::io_context& ioc, beast::tcp_stream&& stream, std::string websocket_address, request_type&& req)
{
 std::make_shared<plain_websocket_session>(ioc, std::move(stream), std::move(websocket_address))->do_accept_in(std::move(req));
}

void make_websocket_session(boost::asio::io_context& ioc, beast::ssl_stream<beast::tcp_stream>&& stream, std::string websocket_address, request_type&& req)
{
 std::make_shared<ssl_websocket_session>(ioc, std::move(stream), std::move(websocket_address))->do_accept_in(std::move(req));
}