#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::cerr << 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;

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

	AVCodecContext* dstVCodecCtx = outVideoStream->codec;
	dstVCodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;
	dstVCodecCtx->width = codecCtx->width;
	dstVCodecCtx->height = codecCtx->height;
	//dstVCodecCtx->time_base = codecCtx->time_base;
	outVideoStream->time_base = codecCtx->time_base;
	outFmtCtx->duration = inFmtCtx->duration;

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

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

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

	AVFrame* frame = nullptr;
    int gotPicture = 0, gotPacket = 0, decodedFrames = 1;
	while(true) {
		AVPacket packet;
		packet.data = nullptr;
		packet.size = 0;
		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);
                frame->pict_type = AV_PICTURE_TYPE_NONE;
                //std::cout << "decoded frame: " << decodedFrames++ << " pts: " << frame->pts << ", dts: " << frame->pkt_dts << std::endl;

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

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

				av_packet_rescale_ts(&encPacket, videoStream->time_base, outVideoStream->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);
	}

    // flush encoder
    if(dstVCodecCtx->codec->capabilities & AV_CODEC_CAP_DELAY) {
        std::cout << "flushing encoder" << std::endl;
        int gotFrame = 0;
        while (1) {
            AVPacket encPacket;
            encPacket.data = nullptr;
            encPacket.size = 0;
            av_init_packet(&encPacket);
            encPacket.stream_index = 0;
            res = avcodec_encode_video2(dstVCodecCtx, &encPacket, nullptr, &gotFrame);
            if (res < 0) {
	            std::cout << "avcodec_encode_video2 failed" << std::endl;
	            break;
            }
            if (gotFrame) {
                //std::cout << "extra frame" << std::endl;
	            av_packet_rescale_ts(&encPacket, videoStream->time_base, outVideoStream->time_base);
                res = av_interleaved_write_frame(outFmtCtx, &encPacket);
	            checkResult(res, "[flush encoder] cannot write frame to output file");
            } else {
	            break;
            }
        }
    }

    res = av_write_trailer(outFmtCtx);
    checkResult(res, "error writing trailer to output file");

	if(codecCtx)
		avcodec_close(codecCtx);

	if(dstVCodecCtx)
		avcodec_close(dstVCodecCtx);

	if(inFmtCtx)
		avformat_close_input(&inFmtCtx);

	avformat_free_context(outFmtCtx);

	return 0;
}