/* 
*本程式主要實現一個自己的decoder並加入到decoder鏈中去,供api呼叫
*作者：繆國凱（MK） 
*[email protected] 
*2015-6-4 
*/ 

#include "stdafx.h"

#ifdef __cplusplus
extern "C"
{
#endif
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/pixdesc.h>
#
#ifdef __cplusplus
};
#endif

#pragma comment(lib, "avcodec.lib")
#pragma comment(lib, "avformat.lib")
#pragma comment(lib, "avutil.lib")
//#pragma comment(lib, "avdevice.lib")
//#pragma comment(lib, "avfilter.lib")
//#pragma comment(lib, "postproc.lib")
//#pragma comment(lib, "swresample.lib")
//#pragma comment(lib, "swscale.lib")

static av_cold int mk_init_decoder(AVCodecContext *avctx)
{
	return 0;
}

static int mk_decode(AVCodecContext *avctx, void *data, int *got_frame,
	AVPacket *avpkt)
{
	AVFrame   *frame   = (AVFrame*)data;
	AVPicture *picture = (AVPicture*)data;
	const uint8_t *buf             = avpkt->data;
	int buf_size                   = avpkt->size;

	int size = avpicture_get_size(avctx->pix_fmt, avctx->width,
		avctx->height);

	frame->pict_type        = AV_PICTURE_TYPE_I;
	frame->key_frame        = 1;
	
	frame->buf[0] = av_buffer_alloc(size);
	
	memcpy(frame->buf[0]->data, buf, buf_size);

	int res = 0;
	if ((res = avpicture_fill(picture, frame->buf[0]->data, avctx->pix_fmt,
		avctx->width, avctx->height)) < 0) 
	{
			av_buffer_unref(&frame->buf[0]);
			return res;
	}

	*got_frame = 1;
	return 0;
}

static av_cold int mk_close_decoder(AVCodecContext *avctx)
{
	return 0;
}

AVCodec ff_mkvideo_decoder = {
	/*.name           = */"mkvideo",
	/*.long_name      = */"mk video",
	/*.type           = */AVMEDIA_TYPE_VIDEO,
	/*.id             = */AV_CODEC_ID_MKVIDEO,
	/*.capabilities	= */CODEC_CAP_PARAM_CHANGE,
	/*.supported_framerates = */NULL,
	/*.pix_fmts		= */NULL,
	/*.supported_samplerates = */NULL,
	/*.sample_fmts	= */NULL,
	/*.channel_layouts = */NULL,
	/*.max_lowres		= */0,
	/*.priv_class		= */NULL,
	/*.profiles		= */NULL,
	/*.priv_data_size	= */0,
	/*.next			= */NULL,
	/*.init_thread_copy = */NULL,
	/*.update_thread_context = */NULL,
	/*.defaults		= */NULL,
	/*.init_static_data = */NULL,	
	/*.init           = */mk_init_decoder,
	/*.encode_sub		= */NULL,
	/*.encode2        = */NULL,
	/*.decode			= */mk_decode,
	/*.close          = */mk_close_decoder,
};

void help()
{
	printf("**********************************************\n");
	printf("Usage:\n");
	printf("    MyMuxer [inputfile.mk] [outputfile] [size]\n");
	printf("\n");
	printf("Examples: \n");
	printf("    MyMuxer a.mk a.avi 1280x720\n");
	printf("**********************************************\n");  
}


