#include "AVTranscoder.h"
extern "C" {
#include <libswscale/swscale.h> 
}
#ifdef WIN32
#pragma comment(lib,"swscale.lib")
#endif

#define SCALED_W 80
#define SCALED_H 60
#define SRC_W 320
#define SRC_H 240

CAVTranscoder::CAVTranscoder():
_start_time(INT64_MAX),
_all_processed(true),
_one2one(true),
_nOutputWidth(320),
_cur_out_v_ts(0),
_cur_out_a_ts(0),
_max_audio(1),
_swsCtx(NULL)
{
	if (_one2one) {
		_nOutputHeight = 480;
	}
	else {
		_nOutputHeight = 240;
		_swsCtx = sws_getContext(SRC_W, SRC_H, PIX_FMT_YUV420P,
			SCALED_W, SCALED_H, PIX_FMT_YUV420P, SWS_BILINEAR,
			NULL, NULL, NULL);
	}
}


CAVTranscoder::~CAVTranscoder()
{
}

int CAVTranscoder::add(media_info & info)
{
	_all_processed = false;
	if (_start_time == INT64_MAX) {
		_start_time = info.start_time_ms;
		_cur_v_time = _start_time;
		_cur_a_time = _start_time;
	}
	vector < CAVDecoder *>::iterator it = _decoders.begin();
	for (; it != _decoders.end(); it++) {
		if ((*it)->getuid() == info.uid){
			(*it)->add(info);
			break;
		}
	}
	if (it == _decoders.end()) {
		CAVDecoder * pVideoDecoder = new CAVDecoder();
		pVideoDecoder->add(info);
		_decoders.push_back(pVideoDecoder);
	}
	return 0;
}

int64_t CAVTranscoder::transcode()
{
	vector<CAVDecoder *> decoders_got_frame;
	vector <CAVDecoder *>::iterator it = _decoders.begin();
	for (; it != _decoders.end();) {
		if((*it)->get_one_v_frame()){
			decoders_got_frame.push_back(*it);
		}
		else {
			it = _decoders.erase(it);
			continue;
		}
		it++;
	}
	
	_all_processed = decoders_got_frame.size() == 0;
	mix_and_output_vframe(decoders_got_frame);

	_cur_v_time += VFRAME_DURATION_MS;
	//sws_scale(img_convert_ctx, (const uint8_t* const*)pFrame->data, pFrame->linesize, 0, pCodecCtx->height, pFrameYUV->data, pFrameYUV->linesize);
	
	while (_cur_a_time < _cur_v_time)
	{
		decoders_got_frame.clear();
		vector < CAVDecoder *>::iterator it = _decoders.begin();
		for (; it != _decoders.end();) {
			if ((*it)->get_one_a_frame()){
				decoders_got_frame.push_back(*it);
			}
			else {
				it = _decoders.erase(it);
				continue;
			}
			it++;
		}
		mix_and_output_aframe(decoders_got_frame);
		_cur_a_time += AFRAME_DURATION_MS;
	}
	
	return _cur_v_time;
}

bool CAVTranscoder::all_processed()
{
	return _all_processed;
}

int CAVTranscoder::close()
{
	if (_swsCtx) {
		sws_freeContext(_swsCtx);
		_swsCtx = NULL;
	}
	flush_encoder(0);
	flush_encoder(1);
	av_write_trailer(_ofmt_ctx);

#if USE_H264BSF  
	av_bitstream_filter_close(h264bsfc);
#endif  
#if USE_AACBSF  
	av_bitstream_filter_close(aacbsfc);
#endif  
	int i;
	for (i = 0; i<2; i++)
	{
		if (_ofmt_ctx && _ofmt_ctx->nb_streams > i && _ofmt_ctx->streams[i] && _ofmt_ctx->streams[i]->codec)
			avcodec_close(_ofmt_ctx->streams[i]->codec);
	}

	if (_ofmt_ctx && !(_ofmt_ctx->oformat->flags & AVFMT_NOFILE))
		avio_close(_ofmt_ctx->pb);
	avformat_free_context(_ofmt_ctx);

	return 0;
}


