#include "../../include/dv-sdk/module.hpp"

#include <boost/filesystem.hpp>
#include <boost/tokenizer.hpp>

#include <thread>

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/imgutils.h>
#include <libavutil/opt.h>
#include <libswscale/swscale.h>
}

class VideoOutput : public dv::ModuleBase {
private:
	// YUV420P seems to have more universal hardware and software
	// support than our own BGR24 for encoding operations.
	static constexpr enum AVPixelFormat PIXEL_FORMAT = AV_PIX_FMT_YUV420P;
	static constexpr auto OPTIONS_SEPARATOR          = ";";
	static constexpr auto OPTIONS_KEYVALUE_SEPARATOR = ":";

	static std::vector<std::string> findEncodersForCodec(const std::string &prefix, const enum AVCodecID id) {
		std::vector<std::string> foundEncoderNames;
		std::vector<std::string> foundEncoderLongNames;

		const AVCodec *currCodec = nullptr;

#if (LIBAVCODEC_VERSION_MAJOR < 58) || (LIBAVCODEC_VERSION_MAJOR == 58 && LIBAVCODEC_VERSION_MINOR < 7)
		avcodec_register_all();

		while ((currCodec = av_codec_next(currCodec)) != nullptr)
#else
		void *iterationState = nullptr;

		while ((currCodec = av_codec_iterate(&iterationState)) != nullptr)
#endif
		{
			if (currCodec->type != AVMEDIA_TYPE_VIDEO) {
				continue;
			}

			if (av_codec_is_encoder(currCodec) == 0) {
				continue;
			}

			if (currCodec->id != id) {
				continue;
			}

			if (currCodec->pix_fmts == nullptr) {
				continue;
			}

			bool compatPixelFormat = false;

			for (size_t i = 0;; i++) {
				const auto pxFormat = currCodec->pix_fmts[i];

				// Detect end of array.
				if (pxFormat == AV_PIX_FMT_NONE) {
					break;
				}

				// If compatible pixel format found, we're done.
				if (pxFormat == PIXEL_FORMAT) {
					compatPixelFormat = true;
					break;
				}
			}

			if (!compatPixelFormat) {
				continue;
			}

			if ((currCodec->name == nullptr) || (currCodec->long_name == nullptr)) {
				continue;
			}

			const auto encoderName     = std::string{currCodec->name};
			const auto encoderLongName = std::string{currCodec->long_name};

			// Check that name doesn't exist in long and short name.
			// Some encoders (Nvidia) have multiple short names but are the
			// same, with same long name, and we don't want them to appear
			// multiple times in the final list.
			if (dv::vectorContains(foundEncoderLongNames, encoderLongName)) {
				continue;
			}

			foundEncoderLongNames.push_back(encoderLongName);

			if (dv::vectorContains(foundEncoderNames, encoderName)) {
				continue;
			}

			foundEncoderNames.push_back(encoderName);
		}

		// Change string representation of all encoders to be 'prefix_[encoder]'.
		std::transform(foundEncoderNames.cbegin(), foundEncoderNames.cend(), foundEncoderNames.begin(),
			[prefix](const std::string &s) {
				return fmt::format("{}_[{}]", prefix, s);
			});

		std::sort(foundEncoderNames.begin(), foundEncoderNames.end());

		return foundEncoderNames;
	}

	boost::filesystem::path path;
	AVCodecContext *codecCtx = nullptr;
	AVFormatContext *fmtCtx  = nullptr;
	AVStream *stream         = nullptr;
	AVFrame *frameOriginal   = nullptr;
	AVFrame *frameEncoding   = nullptr;
	SwsContext *swsCtx       = nullptr;
	AVPacket pkt;
	size_t frameCount = 0;

	void avSetOption(void *object, const std::string &key, const std::string &value) {
		int result = av_opt_set(object, key.c_str(), value.c_str(), 0);
		if (result != 0) {
			// Error! Throw exception to stop initialization.
			char errorMsg[AV_ERROR_MAX_STRING_SIZE];
			av_strerror(result, errorMsg, AV_ERROR_MAX_STRING_SIZE);

			throw std::runtime_error{fmt::format(
				"Failed to set option '{}' to value '{}', error '{}' (code {}).", key, value, errorMsg, result)};
		}
		else {
			log.debug.format("Set option '{}' to value '{}'.", key, value);
		}
	}

