#ifndef DV_OUTPUT_HPP
#define DV_OUTPUT_HPP

#include "dv_io.hpp"

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

#include <chrono>

namespace dv::xml_compat {

inline void transferNodeAttributes(const dv::Config::Node &source, dv::io::support::XMLTreeNode &target) {
	for (const auto &attr : source.getAttributeKeys()) {
		const auto attrType  = source.getAttributeType(attr);
		const auto attrValue = source.getAttribute(attr, attrType);
		switch (attrType) {
			case dv::Config::AttributeType::UNKNOWN:
				break;
			case dv::Config::AttributeType::BOOL:
				target.mAttributes.emplace_back(attr).mValue = attrValue.boolean;
				break;
			case dv::Config::AttributeType::INT:
				target.mAttributes.emplace_back(attr).mValue = attrValue.iint;
				break;
			case dv::Config::AttributeType::LONG:
				target.mAttributes.emplace_back(attr).mValue = attrValue.ilong;
				break;
			case dv::Config::AttributeType::FLOAT:
				target.mAttributes.emplace_back(attr).mValue = attrValue.ffloat;
				break;
			case dv::Config::AttributeType::DOUBLE:
				target.mAttributes.emplace_back(attr).mValue = attrValue.ddouble;
				break;
			case dv::Config::AttributeType::STRING:
				target.mAttributes.emplace_back(attr).mValue = attrValue.string;
				break;
		}
	}
}

inline void transferNodeChildren(const dv::Config::Node &source, dv::io::support::XMLTreeNode &target) {
	for (const auto &child : source.getChildren()) {
		auto &newChild = target.mChildren.emplace_back(child.getName());
		transferNodeAttributes(child, newChild);
		transferNodeChildren(child, newChild);
	}
}

inline void transferNode(const dv::Config::Node &source, dv::io::support::XMLTreeNode &target) {
	transferNodeAttributes(source, target);
	transferNodeChildren(source, target);
}

} // namespace dv::xml_compat

namespace dv {

static void initConfigOptionsOutputCommon(dv::RuntimeConfig &config) {
	std::vector<std::string> compressors;

	for (size_t i = 0; EnumNamesCompressionType()[i] != nullptr; i++) {
		compressors.push_back(EnumNamesCompressionType()[i]);
	}

	config.add("compression", dv::ConfigOption::listOption("Type of data compression to apply.", 0, compressors));
	config.add("timeout",
		dv::ConfigOption::longOption(
			"Recording timeout (in ms), stops recording when timer expires. 0 to disable.", 0, 0, INT64_MAX));

	config.add("writtenPacketsNumber", dv::ConfigOption::statisticOption("Number of packets sent."));
	config.add("writtenPacketsElements", dv::ConfigOption::statisticOption("Number of elements in packets sent."));
	config.add("writtenPacketsSize", dv::ConfigOption::statisticOption("Size in bytes of packets sent."));
	config.add("writtenDataSize", dv::ConfigOption::statisticOption("Size in bytes of data written to disk/network."));
}

struct OutputStatistics {
	int64_t packetsNumber;
	int64_t packetsElements;
	int64_t packetsSize;
	int64_t dataSize;
	std::chrono::steady_clock::time_point startTime;

	OutputStatistics() :
		packetsNumber(0),
		packetsElements(0),
		packetsSize(0),
		dataSize(0),
		startTime(std::chrono::steady_clock::now()) {
	}
};

class OutputData {
private:
	dv::PacketHeader header;
	std::vector<char> buffer;
	flatbuffers::FlatBufferBuilder builder;
	bool isFlatbuffer;

public:
	OutputData() : builder(lz4CompressionChunkSize), isFlatbuffer(true) {
	}

	dv::PacketHeader *getHeader() {
		return (&header);
	}

	const dv::PacketHeader *getHeader() const {
		return (&header);
	}

	flatbuffers::FlatBufferBuilder *getBuilder() {
		return (&builder);
	}

	std::vector<char> *getBuffer() {
		return (&buffer);
	}

	const char *getData() const {
		if (isFlatbuffer) {
			return (reinterpret_cast<const char *>(builder.GetBufferPointer()));
		}
		else {
			return (buffer.data());
		}
	}

	size_t getSize() const {
		if (isFlatbuffer) {
			return (builder.GetSize());
		}
		else {
			return (buffer.size());
		}
	}

