/**

 * Relational pipes

 * Copyright © 2020 František Kučera (Frantovo.cz, GlobalCode.info)

 *

 * This program is free software: you can redistribute it and/or modify

 * it under the terms of the GNU General Public License as published by

 * the Free Software Foundation, version 3 of the License.

 *

 * This program is distributed in the hope that it will be useful,

 * but WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

 * GNU General Public License for more details.

 *

 * You should have received a copy of the GNU General Public License

 * along with this program. If not, see <http://www.gnu.org/licenses/>.

 */

#pragma once

#include <mutex>

#include <mqueue.h>

#include <limits.h>

#include <ext/stdio_filebuf.h>

#include <sys/wait.h>

#include <semaphore.h>

#include "FilesystemCommand.h"

namespace relpipe {

namespace in {

namespace filesystem {

namespace fs = std::filesystem;

using namespace relpipe::writer;

class MQ {

protected:

	/**

	 * Process where this object was created. 

	 * During fork() this object is copied.

	 * Using this variable we can detect the copy.

	 */

	__pid_t originalPid;

	std::string name;

	mqd_t handle;

	bool unlinkAfterClose;

	static const size_t MAX_DATA_LENGTH = PATH_MAX; // TODO: maybe more or configurable

	static const __syscall_slong_t MAX_MESSAGES = 10; // TODO: maybe configurable or derived from the process count

	template<typename... Args> static mqd_t mqOpen(const char *__name, int __oflag, Args... args) {

		mqd_t handle = mq_open(__name, __oflag, args...);

		if (handle < 0) throw RelpipeWriterException(L"Unable to open POSIX MQ.");

		else return handle;

	}

	/**

	 * @param name

	 * @param handle do not call mq_open() directly, use MQ:mqOpen() instead.

	 * @param unlinkAfterClose

	 */

	MQ(const std::string& name, mqd_t handle, bool unlinkAfterClose = false) : originalPid(getpid()), name(name), handle(handle), unlinkAfterClose(unlinkAfterClose) {

	}

public:

	class Message {

	public:

		enum class Type {

			FILENAME,

			END

		};

		Type type;

		size_t dataLength; // TODO: maybe uint16_t from #include <cstdint> would be enough (and shorten the message minimum size from 16 to 4)

		char data[MQ::MAX_DATA_LENGTH];

		void checkDataLength() const {

			if (dataLength > sizeof (data) || dataLength < 0) throw RelpipeWriterException(L"Invalid POSIX MQ message size.");

		}

		size_t getMessageLength() const {

			return sizeof (*this) - sizeof (data) + dataLength;

		}

		std::string getStringData() {

			return std::string(data, dataLength);

		}

		void setStringData(const std::string& s) {

			if (s.size() > sizeof (data)) throw RelpipeWriterException(L"Unable set message data: string too long.");

			::memcpy(data, s.c_str(), s.size());

			dataLength = s.size();

		}

	};

	virtual ~MQ() {

		mq_close(handle);

		if (unlinkAfterClose && originalPid == getpid()) mq_unlink(name.c_str());

	}

	MQ(const MQ& other) = delete;

	void operator=(const MQ& right) = delete;

};

class MQReader : public MQ {

public:

	MQReader(std::string name) : MQ(name, mqOpen(name.c_str(), O_RDONLY)) {

	}

	void receive(Message* m) {

		int result = mq_receive(handle, (char*) m, sizeof (*m), nullptr);

		if (result < 0) throw RelpipeWriterException(L"Unable to receive POSIX MQ message.");

		m->checkDataLength();

	}

};

class MQWriter : public MQ {

private:

	mq_attr attributes;

	mq_attr* getAttributes() {

		attributes.mq_maxmsg = MQ::MAX_MESSAGES;

		attributes.mq_msgsize = sizeof (Message);

		return &attributes;

	}

public:

	MQWriter(std::string name) : MQ(name, mqOpen(name.c_str(), O_RDWR | O_CREAT | O_EXCL, S_IRUSR | S_IWUSR, getAttributes()), true) {

	}