	void encodeFrame(const AVFrame *const frame) {
		// Send the frame to be encoded.
		int ret = avcodec_send_frame(codecCtx, frame);
		if (ret < 0) {
			throw std::runtime_error{"Could not send frame to be encoded."};
		}

	tryReceivePacket:
		av_init_packet(&pkt);

		// Packet data will be allocated by the encoder.
		pkt.data = nullptr;
		pkt.size = 0;

		// Get the encoded data back if there is any.
		ret = avcodec_receive_packet(codecCtx, &pkt);
		if (ret < 0) {
			// Handle errors, some are expected.
			switch (ret) {
				case AVERROR_EOF:
					// Encoder done, all data flushed.
					log.debug("Encoder EOF reached.");
					break;

				case AVERROR(EAGAIN):
					// Need more data sent in.
					break;

				default:
					throw std::runtime_error{"Could not retrieve encoded frame."};
			}
		}
		else {
			// Write the data in the packet to the output format.
			av_packet_rescale_ts(&pkt, codecCtx->time_base, stream->time_base);
			ret = av_interleaved_write_frame(fmtCtx, &pkt);

			// Reset the packet.
			av_packet_unref(&pkt);

			// Fail on error.
			if (ret < 0) {
				throw std::runtime_error{"Could not write frame to output."};
			}

			goto tryReceivePacket;
		}
	}

public:
	static void initInputs(dv::InputDefinitionList &in) {
		in.addFrameInput("frames");
	}

	static const char *initDescription() {
		return ("Make a video file out of frames.");
	}

	static void initConfigOptions(dv::RuntimeConfig &config) {
		// Check if encoders exist on this system for the codecs we want
		// to possibly support. H264 preferred due to encode speed and
		// wide support on playback devices.
		const auto ffv1Encoders = findEncodersForCodec("FFV1", AV_CODEC_ID_FFV1);
		const auto h264Encoders = findEncodersForCodec("H264", AV_CODEC_ID_H264);
		const auto hevcEncoders = findEncodersForCodec("HEVC", AV_CODEC_ID_HEVC);
		const auto vp9Encoders  = findEncodersForCodec("VP9", AV_CODEC_ID_VP9);
		const auto av1Encoders  = findEncodersForCodec("AV1", AV_CODEC_ID_AV1);

		std::vector<std::string> encoders;
		encoders.insert(encoders.end(), ffv1Encoders.cbegin(), ffv1Encoders.cend());
		encoders.insert(encoders.end(), h264Encoders.cbegin(), h264Encoders.cend());
		encoders.insert(encoders.end(), hevcEncoders.cbegin(), hevcEncoders.cend());
		encoders.insert(encoders.end(), vp9Encoders.cbegin(), vp9Encoders.cend());
		encoders.insert(encoders.end(), av1Encoders.cbegin(), av1Encoders.cend());

		config.add("codec", dv::ConfigOption::listOption("How to encode the video data.", 0, encoders));
		config.add("codecOptions",
			dv::ConfigOption::stringOption("Pass custom list of options to encoder if set. Different options are "
										   "separated by ';', an option's name and value are separated by ':'.",
				""));

		config.add("fileName", dv::ConfigOption::fileSaveOption("Where to save the output video file.", "mkv"));
		config.add("fps", dv::ConfigOption::intOption("Frame rate of resulting video (in FPS).", 30, 1, 10000));

		config.add("writtenFramesCount",
			dv::ConfigOption::statisticOption("Number of frames written into video file so far."));

		config.setPriorityOptions({"fileName", "fps", "codec", "writtenFramesCount"});
	}

