#include "http.h"

#include "config.h"
#include "server.h"
#include "response.h"
#include "websocket.h"

#include <fmt/core.h>

#include <openssl/ssl.h>
#include <openssl/crypto.h>

#include <boost/asio/buffer.hpp>
#include <boost/beast/core.hpp>
#include <boost/beast/http.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/beast/websocket/ssl.hpp>
#include <boost/asio/buffers_iterator.hpp>
#include <boost/asio/dispatch.hpp>
#include <boost/asio/ssl/context.hpp>
#include <boost/beast/ssl.hpp>
#include <boost/asio/strand.hpp>
#include <boost/config.hpp>

#include <algorithm>
#include <cstddef>
#include <cstdlib>
#include <filesystem>
#include <functional>
#include <iostream>
#include <memory>
#include <optional>
#include <string>
#include <thread>
#include <vector>

namespace fs = std::filesystem;
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 Reichwein;

namespace {

 // Report a failure
void fail(boost::beast::error_code ec, char const* what)
{
 // ssl::error::stream_truncated, also known as an SSL "short read",
 // indicates the peer closed the connection without performing the
 // required closing handshake (for example, Google does this to
 // improve performance). Generally this can be a security issue,
 // but if your communication protocol is self-terminated (as
 // it is with both HTTP and WebSocket) then you may simply
 // ignore the lack of close_notify.
 //
 // https://github.com/boostorg/beast/issues/38
 //
 // https://security.stackexchange.com/questions/91435/how-to-handle-a-malicious-ssl-tls-shutdown
 //
 // When a short read would cut off the end of an HTTP message,
 // Beast returns the error beast::http::error::partial_message.
 // Therefore, if we see a short read here, it has occurred
 // after the message has been completed, so it is safe to ignore it.

 if (ec == boost::asio::ssl::error::stream_truncated)
     return;

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

// Handles an HTTP server connection
template<class Derived>
class session
{
private:
 Derived& derived()
 {
  return static_cast<Derived&>(*this);
 }

 boost::asio::io_context& ioc_;
 beast::flat_buffer buffer_;
 Server& server_;
 std::optional<http::request_parser<http::string_body>> parser_; // need to reset parser every time, no other mechanism currently
 request_type req_;
 std::shared_ptr<response_type> res_;

 void handle_request()
 {
  beast::get_lowest_layer(derived().stream()).expires_after(std::chrono::seconds(300)); // timeout on write by server much longer than read timeout from client
  auto sp = std::make_shared<response_type>(response::generate_response(req_, server_));

  res_ = sp;

  // Write the response
  http::async_write(
      derived().stream(),
      *sp,
      beast::bind_front_handler(
          &session::on_write,
          derived().shared_from_this(),
          sp->need_eof()));
 }

 void handle_websocket()
 {
  beast::get_lowest_layer(derived().stream()).expires_never();
  make_websocket_session(ioc_, std::move(derived().stream()), response::get_websocket_address(req_, server_), parser_->release());
 }

public:
 explicit
 session(
     boost::asio::io_context& ioc,
     Server& server):
      ioc_(ioc),
      server_(server)
 {
 }

 // Start the asynchronous operation
 void
 run()
 {
  // We need to be executing within a strand to perform async operations
  // on the I/O objects in this session.
  net::dispatch(
      derived().stream().get_executor(),
      beast::bind_front_handler(
          &Derived::on_run,
          derived().shared_from_this()));
 }

 void
 do_read()
 {
  // Make the request empty before reading,
  // otherwise the operation behavior is undefined.
  req_ = {};
 
  // this is the way to reset the parser. it's necessary.
  // https://github.com/boostorg/beast/issues/927
  parser_.emplace();
  parser_->body_limit(1000000000); // 1GB limit

  // Set the timeout.
  beast::get_lowest_layer(derived().stream()).expires_after(std::chrono::seconds(30));

  // Read a request
  http::async_read(derived().stream(), buffer_, *parser_,
      beast::bind_front_handler(
          &session::on_read,
          derived().shared_from_this()));
 }

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

  // This means they closed the connection
  if (ec == http::error::end_of_stream)
   return derived().do_close();

  if (ec == http::error::partial_message)
   return; // ignore

  if (ec)
   return fail(ec, "http read");

  req_ = parser_->get();

  if (websocket::is_upgrade(req_))
  {
   handle_websocket();
   return;
  }
  
  // Send the response
  handle_request();
 }

 void
 on_write(
     bool close,
     beast::error_code ec,
     std::size_t bytes_transferred
     )
 {
  boost::ignore_unused(bytes_transferred);

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

  if (close)
  {
   // This means we should close the connection, usually because
   // the response indicated the "Connection: close" semantic.
   return derived().do_close();
  }

  // We're done with the response so delete it
  res_ = nullptr;

  // Read another request
  do_read();
 }

};

class plain_session:
 public session<plain_session>,
 public std::enable_shared_from_this<plain_session>
{
 beast::tcp_stream stream_;

public:
 explicit plain_session(
     boost::asio::io_context& ioc,
     tcp::socket&& socket,
     Server& server):
  session(ioc, server),
  stream_(std::move(socket))
 {
 }
 
 void on_run()
 {
  // We need to be executing within a strand to perform async operations
  // on the I/O objects in this session. Skip ssl handshake for plain http.
  net::dispatch(stream_.get_executor(),
		beast::bind_front_handler(
		    &session::do_read,
		    shared_from_this()));
 }

 void
 do_close()
 {
  // Send a TCP shutdown
  beast::error_code ec;
  stream_.socket().shutdown(tcp::socket::shutdown_send, ec);
  // At this point the connection is closed gracefully
 }

 beast::tcp_stream& stream()
 {
  return stream_;
 }

}; // class

class ssl_session:
 public session<ssl_session>,
 public std::enable_shared_from_this<ssl_session>
{
 beast::ssl_stream<beast::tcp_stream> stream_;
public:
 explicit ssl_session(
     boost::asio::io_context& ioc,
     tcp::socket&& socket,
     ssl::context& ctx,
     Server& server):
  session(ioc, server),
  stream_(std::move(socket), ctx)
 {
 }
 
 void on_run()
 {
  // Set the timeout
  beast::get_lowest_layer(stream_).expires_after(
      std::chrono::seconds(30));

  // Perform the SSL handshake
  stream_.async_handshake(
      ssl::stream_base::server,
      beast::bind_front_handler(
	  &ssl_session::on_handshake,
	  shared_from_this()));
 }

 void
 on_handshake(beast::error_code ec)
 {
  if (ec)
   return fail(ec, "https handshake");

  do_read();
 }

 void
 do_close()
 {
  // Set the timeout.
  beast::get_lowest_layer(stream_).expires_after(std::chrono::seconds(30));

  // Perform the SSL shutdown
  stream_.async_shutdown(
      beast::bind_front_handler(
          &ssl_session::on_shutdown,
          shared_from_this()));
 }

 void
 on_shutdown(beast::error_code ec)
 {
  if (ec)
   return fail(ec, "https shutdown");

  // At this point the connection is closed gracefully
 }

 beast::ssl_stream<beast::tcp_stream>& stream()
 {
  return stream_;
 }

}; // class

// Accepts incoming connections and launches the sessions
template<class Derived> 
class listener
{
private:
 Derived& derived()
 {
  return static_cast<Derived&>(*this);
 }
    
protected:
 net::io_context& ioc_;
 tcp::acceptor acceptor_;
 ::Server& server_;

public:
 explicit listener(
        net::io_context& ioc,
        tcp::endpoint endpoint,
        Server& server):
        ioc_(ioc),
        acceptor_(ioc),
        server_(server)
 {
  beast::error_code ec;

  // Open the acceptor
  acceptor_.open(endpoint.protocol(), ec);
  if (ec)
  {
   //fail(ec, "http listener open");
   throw std::runtime_error(fmt::format("http listener open unsuccessful: {}:{}", endpoint.address().to_string(), endpoint.port()));
   return;
  }

  // Allow address reuse
  acceptor_.set_option(net::socket_base::reuse_address(true), ec);
  if (ec)
  {
   fail(ec, "http listener set_option");
   return;
  }

  // Bind to the server address
  acceptor_.bind(endpoint, ec);
  if (ec)
  {
   //fail(ec, "http listener bind");
   throw std::runtime_error(fmt::format("http listener bind unsuccessful: {}:{}", endpoint.address().to_string(), endpoint.port()));
   return;
  }

  // Start listening for connections
  acceptor_.listen(net::socket_base::max_listen_connections, ec);
  if (ec)
  {
   //fail(ec, "http listener listen");
   throw std::runtime_error(fmt::format("http listen unsuccessful: {}:{}", endpoint.address().to_string(), endpoint.port()));
   return;
  }
 }