int _tmain(int argc, _TCHAR* argv[])
{
	if(argc < 4|| (!strcmp(argv[1],"--help")))
	{
		help();
		return 0;
	}

	av_register_all();
	avcodec_register(&ff_mkvideo_decoder);

	AVFormatContext *in_fxt = NULL, *out_fxt = NULL;
	AVStream *out_stream = NULL;
	int video_index = -1;

	AVDictionary *param = 0;
	av_dict_set(¶m, "video_size", argv[3], 0);
	if (avformat_open_input(&in_fxt, argv[1], NULL, ¶m) < 0)
	{
		printf("can not open the input file context!\n");
		goto end;
	}
	if (avformat_find_stream_info(in_fxt, NULL) < 0)
	{
		printf("can not find the stream info!\n");
		goto end;
	}

	if(avformat_alloc_output_context2(&out_fxt, NULL, NULL, argv[2]) < 0)
	{
		printf("can not alloc output context!\n");
		goto end;
	}

	for (int i = 0; i < in_fxt->nb_streams; i++)
	{
		if (in_fxt->streams[i]->codec->codec_type == AVMEDIA_TYPE_VIDEO)
		{
			//open decoder
			AVCodec *codec = avcodec_find_decoder(AV_CODEC_ID_MKVIDEO);
			//由於ffmpeg原始碼裡 還沒把mk demuxer的codecid改過來，所以這裡強制賦值,不然avcodec_open2打不開
			in_fxt->streams[i]->codec->codec_id = AV_CODEC_ID_MKVIDEO;
			if(0 > avcodec_open2(in_fxt->streams[i]->codec, codec, NULL))
			{
				printf("can not find or open decoder!\n");
				goto end;
			}
			video_index = i;
			//new stream
			out_stream = avformat_new_stream(out_fxt, NULL);
			if (!out_stream)
			{
				printf("can not new stream for output!\n");
				goto end;
			}
			//set codec context param
			out_stream->codec->codec = avcodec_find_encoder(out_fxt->oformat->video_codec);
			out_stream->codec->height = in_fxt->streams[i]->codec->height;
			out_stream->codec->width = in_fxt->streams[i]->codec->width;

			out_stream->codec->time_base = in_fxt->streams[i]->time_base;
			//out_stream->codec->time_base.den = 25;

			out_stream->codec->sample_aspect_ratio = in_fxt->streams[i]->codec->sample_aspect_ratio;			
			out_stream->codec->pix_fmt = in_fxt->streams[i]->codec->pix_fmt;

			out_stream->avg_frame_rate.den = out_stream->codec->time_base.num;
			out_stream->avg_frame_rate.num = out_stream->codec->time_base.den;
			if (!out_stream->codec->codec)
			{
				printf("can not find the encoder!\n");
				goto end;
			}
			if ((avcodec_open2(out_stream->codec, out_stream->codec->codec, NULL)) < 0)
			{
				printf("can not open the encoder\n");
				goto end;
			}
			if (out_fxt->oformat->flags & AVFMT_GLOBALHEADER)
				out_stream->codec->flags |= CODEC_FLAG_GLOBAL_HEADER;
			break;
		}
	}

	if (-1 == video_index)
	{
		printf("found no video stream in input file!\n");
		goto end;
	}

	if (!(out_fxt->oformat->flags & AVFMT_NOFILE))
	{
		if(avio_open(&out_fxt->pb, argv[2], AVIO_FLAG_WRITE) < 0)
		{
			printf("can not open output file handle!\n");
			goto end;
		}
	}

	if(avformat_write_header(out_fxt, NULL) < 0)
	{
		printf("can not write the header of the output file!\n");
		goto end;
	}

	AVPacket pkt_in, pkt_out;
	
	av_init_packet(&pkt_in);
	av_init_packet(&pkt_out);
	int got_frame, got_picture;
	int i = 0, frame_index = 0;
	while(1)
	{
		AVFrame *frame;
		frame = av_frame_alloc();
		got_frame = -1;
		got_picture = -1;
		if (av_read_frame(in_fxt, &pkt_in) < 0)
		{
			break;
		}frame->buf;
		if (avcodec_decode_video2(in_fxt->streams[video_index]->codec, frame, &got_frame, &pkt_in) < 0)
		{
			printf("can not decoder a frame");
			break;
		}
		av_free_packet(&pkt_in);

		if (got_frame)
		{
			frame->pts = i++;
			pkt_out.data = NULL;//主要這裡必須自己初始化，或者必須置為null，不然ff_alloc_packet2函式會報錯
			pkt_out.size = 0;
			if (avcodec_encode_video2(out_stream->codec, &pkt_out, frame, &got_picture) < 0)
			{
				printf("can not encode a frame!\n");
				break;
			}

			if (got_picture)
			{
				printf("Succeed to encode frame: %5d\tsize:%5d\n",frame_index,pkt_out.size);
				pkt_out.stream_index = out_stream->index;
				frame_index++;
				av_write_frame(out_fxt, &pkt_out);
				av_free_packet(&pkt_out);
			}
		}
		av_frame_free(&frame);
	}
	

	av_write_trailer(out_fxt);

	//clean
	avcodec_close(out_stream->codec);
	avcodec_close(out_fxt->streams[video_index]->codec);
end:
	avformat_close_input(&in_fxt);

	if (out_fxt && !(out_fxt->oformat->flags & AVFMT_NOFILE))
	{
		avio_close(out_fxt->pb);
	}
	avformat_free_context(out_fxt);
	return 0;
}