int CAVTranscoder::open_output_file(const char *filename)
{
	AVStream *out_stream;
	AVCodecContext *enc_ctx;
	AVCodec *encoder;
	int ret;
	unsigned int i;

	_ofmt_ctx = NULL;
	avformat_alloc_output_context2(&_ofmt_ctx, NULL, NULL, filename);
	if (!_ofmt_ctx) {
		av_log(NULL, AV_LOG_ERROR, "Could not create output context\n");
		return AVERROR_UNKNOWN;
	}

	for (i = 0; i < 2; i++) {
		out_stream = avformat_new_stream(_ofmt_ctx, NULL);
		if (!out_stream) {
			av_log(NULL, AV_LOG_ERROR, "Failed allocating output stream\n");
			return AVERROR_UNKNOWN;
		}

		enc_ctx = out_stream->codec;

		if (_ofmt_ctx->oformat->flags & AVFMT_GLOBALHEADER)
			enc_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;

		if (0 == i) {
			encoder = avcodec_find_encoder(AV_CODEC_ID_H264);;
			if (!encoder) {
				av_log(NULL, AV_LOG_FATAL, "Necessary encoder not found\n");
				return AVERROR_INVALIDDATA;
			}

			/* In this example, we transcode to same properties (picture size,
			* sample rate etc.). These properties can be changed for output
			* streams easily using filters */
				enc_ctx->height = _nOutputHeight;
				enc_ctx->width = _nOutputWidth;
				enc_ctx->sample_aspect_ratio.den = 1;
				enc_ctx->sample_aspect_ratio.num = 0;
				/* take first format from list of supported formats */
				enc_ctx->pix_fmt = AV_PIX_FMT_YUV420P;
				/* video time_base can be set to whatever is handy and supported by encoder */
				enc_ctx->time_base.num = 1;
				enc_ctx->time_base.den = 20;

				enc_ctx->me_range = 16;
				enc_ctx->max_qdiff = 4;
				enc_ctx->qmin = 10;
				enc_ctx->qmax = 30;
				enc_ctx->qcompress = 0.6;
				enc_ctx->framerate.den = 20;
				enc_ctx->framerate.num = 1;

				AVDictionary * d = NULL;
				char *k = av_strdup("preset");       // if your strings are already allocated,
				char *v = av_strdup("ultrafast");    // you can avoid copying them like this
				av_dict_set(&d, k, v, AV_DICT_DONT_STRDUP_KEY | AV_DICT_DONT_STRDUP_VAL);
				/* Third parameter can be used to pass settings to encoder */
				ret = avcodec_open2(enc_ctx, encoder, NULL);
				if (ret < 0) {
					av_log(NULL, AV_LOG_ERROR, "Cannot open video encoder for stream #%u\n", i);
					return ret;
				}
			}
			else {
				encoder = avcodec_find_encoder(AV_CODEC_ID_AAC);;
				if (!encoder) {
					av_log(NULL, AV_LOG_FATAL, "Necessary encoder not found\n");
					return AVERROR_INVALIDDATA;
				}
				enc_ctx->sample_rate = 48000;
				enc_ctx->channel_layout = AV_CH_LAYOUT_MONO;
				enc_ctx->channels = av_get_channel_layout_nb_channels(AV_CH_LAYOUT_MONO);
				/* take first format from list of supported formats */
				enc_ctx->sample_fmt = AV_SAMPLE_FMT_S16; //AV_SAMPLE_FMT_FLTP;
				enc_ctx->time_base.num = 1;
				enc_ctx->time_base.den = enc_ctx->sample_rate;
				enc_ctx->bit_rate = 64000;
				/* Third parameter can be used to pass settings to encoder */
				ret = avcodec_open2(enc_ctx, encoder, NULL);
				if (ret < 0) {
					av_log(NULL, AV_LOG_ERROR, "Cannot open video encoder for stream #%u\n", i);
					return ret;
				}
			}
	}

	av_dump_format(_ofmt_ctx, 0, filename, 1);

	if (!(_ofmt_ctx->oformat->flags & AVFMT_NOFILE)) {
		ret = avio_open(&_ofmt_ctx->pb, filename, AVIO_FLAG_WRITE);
		if (ret < 0) {
			av_log(NULL, AV_LOG_ERROR, "Could not open output file '%s'", filename);
			return ret;
		}
	}

	/* init muxer, write output file header */
	ret = avformat_write_header(_ofmt_ctx, NULL);
	if (ret < 0) {
		av_log(NULL, AV_LOG_ERROR, "Error occurred when opening output file\n");
		return ret;
	}

	return 0;
}


	int CAVTranscoder::mix_and_output_vframe(vector<CAVDecoder *> & decoders_got_frame)
	{
		if (_one2one){
			return mix_and_output_one2one_vframe(decoders_got_frame);
		}
		else {
			return  mix_and_output_one2many_vframe(decoders_got_frame);
		}
	}

	int CAVTranscoder::mix_and_output_aframe(vector<CAVDecoder *> & decoders_got_frame)
	{
		AVFrame *pDstFrame = av_frame_alloc();
		pDstFrame->nb_samples = 1024;
		pDstFrame->channel_layout = AV_CH_LAYOUT_MONO;
		pDstFrame->channels = av_get_channel_layout_nb_channels(AV_CH_LAYOUT_MONO);
		pDstFrame->format = AV_SAMPLE_FMT_S16;
		pDstFrame->sample_rate = 48000;

		av_frame_get_buffer(pDstFrame, 0);
		av_samples_set_silence(pDstFrame->data, 0, 1024, pDstFrame->channels, (AVSampleFormat)pDstFrame->format);
		vector < CAVDecoder *>::iterator it = decoders_got_frame.begin();
		for (; it != decoders_got_frame.end(); it++) {
			AVFrame * pFrame = (*it)->_cur_a_frame;
			if (pFrame) {
				int16_t * psrc = (int16_t *)pFrame->extended_data[0];
				int16_t * pdst = (int16_t *)pDstFrame->extended_data[0];
				for (int i = 0; i < 1024; i++,pdst++,psrc++) {
					*pdst += (*psrc/_max_audio);
				}
				(*it)->free_cur_a_frame();
			}
		}


		pDstFrame->pts = _cur_out_a_ts;
		pDstFrame->pkt_dts = _cur_out_a_ts;
		pDstFrame->pkt_pts = _cur_out_a_ts;
		pDstFrame->pkt_duration = 1024;
		_cur_out_a_ts += 1024;

		int got_frame = 0;
		encode_write_frame(pDstFrame, 1, &got_frame);

		return 0;
	}

	int CAVTranscoder::mix_and_output_one2many_vframe(vector<CAVDecoder *> & decoders_got_frame)
	{
		int idxTeacher = -1;
		for (int i = 0; i < decoders_got_frame.size(); i++){
			if (decoders_got_frame[i]->_media_role == mr_teacher) {
				idxTeacher = i;
				break;
			}
		}

		AVFrame *pDstFrame = av_frame_alloc();
		int nDstSize = avpicture_get_size(AV_PIX_FMT_YUV420P, _nOutputWidth, _nOutputHeight);
		uint8_t *dstbuf = new uint8_t[nDstSize];
		avpicture_fill((AVPicture*)pDstFrame, dstbuf, AV_PIX_FMT_YUV420P, _nOutputWidth, _nOutputHeight);

		if (idxTeacher != -1) {
			//copy teacher frame to dest frame
			CAVDecoder * pDecoder = decoders_got_frame[idxTeacher];
			AVFrame * pFrame = pDecoder->_cur_v_frame;
			if (pFrame) {
				fillDestFrame(pDstFrame, pFrame, 0, 0);
			}
			else {//fill with pure color
				memset(pDstFrame->data[0], 0, _nOutputWidth * _nOutputHeight);
				memset(pDstFrame->data[1], 0x80, _nOutputWidth *_nOutputHeight / 4);
				memset(pDstFrame->data[2], 0x80, _nOutputWidth * _nOutputHeight / 4);
			}

			for (int i = 0; i < decoders_got_frame.size(); i++){
				if (i != idxTeacher) {
					//scale eacher frame

					//copy each frame to the dest frame

				}
			}
		}
		else {
			memset(pDstFrame->data[0], 0, _nOutputWidth * _nOutputHeight);
			memset(pDstFrame->data[1], 0x80, _nOutputWidth *_nOutputHeight / 4);
			memset(pDstFrame->data[2], 0x80, _nOutputWidth * _nOutputHeight / 4);

				for (int i = 0; i < decoders_got_frame.size(); i++){
					if (i != idxTeacher) {
						//scale eacher frame

						//copy each frame to the dest frame

					}
				}
		}

		//fill the timestamp of dest frame
		pDstFrame->pts = _cur_out_v_ts;
		pDstFrame->pkt_dts = _cur_out_v_ts;
		pDstFrame->pkt_pts = _cur_out_v_ts;
		pDstFrame->format = AV_PIX_FMT_YUV420P;
		pDstFrame->width = _nOutputWidth;
		pDstFrame->height = _nOutputHeight;
		_cur_out_v_ts++;

		//send to encoder
		int got_frame = 0;
		encode_write_frame(pDstFrame, 0, &got_frame);
		delete dstbuf;

		return 0;
	}

	int CAVTranscoder::fillDestFrame(AVFrame * pDstFrame, AVFrame * pSrcFrame, int x, int y)
	{
		if (!pSrcFrame){
			return 0;
		}

		for (int i = 0; i < pSrcFrame->height; i++)	{
			memcpy(pDstFrame->data[0] + (y + i)*pDstFrame->linesize[0] + x, pSrcFrame->data[0] + i * pSrcFrame->linesize[0], pSrcFrame->linesize[0]>0 ? pSrcFrame->linesize[0] : -pSrcFrame->linesize[0]);
		}

		for (int i = 0; i < pSrcFrame->height / 2; i++){
			memcpy(pDstFrame->data[1] + (y / 2 + i)*pDstFrame->linesize[1] + x / 2, pSrcFrame->data[1] + i * pSrcFrame->linesize[1], pSrcFrame->linesize[1]>0 ? pSrcFrame->linesize[1] : -pSrcFrame->linesize[1]);
			memcpy(pDstFrame->data[2] + (y / 2 + i)*pDstFrame->linesize[2] + x / 2, pSrcFrame->data[2] + i * pSrcFrame->linesize[2], pSrcFrame->linesize[2]>0 ? pSrcFrame->linesize[2] : -pSrcFrame->linesize[2]);
		}
		return 0;
	}

	int CAVTranscoder::mix_and_output_one2one_vframe(vector<CAVDecoder *> & decoders_got_frame)
	{
		//prepare one2one base frame
		AVFrame *pDstFrame = av_frame_alloc();
		int nDstSize = avpicture_get_size(AV_PIX_FMT_YUV420P,_nOutputWidth, _nOutputHeight);
		uint8_t *dstbuf = new uint8_t[nDstSize];
		avpicture_fill((AVPicture*)pDstFrame, dstbuf, AV_PIX_FMT_YUV420P, _nOutputWidth, _nOutputHeight);
		memset(dstbuf, 0x80, nDstSize);

		if (decoders_got_frame.size() == 2){
			fillDestFrame(pDstFrame, decoders_got_frame[0]->_cur_v_frame, 0, decoders_got_frame[0]->_media_role == mr_teacher ? 0 : 240);
			decoders_got_frame[0]->free_cur_v_frame();
			fillDestFrame(pDstFrame, decoders_got_frame[1]->_cur_v_frame, 0, decoders_got_frame[1]->_media_role == mr_teacher ? 0 : 240);
			decoders_got_frame[1]->free_cur_v_frame();
		}
		else if (decoders_got_frame.size() == 1)
		{
			fillDestFrame(pDstFrame, decoders_got_frame[0]->_cur_v_frame, 0, 0);
			decoders_got_frame[0]->free_cur_v_frame();
			//todo: fill the bottom half image with pure color
		}
		else {
			//fill with last image?
		}

		//fill the timestamp of dest frame

		pDstFrame->pts = _cur_out_v_ts;
		pDstFrame->pkt_dts = _cur_out_v_ts;
		pDstFrame->pkt_pts = _cur_out_v_ts;
		pDstFrame->format = AV_PIX_FMT_YUV420P;
		pDstFrame->width = _nOutputWidth;
		pDstFrame->height = _nOutputHeight;
		_cur_out_v_ts++;

		//send to encoder
		int got_frame = 0;
		encode_write_frame(pDstFrame, 0, &got_frame);
		delete dstbuf;

		return 0;
	}