 // Start accepting incoming connections
 void
 run()
 {
  do_accept();
 }

protected:
 void
 do_accept()
 {
  // The new connection gets its own strand
  acceptor_.async_accept(
      net::make_strand(ioc_),
      beast::bind_front_handler(
          &Derived::on_accept,
          derived().shared_from_this()));
 }
}; // class

class plain_listener:
 public listener<plain_listener>,
 public std::enable_shared_from_this<plain_listener>
{
public:
 explicit plain_listener(
     net::io_context& ioc,
     tcp::endpoint endpoint,
     Server& server):
     listener(ioc, endpoint, server)
 {
 }

 void
 on_accept(beast::error_code ec, tcp::socket socket)
 {
  if (ec) {
   fail(ec, "plain listener accept");
  } else {
   // Create the session and run it
   std::make_shared<plain_session>(
       ioc_,
       std::move(socket),
       server_)->run();
  }

  // Accept another connection
  do_accept();
 }
}; // class

class ssl_listener:
 public listener<ssl_listener>,
 public std::enable_shared_from_this<ssl_listener>
{
 ssl::context& ctx_;

public:
 explicit ssl_listener(
     net::io_context& ioc,
     ssl::context& ctx,
     tcp::endpoint endpoint,
     Server& server):
     listener(ioc, endpoint, server),
     ctx_(ctx)
 {
 }

 void
 on_accept(beast::error_code ec, tcp::socket socket)
 {
  if (ec) {
   fail(ec, "ssl listener accept");
  } else {
   // Create the session and run it
   std::make_shared<ssl_session>(
       ioc_,
       std::move(socket),
       ctx_,
       server_)->run();
  }

  // Accept another connection
  do_accept();
 }
}; // class

} // namespace

void make_listener(net::io_context& ioc, net::ip::address address, unsigned short port, Server& server)
{
 std::make_shared<plain_listener>(
  ioc,
  tcp::endpoint{address, port},
  server)->run();
}

void make_listener(net::io_context& ioc, ssl::context& ctx, net::ip::address address, unsigned short port, Server& server)
{
 std::make_shared<ssl_listener>(
  ioc,
  ctx,
  tcp::endpoint{address, port},
  server)->run();
}