	void switchToBuffer() {
		isFlatbuffer = false;
	}
};

class OutputEncoder {
private:
	/// Compression type.
	CompressionType compressionType;
	/// Support LZ4 compression.
	lz4CompressionSupport compressionLZ4;
	/// Support LZ4 compression.
	zstdCompressionSupport compressionZstd;
	/// Output module statistics collection.
	OutputStatistics statistics;
	/// Configuration access.
	dv::RuntimeConfig *config;
	/// Logging access.
	dv::Logger *log;
	/// Support keeping track of written packet's data.
	bool trackDataTable;
	dv::FileDataTable dataTable;

public:
	static void makeOutputNode(
		dv::Config::Node origin, dv::Config::Node destination, const std::string &compressionType) {
		// Required input is always present.
		auto inputInfoNode = origin;
		auto inputNode     = inputInfoNode.getParent();

		auto outputNode     = destination;
		auto outputInfoNode = outputNode.getRelativeNode("info/");

		inputNode.copyTo(outputNode);
		inputInfoNode.copyTo(outputInfoNode);

		// Add original output name.
		outputNode.create<dv::CfgType::STRING>("originalOutputName", inputNode.getName(), {0, 4096},
			dv::CfgFlags::READ_ONLY | dv::CfgFlags::NO_EXPORT, "Name of the origin module's output stream.");

		outputNode.create<dv::CfgType::STRING>("compression", compressionType, {0, 4096},
			dv::CfgFlags::READ_ONLY | dv::CfgFlags::NO_EXPORT, "Type of used compression for this stream.");

		outputNode.create<dv::CfgType::STRING>("originalModuleName", inputNode.getParent().getParent().getName(),
			{0, 4096}, dv::CfgFlags::READ_ONLY | dv::CfgFlags::NO_EXPORT, "Name of the origin module.");
	}

	OutputEncoder(dv::RuntimeConfig *cfg, dv::Logger *logger) :
		compressionType(CompressionType::NONE),
		config(cfg),
		log(logger),
		trackDataTable(false) {
	}

	void setTrackDataTable(bool enable) {
		trackDataTable = enable;

		dataTable.Table.clear();
		dataTable.Table.reserve(1024);
	}

	std::shared_ptr<const OutputData> getTrackDataTable() {
		if (!trackDataTable) {
			return (nullptr);
		}

		auto msgBuild = std::make_shared<OutputData>();

		auto offset = dv::FileDataTableFlatbuffer::Pack(*msgBuild->getBuilder(), &dataTable);

		dv::FinishSizePrefixedFileDataTableBuffer(*msgBuild->getBuilder(), offset);

		compressData(msgBuild.get());

		return (msgBuild);
	}

	void setCompressionType(CompressionType compressType) {
		compressionType = compressType;

		if (compressionType == CompressionType::LZ4 || compressionType == CompressionType::LZ4_HIGH) {
			if (!compressionLZ4.context) {
				// Create LZ4 compression context.
				struct LZ4F_cctx_s *ctx = nullptr;
				auto ret                = LZ4F_createCompressionContext(&ctx, LZ4F_VERSION);
				if (ret != 0) {
					throw std::bad_alloc();
				}

				compressionLZ4.context = std::shared_ptr<struct LZ4F_cctx_s>(ctx, [](struct LZ4F_cctx_s *c) {
					LZ4F_freeCompressionContext(c);
				});
			}

			// Select appropriate compression flags.
			if (compressionType == CompressionType::LZ4_HIGH) {
				compressionLZ4.prefs = &lz4HighCompressionPreferences;
			}
			else {
				compressionLZ4.prefs = &lz4CompressionPreferences;
			}

			compressionLZ4.chunkSize = LZ4F_compressBound(lz4CompressionChunkSize, compressionLZ4.prefs);
			compressionLZ4.endSize   = LZ4F_compressBound(0, compressionLZ4.prefs);
		}

		if (compressionType == CompressionType::ZSTD || compressionType == CompressionType::ZSTD_HIGH) {
			if (!compressionZstd.context) {
				// Create Zstd compression context.
				struct ZSTD_CCtx_s *ctx = ZSTD_createCCtx();
				if (ctx == nullptr) {
					throw std::bad_alloc();
				}

				compressionZstd.context = std::shared_ptr<struct ZSTD_CCtx_s>(ctx, [](struct ZSTD_CCtx_s *c) {
					ZSTD_freeCCtx(c);
				});
			}

			// Select appropriate compression flags.
			if (compressionType == CompressionType::ZSTD_HIGH) {
				compressionZstd.level = ZSTD_maxCLevel();
			}
			else {
				compressionZstd.level = ZSTD_CLEVEL_DEFAULT;
			}
		}
	}

