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

#include "dv_output.hpp"

#include <boost/filesystem.hpp>
#include <boost/nowide/fstream.hpp>

#include <chrono>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>

class OutFile : public dv::ModuleBase {
private:
	struct fileHeader {
		dv::IOHeader data;
		size_t offset = 0;
		size_t length = 0;
	};

	struct fileOutput {
		boost::nowide::ofstream stream;
		fileHeader header;
	};

	fileOutput fileStream;

	dv::OutputEncoder output;

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

	static const char *initDescription() {
		return ("Write AEDAT 4.0 data to a file.");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		// Force LZ4 compression by default for files.
		config.add("compression", dv::ConfigOption::stringOption("Type of data compression to apply.",
									  dv::EnumNameCompressionType(dv::CompressionType::LZ4)));

		initConfigOptionsOutputCommon(config);

		config.add("directory", dv::ConfigOption::directoryOption(
									"Directory in which to save data.", dv::portable_get_user_home_directory()));

		config.add("prefix", dv::ConfigOption::stringOption(
								 "String to add in front of the date part of the final save name.", "dvSave"));

		config.setPriorityOptions({"directory", "prefix", "compression"});
	}

	static void advancedStaticInit(dvModuleData mData) {
		// TODO: flexible number of outputs.

		dv::Cfg::Node(mData->moduleNode)
			.create<dv::CfgType::STRING>("file", "", {0, INT32_MAX}, dv::CfgFlags::READ_ONLY | dv::CfgFlags::NO_EXPORT,
				"Full name of file being currently written.");
	}

	static size_t writeHeader(fileOutput &file) {
		// Construct serialized flatbuffer header.
		flatbuffers::FlatBufferBuilder headerBuild{2 * 1024};

		headerBuild.ForceDefaults(true);

		auto headerOffset = dv::IOHeaderFlatbuffer::Pack(headerBuild, &file.header.data);

		dv::FinishSizePrefixedIOHeaderBuffer(headerBuild, headerOffset);

		uint8_t *headerData   = headerBuild.GetBufferPointer();
		size_t headerDataSize = headerBuild.GetSize();

		if (file.header.offset == 0) {
			// First header write, store offset and length.
			file.header.offset = static_cast<size_t>(file.stream.tellp());
			file.header.length = headerDataSize;

			dv::runtime_assert(
				[&file] {
					return file.header.offset
						== static_cast<std::underlying_type_t<dv::Constants>>(dv::Constants::AEDAT_VERSION_LENGTH);
				},
				[] {
					return "header offset is not right after AEDAT version string";
				});
			dv::runtime_assert(
				[&file] {
					return file.header.length > 14;
				},
				[] {
					return "header must be at least 14 characters due to XML";
				});
		}
		else {
			// Final header write, use offset and length.
			dv::runtime_assert(
				[&file, headerDataSize] {
					return headerDataSize == file.header.length;
				},
				[] {
					return "header size cannot change on successive writes";
				});
		}

		file.stream.seekp(static_cast<std::streamoff>(file.header.offset));
		file.stream.write(reinterpret_cast<char *>(headerData), static_cast<std::streamsize>(headerDataSize));

		file.stream.flush();

		// Return written size.
		return (headerDataSize);
	}