	VideoOutput() {
		path = boost::filesystem::path(config.getString("fileName"));

		if (path.empty() || path.parent_path().empty() || path.filename().empty()) {
			throw std::runtime_error("Output file path not specified.");
		}

		if (path.extension().empty()) {
			throw std::runtime_error("File extension not specified.");
		}

		if (path.extension().string() != ".mkv") {
			throw std::runtime_error(
				"File extension is not MKV (Matroska), given wrong extension is: " + path.extension().string());
		}

		if (!boost::filesystem::exists(path.parent_path()) || !boost::filesystem::is_directory(path.parent_path())) {
			throw std::runtime_error(
				"Output directory doesn't exist or is not a directory: " + path.parent_path().string());
		}

		const auto fps = config.getInt("fps");
		const auto url = "file:" + path.string();

		// Initialize the codec used to encode our given images.
		auto encoder = config.getString("codec");

		encoder.pop_back(); // Remove closing ']'.

		// Extract encoder name, part after '_['.
		const auto encoderName = encoder.substr(encoder.find("_[") + 2);

		const AVCodec *const codec = avcodec_find_encoder_by_name(encoderName.c_str());
		if (codec == nullptr) {
			throw std::runtime_error{"Could not find codec."};
		}

		codecCtx = avcodec_alloc_context3(codec);
		if (codecCtx == nullptr) {
			throw std::runtime_error{"Could not allocate codec context."};
		}

		// Codec configuration: frame size.
		codecCtx->width  = inputs.getFrameInput("frames").sizeX();
		codecCtx->height = inputs.getFrameInput("frames").sizeY();

		// Codec configuration: fps.
		codecCtx->time_base.num = 1;
		codecCtx->time_base.den = fps;

		// Codec configuration: frame format.
		codecCtx->pix_fmt    = PIXEL_FORMAT;
		codecCtx->codec_type = AVMEDIA_TYPE_VIDEO;

		// Codec configuration: threads.
		codecCtx->thread_count = std::thread::hardware_concurrency() / 2;

		// Per-codec options, try to get good quality and performance.
		// If custom options are specified by the user, always prefer those.
		// Else try to find a match against the encoders we know and support.
		auto encoderOptions = config.getString("codecOptions");

		if (!encoderOptions.empty()) {
			boost::tokenizer<boost::char_separator<char>> options(
				encoderOptions, boost::char_separator<char>(OPTIONS_SEPARATOR, nullptr));

			for (const auto &option : options) {
				// Parse custom option to extract key:value pair.
				const auto separatorPosition = option.find(OPTIONS_KEYVALUE_SEPARATOR);
				if (separatorPosition == std::string::npos) {
					// Ignore malformed option.
					log.warning.format("Encoder option '{}' is malformed, no '{}' separator found.", option,
						OPTIONS_KEYVALUE_SEPARATOR);
				}

				const auto key   = option.substr(0, separatorPosition);
				const auto value = option.substr(separatorPosition + 1);

				// Some keys are set directly on codecCtx.
				if ((key == "bit_rate") || (key == "rc_max_rate") || (key == "rc_min_rate") || (key == "rc_buffer_size")
					|| (key == "gop_size") || (key == "max_b_frames")) {
					avSetOption(codecCtx, key, value);
				}
				else {
					// The others we assume are for the specific encoder.
					avSetOption(codecCtx->priv_data, key, value);
				}
			}
		}
		else if (encoderName == "libx264") {
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "preset", "fast");
			avSetOption(codecCtx->priv_data, "crf", "18"); // Constant-quality level, 0 to 51.
		}
		else if (encoderName == "libopenh264") {
			// Cisco Open H.264.
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "rc_mode", "quality"); // Rate-control mode.
		}
		else if ((encoderName == "h264_nvenc") || (encoderName == "nvenc_h264")) {
			// Nvidia GPU encoding.
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "preset", "fast");
			avSetOption(codecCtx->priv_data, "rc", "vbr_hq"); // Variable bit-rate high-quality mode.
			avSetOption(codecCtx->priv_data, "cq", "18");     // Quality level in VBR, 0 to 51.
		}
		else if (encoderName == "h264_mf") {
			// Windows Media Foundation.
			avSetOption(codecCtx->priv_data, "scenario", "camera_record");
			avSetOption(codecCtx->priv_data, "rate_control", "quality"); // Rate-control mode.
			avSetOption(codecCtx->priv_data, "quality", "35");           // Quality level, 0 to 100.
		}
		else if (encoderName == "h264_videotoolbox") {
			// MacOS Video Toolbox.
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "realtime", "true");
		}
		else if (encoderName == "libx265") {
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "preset", "fast");
			avSetOption(codecCtx->priv_data, "crf", "18"); // Constant-quality level, 0 to 51.
		}
		else if ((encoderName == "hevc_nvenc") || (encoderName == "nvenc_hevc")) {
			// Nvidia GPU encoding.
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "preset", "fast");
			avSetOption(codecCtx->priv_data, "rc", "vbr_hq"); // Variable bit-rate high-quality mode.
			avSetOption(codecCtx->priv_data, "cq", "18");     // Quality level in VBR, 0 to 51.
		}
		else if (encoderName == "hevc_mf") {
			// Windows Media Foundation.
			avSetOption(codecCtx->priv_data, "scenario", "camera_record");
			avSetOption(codecCtx->priv_data, "rate_control", "quality"); // Rate-control mode.
			avSetOption(codecCtx->priv_data, "quality", "35");           // Quality level, 0 to 100.
		}
		else if (encoderName == "hevc_videotoolbox") {
			// MacOS Video Toolbox.
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "realtime", "true");
		}
		else if (encoderName == "libaom-av1") {
			avSetOption(codecCtx->priv_data, "usage", "realtime");
			avSetOption(codecCtx->priv_data, "cpu-used", "8"); // Quality/Speed ratio, 0 to 8.
			avSetOption(codecCtx->priv_data, "crf", "20");     // Constant-quality level, 0 to 63.
			codecCtx->bit_rate = 0;                            // For constant-quality mode.
		}
		else if (encoderName == "librav1e") {
			// Rav1e Rust AV1 encoder.
			avSetOption(codecCtx->priv_data, "speed", "10"); // Speed preset, 0 to 10.
		}
		else if (encoderName == "libvpx-vp9") {
			avSetOption(codecCtx->priv_data, "deadline", "realtime");
			avSetOption(codecCtx->priv_data, "cpu-used", "8"); // Target 50% CPU, (100*(16-VALUE)/16)%.
			avSetOption(codecCtx->priv_data, "crf", "20");     // Constant-quality level, 0 to 63.
			codecCtx->bit_rate = 0;                            // For constant-quality mode.
		}

		// Initialize the format context.