	CompressionType getCompressionType() {
		return (compressionType);
	}

	void addStatisticBytes(int64_t bytes) {
		statistics.dataSize += bytes;
	}

	void addStatistics(
		const int64_t dataSize, const int64_t packetsNumber, const int64_t packetsElements, const int64_t packetsSize) {
		statistics.dataSize        += dataSize;
		statistics.packetsNumber   += packetsNumber;
		statistics.packetsElements += packetsElements;
		statistics.packetsSize     += packetsSize;
	}

	void updateStatistics(bool forceUpdate) {
		config->set<dv::CfgType::LONG>("writtenDataSize", static_cast<int64_t>(statistics.dataSize), forceUpdate);
		config->set<dv::CfgType::LONG>(
			"writtenPacketsNumber", static_cast<int64_t>(statistics.packetsNumber), forceUpdate);
		config->set<dv::CfgType::LONG>(
			"writtenPacketsElements", static_cast<int64_t>(statistics.packetsElements), forceUpdate);
		config->set<dv::CfgType::LONG>("writtenPacketsSize", static_cast<int64_t>(statistics.packetsSize), forceUpdate);
	}

	void updateStartTime() {
		statistics.startTime = std::chrono::steady_clock::now();
	}

	bool timeElapsed() {
		int64_t timeout; // in ms
		config->get("timeout", timeout);

		// Timeout disabled.
		if (timeout == 0) {
			return (false);
		}

		// Get current time, calculate and check difference.
		auto timeNow  = std::chrono::steady_clock::now();
		auto timeDiff = std::chrono::duration_cast<std::chrono::milliseconds>(timeNow - statistics.startTime);

		if (timeDiff.count() < timeout) {
			return (false);
		}

		return (true);
	}

	std::shared_ptr<const OutputData> processPacket(
		const dv::Types::TypedObject *packet, int32_t streamId, int64_t byteOffset) {
		const auto typeInfo = dvTypeSystemGetInfoByID(packet->typeId);

		// Construct serialized flatbuffer packet.
		auto msgBuild = std::make_shared<OutputData>();

		auto offset = (*typeInfo.pack)(msgBuild->getBuilder(), packet->obj);

		msgBuild->getBuilder()->FinishSizePrefixed(flatbuffers::Offset<void>(offset), typeInfo.identifier);

		// Uncompressed statistic.
		statistics.packetsSize += msgBuild->getSize();

		// Compress packet as requested.
		compressData(msgBuild.get());

		// Compressed statistic.
		statistics.dataSize += msgBuild->getSize();

		// Create packet header.
		*msgBuild->getHeader() = PacketHeader(streamId, static_cast<int32_t>(msgBuild->getSize()));

		// Add packet header to written bytes statistic.
		statistics.dataSize += sizeof(PacketHeader);

		// Update packet statistics.
		statistics.packetsNumber++;

		// ANYT/NULL types have no information.
		dvTypeTimeElementExtractor timeElementInfo{-1, -1, -1};

		if (typeInfo.timeElementExtractor != nullptr) {
			timeElementInfo = (*typeInfo.timeElementExtractor)(packet->obj);

			statistics.packetsElements += timeElementInfo.numElements;
		}

		updateTrackDataTable(
			timeElementInfo, byteOffset + static_cast<int64_t>(sizeof(PacketHeader)), *msgBuild->getHeader());

		updateStatistics(false);

		return (msgBuild);
	}

private:
	void updateTrackDataTable(
		const dvTypeTimeElementExtractor &timeElementInfo, int64_t byteOffset, const PacketHeader &packetHeader) {
		// Update packet data table.
		if (trackDataTable) {
			auto dataDef = dv::FileDataDefinition();

			dataDef.ByteOffset     = byteOffset;
			dataDef.PacketInfo     = packetHeader;
			dataDef.NumElements    = timeElementInfo.numElements;
			dataDef.TimestampStart = timeElementInfo.startTimestamp;
			dataDef.TimestampEnd   = timeElementInfo.endTimestamp;

			dataTable.Table.push_back(dataDef);

			log->debug.format(
				"Written packet at offset {:d} - StreamID {:d}, Size {:d}, NumElements {:d}, TimestampStart "
				"{:d}, TimestampEnd {:d}.",
				dataDef.ByteOffset, dataDef.PacketInfo.StreamID(), dataDef.PacketInfo.Size(), dataDef.NumElements,
				dataDef.TimestampStart, dataDef.TimestampEnd);
		}
	}