	void send(const Message* m, unsigned int priority = 0) {

		m->checkDataLength();

		int result = mq_send(handle, (const char*) m, m->getMessageLength(), priority);

		if (result) throw RelpipeWriterException(L"Unable to send POSIX MQ message.");

	}

};

/**

 * TODO: move to a common/streamlet library

 */

class NamedMutex {

private:

	/**

	 * Process where this object was created. 

	 * During fork() this object is copied.

	 * Using this variable we can detect the copy.

	 */

	__pid_t originalPid;

	sem_t* handle;

	std::string name;

	NamedMutex() {

	}

public:

	static NamedMutex* create(std::string name) {

		sem_t* handle = sem_open(name.c_str(), O_CREAT | O_EXCL, S_IRUSR | S_IWUSR);

		if (handle == SEM_FAILED) throw RelpipeWriterException(L"Unable to open POSIX semaphore.");

		NamedMutex* result = new NamedMutex();

		result->name = name;

		result->handle = handle;

		result->originalPid = getpid();

		result->unlock();

		return result;

	}

	~NamedMutex() {

		sem_close(handle);

		if (originalPid == getpid()) sem_unlink(name.c_str());

	}

	NamedMutex(const NamedMutex&) = delete;

	NamedMutex& operator=(const NamedMutex&) = delete;

	void lock() {

		int error = sem_wait(handle);

		if (error) throw RelpipeWriterException(L"Unable to lock POSIX semaphore.");

	}

	void unlock() {

		int error = sem_post(handle);

		if (error) throw RelpipeWriterException(L"Unable to unlock POSIX semaphore.");

	}

};

class ParallelFilesystemWorker : FilesystemWorker {

private:

	std::string queueName;

	NamedMutex& stdoutMutex;

	string_t relationName;

	std::map<string_t, std::shared_ptr < AttributeFinder>> attributeFinders;

	Configuration& configuration;

	std::wstring_convert < codecvt_utf8<wchar_t>> convertor; // TODO: support also other encodings.

public:

	ParallelFilesystemWorker(std::string queueName, NamedMutex& stdoutMutex, string_t relationName, std::map<string_t, std::shared_ptr<AttributeFinder> > attributeFinders, Configuration& configuration) : queueName(queueName), stdoutMutex(stdoutMutex), relationName(relationName), attributeFinders(attributeFinders), configuration(configuration) {

	}

	void run() {

		MQ::Message readBuffer;

		MQReader mq(queueName.c_str());

		std::stringstream writeBuffer;

		std::shared_ptr<RelationalWriter> writer(Factory::create(writeBuffer));

		writeHeader(writer.get(), attributeFinders, relationName, &configuration.fields, false);

		while (true) {

			mq.receive(&readBuffer);

			if (readBuffer.type == MQ::Message::Type::FILENAME) {

				std::stringstream originalName(readBuffer.getStringData());

				processSingleFile(writer, originalName, attributeFinders, configuration, relationName);

				{

					std::lock_guard lock(stdoutMutex);

					std::cout << writeBuffer.rdbuf() << std::flush;

					// TODO: optional (configurable) buffering: write multiple records in a single batch

				}

				writeBuffer.str("");

				writeBuffer.clear();

			} else if (readBuffer.type == MQ::Message::Type::END) {

				break;

			} else {

				throw RelpipeWriterException(L"ParallelFilesystemWorker recieved message of unsupported type: " + std::to_wstring((int) readBuffer.type)); // TODO: better exception

			}

		}

	}

};

class ParallelFilesystemProcess {

private:

	__pid_t subPid;

	ParallelFilesystemProcess(__pid_t subPid) : subPid(subPid) {

	}

	/**

	 * TODO: move to a common library (copied from the AWK module) 

	 */

	static void redirectFD(int oldfd, int newfd) {

		int result = dup2(oldfd, newfd);

		if (result < 0) throw ParallelFilesystemProcess::Exception(L"Unable redirect FD.");

	}

	/**

	 * TODO: move to a common library (copied from the AWK module) 

	 */

	static void closeOrThrow(int fd) {

		int error = close(fd);

		if (error) throw ParallelFilesystemProcess::Exception(L"Unable to close FD: " + std::to_wstring(fd) + L" from PID: " + std::to_wstring(getpid()));

	}

public:

	class Exception : public relpipe::writer::RelpipeWriterException {

	public:

		Exception(std::wstring message) : relpipe::writer::RelpipeWriterException(message) {

		}

	};

	static ParallelFilesystemProcess* create(std::string queueName, int outputFD, NamedMutex& stdoutMutex, string_t relationName, std::map<string_t, std::shared_ptr < AttributeFinder>> attributeFinders, Configuration& configuration) {

		__pid_t subPid = fork();

		if (subPid < 0) {

			throw SubProcess::Exception(L"Unable to fork the hash process.");

		} else if (subPid == 0) {

			// Child process

			closeOrThrow(STDIN_FILENO); // strace -cf will show failed close() calls (same as number of processes)

			if (outputFD != STDOUT_FILENO) redirectFD(outputFD, STDOUT_FILENO);

			ParallelFilesystemWorker w(queueName, stdoutMutex, relationName, attributeFinders, configuration);

			w.run();

			return nullptr;

		} else {

			// Parent process

			return new ParallelFilesystemProcess(subPid);

		}

	}

	int wait() {

		int status = -1;

		::waitpid(subPid, &status, 0);

		return status;

	}

	__pid_t getPid() const {

		return subPid;

	}

};

class ParallelFilesystemCommand : public FilesystemCommand {

public:

	void process(int inputFD, int outputFD, Configuration& configuration) {

		__gnu_cxx::stdio_filebuf<char> inputBuffer(inputFD, std::ios::in);

		__gnu_cxx::stdio_filebuf<char> outputBuffer(outputFD, std::ios::out);

		std::istream input(&inputBuffer);

		std::ostream output(&outputBuffer);

		// Write relation header:

		string_t relationName = fetchRelationName(&configuration);

		std::map<string_t, std::shared_ptr < AttributeFinder>> attributeFinders = createAttributeFinders();

		std::shared_ptr<RelationalWriter> writer(Factory::create(output));

		writeHeader(writer.get(), attributeFinders, relationName, &configuration.fields);

		output.flush();

		// Create queue:

		std::string queueName = "/relpipe-in-filesystem_parallel_" + std::to_string(getpid());

		MQWriter mq(queueName.c_str());

		MQ::Message writeBuffer;

		// Create lock for STDOUT synchronization:

		std::unique_ptr<NamedMutex> stdoutMutex(NamedMutex::create(queueName));

		// Start workers:

		std::vector<std::shared_ptr < ParallelFilesystemProcess>> workerProcesses;

		bool inMainProcess = true;

		for (int i = 0; i < configuration.parallelism; i++) {

			std::shared_ptr<ParallelFilesystemProcess> workerProcess(ParallelFilesystemProcess::create(queueName, outputFD, *stdoutMutex, relationName, createAttributeFinders(), configuration));

			if (workerProcess) {

				workerProcesses.push_back(workerProcess);

			} else {

				inMainProcess = false;

				break;

			}

		}

		if (inMainProcess) {

			// Distribute file names to the workers:

			for (std::stringstream originalName; readNext(input, originalName); reset(originalName)) {

				writeBuffer.type = MQ::Message::Type::FILENAME;

				writeBuffer.setStringData(originalName.str());

				mq.send(&writeBuffer);

			}

			// Tell workers that everything is done:

			writeBuffer.type = MQ::Message::Type::END;

			writeBuffer.setStringData("");

			for (int i = 0; i < configuration.parallelism; i++) mq.send(&writeBuffer);

			// Wait for workers exit:

			std::map<__pid_t, int> failedProcesses;

			for (std::shared_ptr<ParallelFilesystemProcess> p : workerProcesses) {

				int result = p->wait();

				if (result) failedProcesses[p->getPid()] = result;

			}

			if (failedProcesses.size()) {

				std::wstringstream errorMessage;

				errorMessage << L"One or more processes failed: ";

				for (auto failed : failedProcesses) errorMessage << failed.first << L":" << failed.second << L", ";

				throw ParallelFilesystemProcess::Exception(errorMessage.str());

			}

		} else {

			// we are in a worker process → do nothing, finished

		}

	}

};

}

}

}