#if (LIBAVCODEC_VERSION_MAJOR >= 59)
		const
#endif
			AVOutputFormat *const format = av_guess_format("matroska", nullptr, nullptr);
		if (format == nullptr) {
			throw std::runtime_error{"Could not find format."};
		}

		int ret = avformat_alloc_output_context2(&fmtCtx, format, format->name, url.c_str());
		if (ret < 0) {
			throw std::runtime_error{"Could not allocate output format context."};
		}

		// Some formats such as mkv require stream headers to be separate (not for all codecs).
		if ((fmtCtx->oformat->flags & AVFMT_GLOBALHEADER) != 0) {
			log.debug("Output format requires global header.");

			codecCtx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
		}

		// FFV1 is used as lossless archival format, so try to not have parts of
		// the file, such as the headers, change randomly.
		if (encoderName == "ffv1") {
			log.debug("Enabling bitexact encoding.");

			fmtCtx->flags   |= AVFMT_FLAG_BITEXACT;
			codecCtx->flags |= AV_CODEC_FLAG_BITEXACT;
		}

		ret = avcodec_open2(codecCtx, codec, nullptr);
		if (ret < 0) {
			throw std::runtime_error{"Could not initialize codec context."};
		}

		swsCtx = sws_getCachedContext(swsCtx, codecCtx->width, codecCtx->height, AV_PIX_FMT_GRAY8, codecCtx->width,
			codecCtx->height, codecCtx->pix_fmt, SWS_BICUBIC, nullptr, nullptr, nullptr);
		if (swsCtx == nullptr) {
			throw std::runtime_error{"Could not allocate scaling context."};
		}

		// Initialize the frame containing our raw data.
		frameOriginal = av_frame_alloc();
		if (frameOriginal == nullptr) {
			throw std::runtime_error{"Could not allocate original frame."};
		}

		// Initialize the frame containing the data for encoding.
		frameEncoding = av_frame_alloc();
		if (frameEncoding == nullptr) {
			throw std::runtime_error{"Could not allocate encoding frame."};
		}

		frameEncoding->format = codecCtx->pix_fmt;
		frameEncoding->width  = codecCtx->width;
		frameEncoding->height = codecCtx->height;

		ret = av_image_alloc(frameEncoding->data, frameEncoding->linesize, frameEncoding->width, frameEncoding->height,
			codecCtx->pix_fmt, 1);
		if (ret < 0) {
			throw std::runtime_error{"Could not allocate encoding frame image data."};
		}

		// Open URL for writing video data.
#if (LIBAVFORMAT_VERSION_MAJOR < 58) || (LIBAVFORMAT_VERSION_MAJOR == 58 && LIBAVFORMAT_VERSION_MINOR < 7)
		ret = avio_open(&fmtCtx->pb, fmtCtx->filename, AVIO_FLAG_WRITE);