int CAVTranscoder::encode_write_frame(AVFrame *filt_frame, unsigned int stream_index, int *got_frame) {
	int ret;
	int got_frame_local;
	AVPacket enc_pkt;

	int(*enc_func)(AVCodecContext *, AVPacket *, const AVFrame *, int *) =
		stream_index == 0 ? avcodec_encode_video2 : avcodec_encode_audio2;

	if (!got_frame)
		got_frame = &got_frame_local;

	/* encode filtered frame */
	enc_pkt.data = NULL;
	enc_pkt.size = 0;
	av_init_packet(&enc_pkt);
	ret = enc_func(_ofmt_ctx->streams[stream_index]->codec, &enc_pkt,
		filt_frame, got_frame);
	av_frame_free(&filt_frame);
	if (ret < 0)
		return ret;
	if (!(*got_frame))
		return 0;

	/* prepare packet for muxing */
	enc_pkt.stream_index = stream_index;
	av_packet_rescale_ts(&enc_pkt,
		_ofmt_ctx->streams[stream_index]->codec->time_base,
		_ofmt_ctx->streams[stream_index]->time_base);

	/* mux encoded frame */
	ret = av_interleaved_write_frame(_ofmt_ctx, &enc_pkt);

	return ret;
}


int CAVTranscoder::flush_encoder(unsigned int stream_index)
{
	int ret;
	int got_frame;

	if (!(_ofmt_ctx->streams[stream_index]->codec->codec->capabilities &
		CODEC_CAP_DELAY))
		return 0;

	while (1) {
		av_log(NULL, AV_LOG_INFO, "Flushing stream #%u encoder\n", stream_index);
		ret = encode_write_frame(NULL, stream_index, &got_frame);
		if (ret < 0)
			break;
		if (!got_frame)
			return 0;
	}
	return ret;
}

void CAVTranscoder::set_max_audio(int max_audio)
{
	_max_audio = max_audio;
}