#include <iostream>
#include "ffcpp/ffcpp.h"
#include "ffcpp/MediaFile.h"

constexpr int VIDEO_STREAM_INDEX = 0;
constexpr int AUDIO_STREAM_INDEX = 1;

namespace ff = ffcpp;

void flushEncoder(ff::MediaFile& file, ff::Codec& encoder, const ff::Stream& inStream, const ff::Stream& outStream, int streamIndex) {
	if(encoder.capabilities() & AV_CODEC_CAP_DELAY) {
		while (1) {
			auto packet = encoder.encode(nullptr);
			if(!packet) break;

			packet.setStreamIndex(streamIndex);
			packet.rescaleTimestamps(inStream.timeBase(), outStream.timeBase());
			file.writePacket(packet);
		}
	}
}

int main(int argc, char** argv) {
	ff::init();
	ff::MediaFile input(argv[1], ff::Mode::Read);
	ff::MediaFile output(argv[2], ff::Mode::Write);

	ff::Stream& vStream = input.videoStream();

	ff::Codec& vDecoder = vStream.codec();
	ff::Stream& outVStream = output.addStream(AV_CODEC_ID_H264, vDecoder.width(), vDecoder.height(), AV_PIX_FMT_YUV420P);
	outVStream.setTimeBase(vDecoder.timeBase());
	ff::Codec& vEncoder = outVStream.codec();
	output.writeHeader();

	int64_t oldPts = 0, oldDts = 0;
	while(auto packet = input.readPacket()) {
		AVMediaType packetType = input.packetType(packet);
		if(packetType == AVMEDIA_TYPE_VIDEO) {
			auto frame = vDecoder.decode(packet);
			frame.setPictureType(AV_PICTURE_TYPE_NONE);
			auto encPacket = vEncoder.encode(frame);
			if(!encPacket) continue;

			encPacket.setStreamIndex(packetType == AVMEDIA_TYPE_VIDEO ? VIDEO_STREAM_INDEX : AUDIO_STREAM_INDEX);

			/*
			// try to recover in case of bad pts/dts
			if(encPacket.pts < encPacket.dts) {
				encPacket.dts = encPacket.pts;
			}

			if(encPacket.pts < oldPts)
				encPacket.pts = oldPts;

			if(encPacket.dts < oldDts)
				encPacket.dts = oldDts;

			oldPts = encPacket.pts;
			oldDts = encPacket.dts;
			 */

			encPacket.rescaleTimestamps(vStream.timeBase(), outVStream.timeBase());
			output.writePacket(encPacket);
		}
	}

	flushEncoder(output, vEncoder, vStream, outVStream, VIDEO_STREAM_INDEX);
	output.writeTrailer();

	return 0;
}