#else
		ret = avio_open(&fmtCtx->pb, fmtCtx->url, AVIO_FLAG_WRITE);
#endif
		if (ret < 0) {
			throw std::runtime_error{"Could not open file resource."};
		}

		// Configure the AVStream for the output format context.
		stream = avformat_new_stream(fmtCtx, codec);
		if (stream == nullptr) {
			throw std::runtime_error{"Could not allocate stream."};
		}

		ret = avcodec_parameters_from_context(stream->codecpar, codecCtx);
		if (ret < 0) {
			throw std::runtime_error{"Could not initialize codec parameters."};
		}

		// Set timebase again for muxer.
		stream->time_base.num = 1;
		stream->time_base.den = fps;

		// Rewrite the header.
		ret = avformat_write_header(fmtCtx, nullptr);
		if (ret < 0) {
			throw std::runtime_error{"Could not write format header."};
		}
	}

	void run() override {
		const auto inFrame = inputs.getFrameInput("frames").frame();

		// Update scaling context with correct format.
		enum AVPixelFormat framePixelFormat;

		switch (inFrame.format()) {
			case dv::FrameFormat::OPENCV_8U_C3:
				framePixelFormat = AV_PIX_FMT_BGR24;
				break;

			case dv::FrameFormat::OPENCV_8U_C4:
				framePixelFormat = AV_PIX_FMT_BGRA;
				break;

			case dv::FrameFormat::OPENCV_8U_C1:
			default:
				framePixelFormat = AV_PIX_FMT_GRAY8;
				break;
		}

		cv::Mat frame;

		if (inFrame.size() != inputs.getFrameInput("frames").size()) {
			// ROI frame. Only the actual pixels are in data(), but encoders
			// expect a fixed size frame as given during setup. So if ROI is
			// active, we make an original sized black frame and copy the ROI
			// region to it at the right position.
			frame = cv::Mat(inputs.getFrameInput("frames").size(), static_cast<int>(inFrame.format()), cv::Scalar{0});
			(*inFrame.getMatPointer()).copyTo(frame(inFrame.roi()));
		}
		else {
			// Use input frame directly if full size.
			frame = *inFrame.getMatPointer();
		}

		int ret = av_image_fill_arrays(
			frameOriginal->data, frameOriginal->linesize, frame.data, framePixelFormat, frame.cols, frame.rows, 1);
		if (ret < 0) {
			throw std::runtime_error{"Could not fill original frame image data."};
		}

		swsCtx = sws_getCachedContext(swsCtx, frame.cols, frame.rows, framePixelFormat, frameEncoding->width,
			frameEncoding->height, codecCtx->pix_fmt, SWS_BILINEAR, nullptr, nullptr, nullptr);
		if (swsCtx == nullptr) {
			throw std::runtime_error{"Could not get updated scaling context."};
		}

		// Use FFMPEG swscale to convert from original DV format to what the encoder needs.
		sws_scale(swsCtx, frameOriginal->data, frameOriginal->linesize, 0, frame.rows, frameEncoding->data,
			frameEncoding->linesize);

		// Which frame is it?
		frameEncoding->pts = static_cast<long int>(frameCount);
		frameCount++;

		encodeFrame(frameEncoding);

		config.setLong("writtenFramesCount", static_cast<int64_t>(frameCount));
	}

	~VideoOutput() override {
		// End streaming.
		try {
			encodeFrame(nullptr);
		}
		catch (const std::runtime_error &ex) {
			log.error.format("Flushing encoder - Error: '{:s}'.", ex.what());
		}

		int ret = av_write_trailer(fmtCtx);
		if (ret < 0) {
			log.error("Could not write format trailer.");
		}

		// Free all resources.
		ret = avio_close(fmtCtx->pb);
		if (ret < 0) {
			log.error("Could not close file resource.");
		}

		avformat_free_context(fmtCtx);

		avcodec_free_context(&codecCtx);

		av_freep(&frameEncoding->data[0]);
		av_frame_free(&frameEncoding);

		av_frame_free(&frameOriginal);

		sws_freeContext(swsCtx);

		config.set<dv::CfgType::LONG>("writtenFramesCount", static_cast<int64_t>(frameCount), true);

		log.info << "Video saved in: " << path << std::endl;
	}
};

registerModuleClass(VideoOutput)
