#include <iostream>

extern "C" {
	#include <libavformat/avformat.h>
}

void checkResult(int res, const char* msg) {
	if(res < 0) {
		char errStr[260];
		av_strerror(res, errStr, 260);
		std::cout << msg << ": " << errStr << std::endl;
	}
}

int main(int argc, char** argv) {
	AVFormatContext* inFmtCtx = nullptr;
	AVFormatContext* outFmtCtx = nullptr;

	av_register_all();
	int res = avformat_open_input(&inFmtCtx, argv[1], nullptr, nullptr);
	checkResult(res, "cannot open input");

	res = avformat_find_stream_info(inFmtCtx, nullptr);
	checkResult(res, "cannot find stream info");

	res = avformat_alloc_output_context2(&outFmtCtx, nullptr, nullptr, argv[2]);
	checkResult(res, "cannot allocate output context");

	AVStream *videoStream = nullptr, *audioStream = nullptr;
	for(int i = 0; i < inFmtCtx->nb_streams; ++i) {
		AVStream* stream = inFmtCtx->streams[i];
		if(stream->codec->codec_type == AVMEDIA_TYPE_VIDEO) {
			std::cout << "found video stream " << avcodec_get_name(stream->codec->codec_id) << std::endl;
			videoStream = stream;
		} else if(stream->codec->codec_type == AVMEDIA_TYPE_AUDIO) {
			std::cout << "found audio stream " << avcodec_get_name(stream->codec->codec_id) << std::endl;
			audioStream = stream;
		}
	}

	AVCodecContext* codecCtx = videoStream->codec;
	AVCodec* codec = avcodec_find_decoder(codecCtx->codec_id);
	if(!codec) std::cout << "cannot find decoder" << std::endl;

	res = avcodec_open2(codecCtx, codec, nullptr);
	checkResult(res, "cannot open codec for decoding");

	// coder
	AVCodec* dstVCodec = avcodec_find_encoder(AV_CODEC_ID_H264);
	if(!dstVCodec) std::cout << "cannot find encoder" << std::endl;

	AVCodecContext* dstVCodecCtx = avcodec_alloc_context3(dstVCodec);
	if(!dstVCodecCtx) std::cout << "cannot allocate context" << std::endl;

	dstVCodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;
	dstVCodecCtx->width = codecCtx->width;
	dstVCodecCtx->height = codecCtx->height;
	dstVCodecCtx->time_base = codecCtx->time_base;
    outFmtCtx->duration = inFmtCtx->duration;

	res = avcodec_open2(dstVCodecCtx, dstVCodec, nullptr);
	checkResult(res, "cannot open codec for encoding");

	if(outFmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
		dstVCodecCtx->flags |= CODEC_FLAG_GLOBAL_HEADER;

	AVStream* outVideoStream = avformat_new_stream(outFmtCtx, dstVCodec);
	if(!outVideoStream) std::cout << "cannot create output video stream" << std::endl;

	if (!(outFmtCtx->oformat->flags & AVFMT_NOFILE)) {
		res = avio_open(&outFmtCtx->pb, argv[2], AVIO_FLAG_WRITE);
		checkResult(res, "cannot open output file");
	}

	av_dump_format(outFmtCtx, 0, argv[2], 1);

	res = avformat_write_header(outFmtCtx, nullptr);
	checkResult(res, "error writing header to output file");

	AVPacket packet;
	packet.data = nullptr;
	packet.size = 0;
	AVFrame* frame = nullptr;
    int gotPicture = 0, gotPacket = 0, decodedFrames = 1;
	while(true) {
		res = av_read_frame(inFmtCtx, &packet);
		if(res < 0) break;

		frame = av_frame_alloc();
		if(!frame) break;

		AVMediaType packetType = inFmtCtx->streams[packet.stream_index]->codec->codec_type;
		if(packetType == AVMEDIA_TYPE_VIDEO) {
			res = avcodec_decode_video2(codecCtx, frame, &gotPicture, &packet);
			if(res < 0) {
				av_frame_free(&frame);
				break;
			}

			if(gotPicture) {
				frame->pts = av_frame_get_best_effort_timestamp(frame);
				std::cout << "decoded frame: " << decodedFrames++ << std::endl;

				AVPacket encPacket;
				encPacket.data = nullptr;
				encPacket.size = 0;
				av_init_packet(&encPacket);
				encPacket.stream_index = 0;

				encPacket.pts = decodedFrames;
				encPacket.dts = decodedFrames;

				res = avcodec_encode_video2(dstVCodecCtx, &encPacket, frame, &gotPacket);
				av_frame_free(&frame);
				if(res < 0) break;
				if(!gotPacket) continue;

				AVStream *st = outFmtCtx->streams[encPacket.stream_index];
				auto x = st->codec->has_b_frames;
				auto y = st->codec->max_b_frames;

				//av_packet_rescale_ts(&encPacket, st->codec->time_base, st->time_base);

				res = av_interleaved_write_frame(outFmtCtx, &encPacket);
				checkResult(res, "cannot write frame to output file");
				if(res < 0) break;
			}
		}

		av_free_packet(&packet);
	}

	av_write_trailer(outFmtCtx);


	if(codecCtx)
		avcodec_close(codecCtx);

	if(dstVCodecCtx)
		avcodec_close(dstVCodecCtx);

	if(inFmtCtx)
		avformat_close_input(&inFmtCtx);

	avformat_free_context(outFmtCtx);

	return 0;
}