#include "AVTranscoder.h"


CAVTranscoder::CAVTranscoder():
_start_time(INT64_MAX),
_all_processed(true),
_one2one(false),
_nOutputWidth(320)
{
	if (_one2one) {
		_nOutputHeight = 480;
	}
	else {
		_nOutputHeight = 240;
	}
}


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 += 50;
	//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()
{
	return 0;
}


int CAVTranscoder::open_output_file(const char *filename)
{
	AVStream *out_stream;
	AVCodecContext *dec_ctx, *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 = 1;
				/* 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;
				/* 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;
				/* 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)
	{
		vector < CAVDecoder *>::iterator it = decoders_got_frame.begin();
		for (; it != decoders_got_frame.end(); it++) {
			(*it)->free_cur_a_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;
			}
		}

		if (idxTeacher != -1) {
			// the dest frame is the teacher frame

			for (int i = 0; i < decoders_got_frame.size(); i++){
				if (i != idxTeacher) {
					//scale eacher frame
					//copy each frame to the dest frame
				}
			}

			// set the timestamp of teacher frame

			//send to encoder
		}
		else {
			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(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) {
						idxTeacher = i;
						break;
					}
				}

			//fill the timestamp of dest frame


			//send to encoder
		}

		return 0;
	}

	int CAVTranscoder::fillDestFrame(AVFrame * pDstFrame, AVFrame * pSrcFrame, int x, int y)
	{
		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]);
		}

		for (int i = 0; i < pSrcFrame->height / 2; i++){
			memcpy(pDstFrame->data[1] + (y / 2)*pDstFrame->linesize[1] + x / 2, pSrcFrame->data[1] + i * pSrcFrame->linesize[1], pSrcFrame->linesize[1]);
			memcpy(pDstFrame->data[2] + (y / 2)*pDstFrame->linesize[2] + x / 2, pSrcFrame->data[2] + i * 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);
		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);
			fillDestFrame(pDstFrame, decoders_got_frame[1]->_cur_v_frame, 0, decoders_got_frame[1]->_media_role == mr_teacher ? 0 : 240);
		}
		else {
			fillDestFrame(pDstFrame, decoders_got_frame[0]->_cur_v_frame, 0, 0);
			//todo: fill the bottom half image with pure color
		}

		//fill the timestamp of dest frame


		//send to encoder


		return 0;
	}

int encode_write_frame(AVFrame *filt_frame, unsigned int stream_index, int *got_frame) {
	int ret;
	int got_frame_local;
	AVPacket enc_pkt;
#if 0
	int(*enc_func)(AVCodecContext *, AVPacket *, const AVFrame *, int *) =
		(ifmt_ctx->streams[stream_index]->codec->codec_type ==
		AVMEDIA_TYPE_VIDEO) ? avcodec_encode_video2 : avcodec_encode_audio2;

	if (!got_frame)
		got_frame = &got_frame_local;

	av_log(NULL, AV_LOG_INFO, "Encoding frame\n");
	/* 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);

	av_log(NULL, AV_LOG_DEBUG, "Muxing frame\n");
	/* mux encoded frame */
	ret = av_interleaved_write_frame(ofmt_ctx, &enc_pkt);
#endif
	return ret;
}

#if 0
static int 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;
}
#endif