#include "../../include/dv-sdk/cross/asio_tcptlssocket.hpp"

#include "dv_output.hpp"

#include <dv-processing/io/support/xml_config_io.hpp>
#include <dv-processing/io/writer.hpp>

class NetSocket;

class Connection : public std::enable_shared_from_this<Connection> {
private:
	NetSocket *parent;
	WriteOrderedSocket<LocalSocket> socket;
	uint8_t keepAliveReadSpace;

public:
	Connection(asioLocal::socket s, NetSocket *server);
	~Connection();

	void close();
	void start(dv::io::Writer *writer, const std::string_view &infoNode);
	void writeMessage(const std::shared_ptr<const dv::io::support::IODataBuffer> &packet);
	void writeIOHeader(const std::shared_ptr<const dv::io::support::IODataBuffer> &ioHeader);

private:
	void keepAliveByReading();
	void handleError(const boost::system::error_code &error, const char *message);
};

class NetSocket : public dv::ModuleBase {
private:
	asio::io_context ioService;
	asioLocal::acceptor acceptor;

	std::vector<Connection *> clients;
	dv::OutputEncoder output;
	size_t queuedPackets;

	std::string mInfoNode;
	dv::ModuleStatistics *statistics = nullptr;
	std::unique_ptr<dv::io::Writer> mAedat4Writer;

public:
	static void initInputs(dv::InputDefinitionList &in) {
		in.addInput("output0", dv::Types::anyIdentifier, false);
	}

	static const char *initDescription() {
		return ("Send AEDAT 4 data out via local sockets to connected clients (server mode).");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		initConfigOptionsOutputCommon(config);

		config.add("socketPath",
			dv::ConfigOption::stringOption("Socket file path to listen on (server mode).", "/tmp/dv-runtime.sock"));
		config.add(
			"maxConnectionsBacklog", dv::ConfigOption::intOption("Maximum number of pending connections.", 5, 1, 32));
		config.add("maxConcurrentConnections",
			dv::ConfigOption::intOption("Maximum number of concurrent active connections (clients).", 10, 1, 128));
		config.add("maxDataBacklog",
			dv::ConfigOption::intOption(
				"Maximum number of data packets in queue waiting to be sent to clients (-1 for infinite).", 5, -1,
				256));

		config.setPriorityOptions({"socketPath", "compression"});
	}

	NetSocket() : acceptor(ioService), output(&config, &log), queuedPackets(0) {
		// Add compression information at startup, won't change during execution.
		const auto compressionType = config.get<dv::CfgType::STRING>("compression");
		const auto compressionEnum = dv::parseCompressionTypeFromString(compressionType);

		auto stats = std::make_unique<dv::ModuleStatistics>(&output);
		// Intentional pointer leak, so statistics multiplier could be updated
		statistics = stats.get();

		mAedat4Writer = std::make_unique<dv::io::Writer>(compressionEnum, std::move(stats));

		output.setCompressionType(compressionEnum);

		// Required input is always present.
		dv::OutputEncoder::makeOutputNode(
			inputs.infoNode("output0"), moduleNode.getRelativeNode("outInfo/0/"), compressionType);

		// --- Generate IO header starts
		dv::io::support::XMLTreeNode root("outInfo");
		dv::xml_compat::transferNode(moduleNode.getRelativeNode("outInfo/"), root);

		dv::io::support::XMLConfigWriter xml(root);
		mInfoNode = xml.getXMLContent();
		// --- Generate IO header ends

		// Configure acceptor.
		std::filesystem::remove(config.get<dv::CfgType::STRING>("socketPath"));

		auto endpoint = asioLocal::endpoint(config.get<dv::CfgType::STRING>("socketPath"));

		try {
			acceptor.open(endpoint.protocol());
			// REUSE_ADDRESS does not exist for local Unix sockets.
			acceptor.bind(endpoint);
			acceptor.listen(config.get<dv::CfgType::INT>("maxConnectionsBacklog"));
		}
		catch (const std::exception &) {
			acceptor.close();
			moduleNode.getRelativeNode("outInfo/").removeNode();
			throw;
		}

		acceptStart();

		log.debug.format("Output server ready on {:s}.", config.get<dv::CfgType::STRING>("socketPath"));

		output.updateStartTime();
	}

	~NetSocket() override {
		acceptor.close();

		// Post 'close all connections' to end of async queue,
		// so that any other callbacks, such as pending accepts,
		// are executed first, and we really close all sockets.
		asio::post(ioService, [this]() {
			// Close all open connections, hard.
			for (const auto client : clients) {
				client->close();
			}
		});

		// Wait for all clients to go away.
		while (!clients.empty()) {
			handleIO();
		}

		// Ensure statistics are up-to-date.
		output.updateStatistics(true);

		// Cleanup manually added output info nodes.
		moduleNode.getRelativeNode("outInfo/").removeNode();
	}

	void handleIO() {
		ioService.poll();
#if BOOST_ASIO_NET_NEW_INTERFACE == 1
		ioService.restart();
#else
		ioService.reset();
#endif
	}

	void removeClient(Connection *client) {
		dv::vectorRemove(clients, client);
	}

