module;
#include <fstream>
#include <ios>
#include <istream>
#define ASIO_DISABLE_VISIBILITY
#include <asio.hpp>
#include <cstdint>
#include <cstddef>
#include <functional>
#include <memory>
#include <ostream>
#include <sstream>
#include <string>
#include <vector>

export module eu07.network.backend.asio;

import eu07.network.message;
import eu07.network.network;
import eu07.gl.buffer;
import eu07.rendering.geometrybank;
import eu07.scene.sn_utils;
import eu07.utilities.logs;


export {
#ifdef DBG_NEW
#undef new
#endif



namespace network::tcp
{
    const uint32_t NETWORK_MAGIC = 0x37305545;
	const uint32_t MAX_MSG_SIZE = 100000;

	class connection : public network::connection
	{
		friend class server;
		friend class client;

	public:
                connection(asio::io_context &io_ctx, bool client = false, size_t counter = 0);
		~connection();

		virtual void connected() override;
		virtual void disconnect() override;
		virtual void send_messages(const std::vector<std::shared_ptr<message> > &messages) override;
		virtual void send_message(const message &msg) override;

		asio::ip::tcp::socket m_socket;

	private:
		std::string m_header_buffer;
		std::string m_body_buffer;

		void write_message(const message &msg, std::ostream &stream);
		void send_data(std::shared_ptr<std::string> buffer);
		void read_header();
		void handle_header(const asio::error_code &err, size_t bytes_transferred);
		void handle_data(const asio::error_code &err, size_t bytes_transferred);
	};

	class server : public network::server
	{
	private:
		void accept_conn();
		void handle_accept(std::shared_ptr<connection> conn, const asio::error_code &err);

		asio::ip::tcp::acceptor m_acceptor;
                asio::io_context &m_io_ctx;

	public:
		server(std::shared_ptr<std::istream> buf, asio::io_context &io_ctx, const std::string &host, uint32_t port);
	};

	class client : public network::client
	{
	private:
		void handle_accept(const asio::error_code &err);
		asio::io_context &io_ctx;
		std::string host;
		uint32_t port;

	protected:
		virtual void connect() override;

	public:
		client(asio::io_context &io_ctx, const std::string &host, uint32_t port);
	};

	class asio_manager : public network::backend_manager {
		asio::io_context io_context;

	public:
		asio_manager();

		virtual std::shared_ptr<network::server> create_server(std::shared_ptr<std::fstream>, const std::string &conf) override;
		virtual std::shared_ptr<network::client> create_client(const std::string &conf) override;
		virtual void update() override;
	};

	asio_manager manager;
}

}  // export

// The implementation lives here rather than in a separate implementation unit:
// asio defines its customization points as objects in an anonymous namespace, and
// GCC emits those a second time in any implementation unit that also includes
// asio.hpp, which the assembler rejects as duplicate symbols. One translation unit
// including asio is the way around it.

network::tcp::connection::connection(asio::io_context &io_ctx, bool client, size_t counter)
    : network::connection(client, counter),
      m_socket(io_ctx)
{
	m_header_buffer.resize(8);
}

network::tcp::connection::~connection()
{
	m_socket.shutdown(m_socket.shutdown_both);
	m_socket.close();
}

void network::tcp::connection::disconnect()
{
	network::connection::disconnect();
}

void network::tcp::connection::send_data(std::shared_ptr<std::string> buffer)
{
	asio::async_write(m_socket, asio::buffer(*buffer.get()), std::bind(&connection::send_complete, this, buffer));
}

void network::tcp::connection::connected()
{
	network::connection::connected();
	read_header();
}

void network::tcp::connection::read_header()
{
	asio::async_read(m_socket, asio::buffer(m_header_buffer),
	                    std::bind(&connection::handle_header, this,
	                              std::placeholders::_1, std::placeholders::_2));
}

void network::tcp::connection::handle_header(const asio::error_code &err, size_t bytes_transferred)
{
	if (err) {
		disconnect();
		return;
	}

	std::istringstream header(m_header_buffer);
	if (m_header_buffer.size() != bytes_transferred) {
		disconnect();
		return;
	}

	uint32_t sig = sn_utils::ld_uint32(header);
	if (sig != NETWORK_MAGIC) {
		disconnect();
		return;
	}

	uint32_t len = sn_utils::ld_uint32(header);
	m_body_buffer.resize(len);
	if (len > MAX_MSG_SIZE) {
		disconnect();
		return;
	}

	asio::async_read(m_socket, asio::buffer(m_body_buffer),
	                 std::bind(&connection::handle_data, this, std::placeholders::_1, std::placeholders::_2));
}

