#include "config_server_connection.h"

#include "../../include/dv-sdk/cross/portable_endian.h"

#include "config_server.h"
#include "config_server_actions.h"

// Start at 1, 0 is reserved for system.
std::atomic_uint64_t ConfigServerConnection::clientIDGenerator{1};

ConfigServerConnection::ConfigServerConnection(
	asioTCP::socket s, bool tlsEnabled, asioSSL::context *tlsContext, ConfigServer *server) :
	parent(server),
	socket(std::move(s), tlsEnabled, tlsContext) {
	clientID = clientIDGenerator.fetch_add(1);

	parent->setCurrentClientID(clientID);

	dv::Log(dv::logLevel::DEBUG, "New connection from client {:d} ({:s}:{:d}).", clientID,
		socket.remote_address().to_string(), socket.remote_port());
}

ConfigServerConnection::~ConfigServerConnection() {
	parent->setCurrentClientID(clientID);

	parent->removeClient(this);

	dv::Log(dv::logLevel::DEBUG, "Closing connection from client {:d} ({:s}:{:d}).", clientID,
		socket.remote_address().to_string(), socket.remote_port());
}

void ConfigServerConnection::start() {
	auto self(shared_from_this());

	socket.start(
		[this, self](const boost::system::error_code &error) {
			if (error) {
				handleError(error, "Failed startup (TLS handshake)");
			}
			else {
				readMessageSize();
			}
		},
		asioSSL::stream_base::server);
}

void ConfigServerConnection::close() {
	socket.close();
}

uint64_t ConfigServerConnection::getClientID() {
	return (clientID);
}

void ConfigServerConnection::addPushClient() {
	parent->setCurrentClientID(clientID);

	parent->addPushClient(this);
}

void ConfigServerConnection::removePushClient() {
	parent->setCurrentClientID(clientID);

	parent->removePushClient(this);
}

void ConfigServerConnection::writePushMessage(std::shared_ptr<const flatbuffers::FlatBufferBuilder> message) {
	auto self(shared_from_this());

	socket.write(asio::buffer(message->GetBufferPointer(), message->GetSize()),
		[this, self, message](const boost::system::error_code &error, size_t /*length*/) {
			if (error) {
				handleError(error, "Failed to write push message");
			}
		});
}

void ConfigServerConnection::writeMessage(std::shared_ptr<const flatbuffers::FlatBufferBuilder> message) {
	auto self(shared_from_this());

	socket.write(asio::buffer(message->GetBufferPointer(), message->GetSize()),
		[this, self, message](const boost::system::error_code &error, size_t /*length*/) {
			if (error) {
				handleError(error, "Failed to write message");
			}
			else {
				// Restart, wait for next message.
				readMessageSize();
			}
		});
}

void ConfigServerConnection::readMessageSize() {
	auto self(shared_from_this());

	socket.read(asio::buffer(&incomingMessageSize, sizeof(incomingMessageSize)),
		[this, self](const boost::system::error_code &error, size_t /*length*/) {
			if (error) {
				handleError(error, "Failed to read message size");
			}
			else {
				parent->setCurrentClientID(clientID);

				// Message size is little-endian.
				incomingMessageSize = le32toh(incomingMessageSize);

				// Check for wrong (excessive) message length.
				// Close connection by falling out of scope.
				if (incomingMessageSize > DV_CONFIG_SERVER_MAX_INCOMING_SIZE) {
					dv::Log(dv::logLevel::WARNING, "Client {:d}: message length error ({:d} bytes).", clientID,
						incomingMessageSize);
					return;
				}

				readMessage();
			}
		});
}

void ConfigServerConnection::readMessage() {
	auto self(shared_from_this());

	auto messageBufferPtr = new uint8_t[incomingMessageSize];

	socket.read(asio::buffer(messageBufferPtr, incomingMessageSize),
		[this, self, messageBufferPtr](const boost::system::error_code &error, size_t /*length*/) {
			auto messageBuffer = std::unique_ptr<uint8_t[]>(messageBufferPtr);

			if (error) {
				handleError(error, "Failed to read message");
			}
			else {
				// Any changes coming as a result of clients doing something
				// must originate as a result of a call to this function.
				// So we set the client ID for the current thread to the
				// current client value, so that any listeners will see it too.
				parent->setCurrentClientID(clientID);

				// Now we have the flatbuffer message and can verify it.
				flatbuffers::Verifier verifyMessage(messageBuffer.get(), incomingMessageSize);

				if (!dv::VerifyConfigActionDataBuffer(verifyMessage)) {
					// Failed verification, close connection by falling out of scope.
					dv::Log(dv::logLevel::WARNING, "Client {:d}: message verification error.", clientID);
					return;
				}

				ConfigServerHandleRequest(self, std::move(messageBuffer));
			}
		});
}

void ConfigServerConnection::handleError(const boost::system::error_code &error, const char *message) {
	parent->setCurrentClientID(clientID);

#if defined(OS_MACOS)
	// On MacOS, trying to write to a socket that's being closed can fail
	// with EPROTOTYPE, so we must filter it out here too, details at:
	// https://erickt.github.io/blog/2014/11/19/adventures-in-debugging-a-potential-osx-kernel-bug/
	if ((error == asio::error::eof) || (error == asio::error::broken_pipe) || (error == asio::error::operation_aborted)
		|| (error == boost::system::errc::wrong_protocol_type)) {
#else
	if ((error == asio::error::eof) || (error == asio::error::broken_pipe)
		|| (error == asio::error::operation_aborted)) {
#endif
		// Handle EOF/shutdown separately.
		dv::Log(dv::logLevel::DEBUG, "Client {:d}: connection closed ({:d}).", clientID, error.value());
	}
	else {
		dv::Log(dv::logLevel::ERROR, "Client {:d}: {:s}. Error: {:s} ({:d}).", clientID, message, error.message(),
			error.value());
	}
}