	void handleData() {
		if (clients.empty()) {
			return;
		}

		auto input0 = dvModuleInputGet(moduleData, "output0");

		if (input0 != nullptr) {
			// TODO: this only works with default, known types!
			dv::types::TypedObject input(*dv::io::support::defaultTypeResolver(input0->typeId));
			input.obj = input0->obj;

			auto maxDataBacklog = config.get<dv::CfgType::INT>("maxDataBacklog");

			if ((maxDataBacklog == -1) || (queuedPackets <= static_cast<size_t>(maxDataBacklog))) {
				queuedPackets++;

				mAedat4Writer->writePacket(
					&input, 0, [this](const std::shared_ptr<const dv::io::support::IODataBuffer> buffer) {
						auto outBuffer = dv::shared_ptr_wrap_extra_deleter<const dv::io::support::IODataBuffer>(
							buffer, [this](const dv::io::support::IODataBuffer *) {
								queuedPackets--;
							});

						for (const auto client : clients) {
							client->writeMessage(outBuffer);
						}
					});
			}

			// Prevent the typed object destructor freeing the memory
			input.obj = nullptr;

			dvModuleInputDismiss(moduleData, "output0", input0);
		}
	}

	void run() override {
		// Process waiting connections before adding new data.
		handleIO();

		// Push data into the socket
		handleData();

		// Send data out, process writes.
		handleIO();

		// Check if timeout expired.
		if (output.timeElapsed()) {
			log.warning << "Local socket network output timeout elapsed, terminating data transmission server."
						<< std::endl;
			config.setBool("running", false);
		}

		// Write out statistics
		output.updateStatistics(false);
	}

private:
	void acceptStart() {
		acceptor.async_accept([this](const boost::system::error_code &error, asioLocal::socket peer) {
			if (error) {
				// Ignore cancel error, normal on shutdown.
				if (error != asio::error::operation_aborted) {
					log.error.format(
						"Failed to accept connection. Error: {:s} ({:d}).", error.message(), error.value());
				}
			}
			else {
				if (clients.size() >= static_cast<size_t>(config.get<dv::CfgType::INT>("maxConcurrentConnections"))) {
					log.warning.format(
						"Maximum number of clients reached, denying {:s}.", peer.remote_endpoint().path());
					peer.close();
				}
				else {
					auto client = std::make_shared<Connection>(std::move(peer), this);

					clients.push_back(client.get());

					client->start(mAedat4Writer.get(), mInfoNode);
				}

				acceptStart();
			}
		});
	}
};

Connection::Connection(asioLocal::socket s, NetSocket *server) :
	parent(server),
	socket(std::move(s)),
	keepAliveReadSpace(0) {
	parent->log.debug.format("New connection from client {:s}.", socket.remote_endpoint().path());
}

Connection::~Connection() {
	parent->removeClient(this);

	parent->log.debug.format("Closing connection from client {:s}.", socket.remote_endpoint().path());
}

void Connection::start(dv::io::Writer *writer, const std::string_view &infoNode) {
	auto self(shared_from_this());

	socket.start([this, self, writer, &infoNode](const boost::system::error_code &) {
		writer->writeHeader(-1, infoNode, [this](const std::shared_ptr<const dv::io::support::IODataBuffer> header) {
			writeIOHeader(header);
		});

		keepAliveByReading();
	});
}

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

void Connection::writeIOHeader(const std::shared_ptr<const dv::io::support::IODataBuffer> &ioHeader) {
	const auto self(shared_from_this());

	socket.write(asio::buffer(ioHeader->getData(), ioHeader->getDataSize()),
		[this, self, ioHeader = ioHeader](const boost::system::error_code &error, const size_t /*length*/) {
			if (error) {
				handleError(error, "Failed to write IOHeader data");
			}
		});
}

void Connection::writeMessage(const std::shared_ptr<const dv::io::support::IODataBuffer> &packet) {
	auto self(shared_from_this());

	// Write packet header first.
	socket.write(asio::buffer(reinterpret_cast<const std::byte *>(packet->getHeader()), sizeof(dv::PacketHeader)),
		[this, self, packet = packet](const boost::system::error_code &error, const size_t) {
			if (error) {
				handleError(error, "Failed to write message header");
			}
		});

	// Then write packet content.
	socket.write(asio::buffer(packet->getData(), packet->getDataSize()),
		[this, self, packet = packet](const boost::system::error_code &error, const size_t) {
			if (error) {
				handleError(error, "Failed to write message data");
			}
		});
}

void Connection::keepAliveByReading() {
	auto self(shared_from_this());

	socket.read(asio::buffer(&keepAliveReadSpace, sizeof(keepAliveReadSpace)),
		[this, self](const boost::system::error_code &error, size_t /*length*/) {
			if (error) {
				handleError(error, "Read keep-alive failure");
			}
			else {
				handleError(error, "Detected illegal incoming data");
			}
		});
}

void Connection::handleError(const boost::system::error_code &error, const char *message) {
#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.
		parent->log.debug.format(
			"Client {:s}: connection closed ({:d}).", socket.remote_endpoint().path(), error.value());
	}
	else {
		parent->log.error.format("Client {:s}: {:s}. Error: {:s} ({:d}).", socket.remote_endpoint().path(), message,
			error.message(), error.value());
	}
}

registerModuleClass(NetSocket)