void network::tcp::connection::handle_data(const asio::error_code &err, size_t bytes_transferred)
{
	if (err) {
		disconnect();
		return;
	}

	if (m_body_buffer.size() != bytes_transferred) {
		disconnect();
		return;
	}

	std::istringstream stream(m_body_buffer);
	std::shared_ptr<message> msg = deserialize_message(stream);
	if (message_handler)
		message_handler(*msg);

	read_header();
}

void network::tcp::connection::write_message(const message &msg, std::ostream &stream)
{
	size_t beg = (size_t)stream.tellp();

	sn_utils::ls_uint32(stream, NETWORK_MAGIC);
	sn_utils::ls_uint32(stream, 0);

	serialize_message(msg, stream);

	size_t size = (size_t)stream.tellp() - beg - 8;
	if (size > MAX_MSG_SIZE) {
		ErrorLog("net: message too big", logtype::net);
		return;
	}
	stream.seekp(beg + 4, std::ios_base::beg);
	sn_utils::ls_uint32(stream, size);

	stream.seekp(0, std::ios_base::end);
}

void network::tcp::connection::send_messages(const std::vector<std::shared_ptr<message> > &messages)
{
	if (messages.size() == 0)
		return;

	std::ostringstream stream;
	for (auto const &msg : messages)
		write_message(*msg.get(), stream);

	stream.flush();
	send_data(std::make_shared<std::string>(stream.str()));
}

void network::tcp::connection::send_message(const message &msg)
{
	std::ostringstream stream;
	write_message(msg, stream);

	stream.flush();
	send_data(std::make_shared<std::string>(stream.str()));
}

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

network::tcp::server::server(std::shared_ptr<std::istream> buf, asio::io_context &io_ctx, const std::string &host, uint32_t port)
    : network::server(buf), m_acceptor(io_ctx), m_io_ctx(io_ctx)
{
	auto endpoint = asio::ip::tcp::endpoint(asio::ip::make_address(host), port);
	m_acceptor.open(endpoint.protocol());
	m_acceptor.set_option(asio::socket_base::reuse_address(true));
	m_acceptor.set_option(asio::ip::tcp::no_delay(true));
	m_acceptor.bind(endpoint);
	m_acceptor.listen(10);

	accept_conn();
}

void network::tcp::server::accept_conn()
{
    std::shared_ptr<connection> conn = std::make_shared<connection>(m_io_ctx);
    conn->set_handler(std::bind(&server::handle_message, this, conn, std::placeholders::_1));

    m_acceptor.async_accept(conn->m_socket, std::bind(&server::handle_accept, this, conn, std::placeholders::_1));
}

void network::tcp::server::handle_accept(std::shared_ptr<connection> conn, const asio::error_code &err)
{
	if (!err)
	{
		clients.emplace_back(conn);
		conn->connected();
	}
	else
	{
		conn->state = connection::DEAD;
		WriteLog(std::string("net: failed to accept client: " + err.message()), logtype::net);
	}

	accept_conn();
}

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

network::tcp::client::client(asio::io_context &io_ctxx, const std::string &hostarg, uint32_t portarg)
    : host(hostarg), port(portarg), io_ctx(io_ctxx)
{
}

void network::tcp::client::connect()
{
	if (this->conn)
		return;

	std::shared_ptr<connection> conn = std::make_shared<connection>(io_ctx, true, resume_frame_counter);
	conn->set_handler(std::bind(&client::handle_message, this, conn, std::placeholders::_1));

	asio::ip::tcp::endpoint endpoint(
	            asio::ip::make_address(host), port);
	conn->m_socket.open(endpoint.protocol());
	conn->m_socket.set_option(asio::ip::tcp::no_delay(true));
	conn->m_socket.async_connect(endpoint,
	                    std::bind(&client::handle_accept, this, std::placeholders::_1));

	this->conn = conn;

}

void network::tcp::client::handle_accept(const asio::error_code &err)
{
	if (!err)
	{
		conn->connected();
	}
	else
	{
		WriteLog(std::string("net: failed to connect: " + err.message()), logtype::net);
		conn->disconnect();
	}
}

network::tcp::asio_manager::asio_manager() {
	backend_list().emplace("tcp", this);
}

std::shared_ptr<network::server> network::tcp::asio_manager::create_server(std::shared_ptr<std::fstream> backbuffer, const std::string &conf) {
	std::istringstream stream(conf);

	std::string host;
	std::getline(stream, host, ':');

	int port;
	stream >> port;

	return std::make_shared<tcp::server>(backbuffer, io_context, host, port);
}

std::shared_ptr<network::client> network::tcp::asio_manager::create_client(const std::string &conf) {
	std::istringstream stream(conf);

	std::string host;
	std::getline(stream, host, ':');

	int port;
	stream >> port;

	return std::make_shared<tcp::client>(io_context, host, port);
}

void network::tcp::asio_manager::update() {
	io_context.restart();
	io_context.poll();
}