#include "../../include/dv-sdk/module.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 LiveStream : 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;

	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 const char *initDescription() {
		return "Stream live video via FFMPEG.";
	}

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

	static void initConfigOptions(dv::RuntimeConfig &cfg) {
		cfg.add("ip", dv::ConfigOption::stringOption("IP address for stream.", "127.0.0.1"));
		cfg.add("port", dv::ConfigOption::intOption("Port number for stream.", 8554, 1, UINT16_MAX));
		cfg.add("fps", dv::ConfigOption::intOption("Number of frames per second to stream.", 30, 1, 10000));
	}

	LiveStream() {
		const auto fps = config.getInt("fps");

		// See: https://trac.ffmpeg.org/wiki/StreamingGuide#Latency
		// Play with: mpv -hwdec=no --no-cache --untimed --no-demuxer-thread --vd-lavc-threads=1 "rtp://127.0.0.1:8554"

		const auto url = "rtp://" + config.getString("ip") + ":" + std::to_string(config.getInt("port"));

		// Initialize the codec used to encode our given images.
		const AVCodecID codecID = AV_CODEC_ID_H264;

		const AVCodec *const codec = avcodec_find_encoder(codecID);
		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;

		if (codecID == AV_CODEC_ID_H264) {
			avSetOption(codecCtx->priv_data, "profile", "main");
			avSetOption(codecCtx->priv_data, "preset", "ultrafast");
			avSetOption(codecCtx->priv_data, "tune", "zerolatency");
		}

		// Initialize the format context.
#if (LIBAVCODEC_VERSION_MAJOR >= 59)
		const
#endif
			AVOutputFormat *const format = av_guess_format("rtp_mpegts", 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;
		}

		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 resource (FFMPEG built with no streaming support?)."};
		}

		// 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);
	}

	~LiveStream() 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);
	}
};

registerModuleClass(LiveStream)