	OutFile() : output(&config, &log) {
		// Get current time for suffix part. Determine new save name.
		const auto saveName = fmt::format("{:s}-{:%Y_%m_%d_%H_%M_%S}{:s}", config.get<dv::CfgType::STRING>("prefix"),
			dv::localTime(std::chrono::system_clock::now()), dv::io::support::AEDAT4_FILE_EXTENSION);

		// Create directory.
		boost::filesystem::path saveDir{config.get<dv::CfgType::STRING>("directory")};

		if (saveDir.empty()) {
			saveDir = dv::portable_get_user_home_directory();
		}

		boost::filesystem::create_directories(saveDir);

		// Add compression information at startup, won't change during execution.
		auto compressionType = config.get<dv::CfgType::STRING>("compression");

		output.setCompressionType(dv::parseCompressionTypeFromString(compressionType));
		output.setTrackDataTable(true);

		// Setup output node and convert it to string for file output.
		dv::OutputEncoder::makeOutputNode(
			inputs.infoNode("output0"), moduleNode.getRelativeNode("outInfo/0/"), compressionType);

		if (inputs.isConnected("output1")) {
			dv::OutputEncoder::makeOutputNode(
				inputs.infoNode("output1"), moduleNode.getRelativeNode("outInfo/1/"), compressionType);
		}

		if (inputs.isConnected("output2")) {
			dv::OutputEncoder::makeOutputNode(
				inputs.infoNode("output2"), moduleNode.getRelativeNode("outInfo/2/"), compressionType);
		}

		if (inputs.isConnected("output3")) {
			dv::OutputEncoder::makeOutputNode(
				inputs.infoNode("output3"), moduleNode.getRelativeNode("outInfo/3/"), compressionType);
		}

		auto outputNodeString = moduleNode.getRelativeNode("outInfo/").exportSubTreeToXMLString(true);

		// Open file for writing.
		auto fileName     = (saveDir / saveName);
		fileStream.stream = boost::nowide::ofstream(fileName, std::ios::out | std::ios::binary);

		moduleNode.updateReadOnly<dv::CfgType::STRING>("file", fileName.string());

		// Write AEDAT 4.0 header.
		fileStream.stream.write("#!AER-DAT4.0\r\n",
			static_cast<std::underlying_type_t<dv::Constants>>(dv::Constants::AEDAT_VERSION_LENGTH));

		output.addStatisticBytes(
			static_cast<std::underlying_type_t<dv::Constants>>(dv::Constants::AEDAT_VERSION_LENGTH));

		fileStream.header.data.infoNode          = outputNodeString;
		fileStream.header.data.compression       = dv::parseCompressionTypeFromString(compressionType);
		fileStream.header.data.dataTablePosition = -1; // Data jump table not present yet.

		size_t headerSize = writeHeader(fileStream);

		output.addStatisticBytes(static_cast<int64_t>(headerSize));

		log.info.format("Writing data to file '{:s}'.", fileName.string());

		output.updateStartTime();
	}

	~OutFile() override {
		// Ensure data written out until now is committed.
		fileStream.stream.flush();

		// Ensure statistics are flushed to config-tree.
		output.updateStatistics(true);

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

		// Write out file data table and update header.
		auto fileDataTable = output.getTrackDataTable();

		fileStream.header.data.dataTablePosition = fileStream.stream.tellp();

		fileStream.stream.write(fileDataTable->getData(), static_cast<std::streamsize>(fileDataTable->getSize()));

		writeHeader(fileStream);

		// Ensure statistics are up-to-date with file table size.
		output.addStatisticBytes(static_cast<int64_t>(fileDataTable->getSize()));
		output.updateStatistics(true);
	}

	void writePacket(const char *outputName, int32_t streamId) {
		auto input = dvModuleInputGet(moduleData, outputName);

		if (input != nullptr) {
			// outData.size is in bytes, not elements.
			auto outMessage = output.processPacket(input, streamId, fileStream.stream.tellp());

			// Write packet header first.
			fileStream.stream.write(reinterpret_cast<const char *>(outMessage->getHeader()), sizeof(dv::PacketHeader));

			// Then write packet content.
			fileStream.stream.write(outMessage->getData(), static_cast<std::streamsize>(outMessage->getSize()));

			// Dismiss data, freeing it.
			dvModuleInputDismiss(moduleData, outputName, input);
		}
	}

	void run() override {
		writePacket("output0", 0);

		if (inputs.isConnected("output1")) {
			writePacket("output1", 1);
		}

		if (inputs.isConnected("output2")) {
			writePacket("output2", 2);
		}

		if (inputs.isConnected("output3")) {
			writePacket("output3", 3);
		}

		// Check if timeout expired.
		if (output.timeElapsed() && config.getBool("running")) {
			log.warning << "File output timeout elapsed, terminating recording." << std::endl;
			config.setBool("running", false);
		}
	}
};

registerModuleClass(OutFile)