	void compressData(OutputData *msgBuild) {
		// Flatbuffer encoded packet.
		auto fbData     = msgBuild->getData();
		auto fbDataSize = static_cast<ssize_t>(msgBuild->getSize());

		if (compressionType != CompressionType::NONE) {
			// Switch to vector instead of flatbuffer.
			// The vector will contain the compressed packet.
			msgBuild->switchToBuffer();

			if (compressionType == CompressionType::LZ4 || compressionType == CompressionType::LZ4_HIGH) {
				// Write out header of compressed frame.
				msgBuild->getBuffer()->resize(LZ4F_HEADER_SIZE_MAX);

				auto written = LZ4F_compressBegin(compressionLZ4.context.get(), msgBuild->getBuffer()->data(),
					LZ4F_HEADER_SIZE_MAX, compressionLZ4.prefs);
				if (LZ4F_isError(written)) {
					// Compression error.
					log->error << "LZ4 compression error: " << LZ4F_getErrorName(written) << std::endl;
					msgBuild->getBuffer()->clear();
					return;
				}

				while (fbDataSize > 0) {
					// Write out compressed data in chunks.
					auto chunkSize         = lz4CompressionChunkSize;
					auto maxCompressedSize = compressionLZ4.chunkSize;

					if (static_cast<size_t>(fbDataSize) < lz4CompressionChunkSize) {
						chunkSize         = static_cast<size_t>(fbDataSize);
						maxCompressedSize = LZ4F_compressBound(chunkSize, compressionLZ4.prefs);
					}

					msgBuild->getBuffer()->resize(written + maxCompressedSize);

					auto retVal = LZ4F_compressUpdate(compressionLZ4.context.get(),
						msgBuild->getBuffer()->data() + written, maxCompressedSize, fbData, chunkSize, nullptr);
					if (LZ4F_isError(retVal)) {
						// Compression error.
						log->error << "LZ4 compression error: " << LZ4F_getErrorName(retVal) << std::endl;
						msgBuild->getBuffer()->clear();
						return;
					}
					else {
						written += retVal;
					}

					// Update counters.
					fbData     += lz4CompressionChunkSize;
					fbDataSize -= lz4CompressionChunkSize;
				}

				// Write out end of compressed frame.
				msgBuild->getBuffer()->resize(written + compressionLZ4.endSize);

				auto retVal = LZ4F_compressEnd(compressionLZ4.context.get(), msgBuild->getBuffer()->data() + written,
					compressionLZ4.endSize, nullptr);
				if (LZ4F_isError(retVal)) {
					// Compression error.
					log->error << "LZ4 compression error: " << LZ4F_getErrorName(retVal) << std::endl;
					msgBuild->getBuffer()->clear();
					return;
				}
				else {
					written += retVal;
				}

				// Set final size.
				msgBuild->getBuffer()->resize(written);
			}

			if (compressionType == CompressionType::ZSTD || compressionType == CompressionType::ZSTD_HIGH) {
				// Allocate maximum needed memory for compressed data block.
				auto maxCompressedSize = ZSTD_compressBound(static_cast<size_t>(fbDataSize));
				msgBuild->getBuffer()->resize(maxCompressedSize);

				// Compress data using Zstd algorithm.
				auto retVal = ZSTD_compressCCtx(compressionZstd.context.get(), msgBuild->getBuffer()->data(),
					maxCompressedSize, fbData, static_cast<size_t>(fbDataSize), compressionZstd.level);
				if (ZSTD_isError(retVal)) {
					// Compression error.
					log->error << "Zstd compression error: " << ZSTD_getErrorName(retVal) << std::endl;
					msgBuild->getBuffer()->clear();
					return;
				}

				// Update size.
				msgBuild->getBuffer()->resize(retVal);
			}
		}
	}
};

class ModuleStatistics : public dv::io::support::IOStatistics {
protected:
	dv::OutputEncoder *encoder;

	uint64_t mStatsMultiplier = 1;

public:
	explicit ModuleStatistics(OutputEncoder *encoder) : encoder(encoder) {
	}

	void setStatsMultiplier(const uint64_t statsMultiplier) {
		mStatsMultiplier = statsMultiplier;
	}

	void publish() override {
		encoder->addStatistics(static_cast<int64_t>(mDataSize * mStatsMultiplier),
			static_cast<int64_t>(mPacketsNumber * mStatsMultiplier),
			static_cast<int64_t>(mPacketsElements * mStatsMultiplier),
			static_cast<int64_t>(mPacketsSize * mStatsMultiplier));
	}
};

} // namespace dv

#endif // DV_OUTPUT_HPP
