#include <stdio.h>
#include <avcodec.h>
#include <avformat.h>

#include "writewebm.h"

#define STREAM_FRAME_RATE 25
#define STREAM_PIX_FMT AV_PIX_FMT_YUV420P 

OutputStream audioStream = { 0 }, videoStream = { 0 };

WebMWriter *init(char *filename)
{
	printf("@@@@@@@@@@@ INITING libavformat ....\n");
	av_register_all();
	
	AVFormatContext *ctx ;
	AVCodec *audioCodec, *videoCodec;
	const char *fmt_name = NULL;
	const char *file_name = filename;

	int alloc_status = avformat_alloc_output_context2(&ctx, NULL, fmt_name, file_name);
	printf("@@@@ result : %d\n",alloc_status);
	if(!ctx)
		return NULL;

	AVOutputFormat *fmt = (*ctx).oformat;
	printf("@@@ format : %s\t%d\t%d\n",(*fmt).long_name,(*fmt).video_codec,(*fmt).audio_codec);

	if(fmt->video_codec != AV_CODEC_ID_NONE)
	{
		addStream(&videoStream, ctx, &videoCodec, AV_CODEC_ID_VP8);
	}

	if(fmt->audio_codec != AV_CODEC_ID_NONE)
	{
		addStream(&audioStream, ctx, &audioCodec, fmt->audio_codec);
	}
	
	//printf("@@@@@@@@@@@@ JUST NOW ALLOCATED STREAMS %s\t%s\n",videoStream.enc->codec->long_name,audioStream.enc->codec->long_name);
	
	printf("@@@@@@ BEFORE openAudio, strict_std_compliance : %d\t%d\n",audioStream.enc->strict_std_compliance, videoStream.enc->strict_std_compliance); 
	if(videoStream.st)
		openVideo1(&videoStream, videoCodec);
	if(audioStream.st)
		openAudio(&audioStream, audioCodec);

	av_dump_format(ctx, 0, file_name, 1);

	int ret;
	/* open the output file, if needed */
	if (!(fmt->flags & AVFMT_NOFILE)) {
		ret = avio_open(&ctx->pb, file_name, AVIO_FLAG_WRITE);
		if (ret < 0) {
			printf("Could not open '%s': %s\n", file_name, av_err2str(ret));
			return NULL;
		}
	}

	if(fmt->flags & AVFMT_NOFILE)
		printf("@@@@@ AVFMT_NOFILE is set \n");

	/* Write the stream header, if any. */
	ret = avformat_write_header(ctx, NULL);
	if (ret < 0) {
		fprintf(stderr, "Error occurred when opening output file: %s\n",
				av_err2str(ret));
		return NULL;
	}
	printf("@@@@@ frame_size : %d\t%d\n",audioStream.enc->frame_size, videoStream.enc->frame_size);
	printf("@@@@@@ HEADERS Written %zd\n",sizeof(struct WebMWriter));

	WebMWriter *webmWriter = malloc(sizeof(struct WebMWriter));
	webmWriter->ctx = ctx;
	webmWriter->outfmt = fmt;
	webmWriter->audioStream = &audioStream;
	webmWriter->videoStream = &videoStream;
	webmWriter->audioCodec = audioCodec;
	webmWriter->videoCodec = videoCodec;

	/*
	uint8_t *data;
	int size;
	int frameTimeStamp = 1;
	writeVideoStream(ctx, videoStream.st, data, size, frameTimeStamp);	

	int sampleRate = audioStream.enc->sample_rate;
	int samplesPerFrame = audioStream.enc->frame_size;
	writeAudioStream(ctx, audioStream.st, data, size, frameTimeStamp, samplesPerFrame, sampleRate, audioStream.enc); 
	*/

	/* CLOSE */

	if(webmWriter->videoStream->st == NULL)
		printf("@@@@@@ writewebM videoStream is NULL \n");
	else
		printf("@@@@@@@ writewebM videoStream is NOT NULL \n");

	return webmWriter;
}

void close(AVFormatContext *ctx, OutputStream videoStream, OutputStream audioStream, AVOutputFormat *outfmt)
{
	int ret = av_write_trailer(ctx);
	printf("@@@@@@ ret : %d\n",ret);
 
 	/*
 	if (videoStream.st)
                //avcodec_close(videoStream.enc);
		avcodec_free_context(&videoStream.enc); //this calls avcodec_close
        if (audioStream.st)
                //avcodec_close(audioStream.enc);
		avcodec_free_context(&audioStream.enc); //this calls avcodec_close

        if (!(outfmt->flags & AVFMT_NOFILE))
                //avio_close(ctx->pb);
                avio_closep(&ctx->pb);  //this calls avio_close
	*/

        /* free the stream */
        //avformat_free_context(ctx);

        /* free the streams */
	/*
        for(int i = 0; i < ctx->nb_streams; i++) {
                av_freep(&ctx->streams[i]->codec);
                av_freep(&ctx->streams[i]);
        }
        av_free(ctx);
	*/
	//return;
}

int writeVideoStream(AVFormatContext *ctx, AVStream *st, uint8_t *data, int size, long frameTimeStamp, int isKeyFrame, AVCodecContext *codec_ctx)
{
	AVRational rat = st->time_base;
	printf("@@@@@ writeVideoStream TIMEBASE : %d\t%d\t%ld\n",rat.num, rat.den, frameTimeStamp);
	AVPacket pkt;
	av_init_packet(&pkt);
	
	void *opaque;
	int flags = AV_BUFFER_FLAG_READONLY;
	AVBufferRef *bufferRef = av_buffer_create(data, size, NULL, opaque, flags);
	
	pkt.buf = bufferRef;
	pkt.data = data;
	pkt.size = size;
	pkt.stream_index  = st->index;
	pkt.pts = pkt.dts = frameTimeStamp;

	/* rescale output packet timestamp values from codec to stream timebase */
    	//av_packet_rescale_ts(&pkt, codec_ctx->time_base, st->time_base);
	
	if(isKeyFrame == 1)
		pkt.flags |= AV_PKT_FLAG_KEY;

	//int ret = av_write_frame(ctx, &pkt);
	int ret = av_interleaved_write_frame(ctx, &pkt);
	return ret;
}

int writeAudioStream(AVFormatContext *ctx, AVStream *st, uint8_t *data, int size, long frameTimeStamp, AVCodecContext *codec_ctx)
{
	int sampleRate = codec_ctx->sample_rate;
	int samplesPerFrame = codec_ctx->frame_size;
        AVRational rat = st->time_base;
        printf("@@@@@ writeAudioStream TIMEBASE : %d\t%d\t%ld\n",rat.num, rat.den, frameTimeStamp);
        AVPacket pkt;
        av_init_packet(&pkt);

        void *opaque;
        int flags = AV_BUFFER_FLAG_READONLY;
        AVBufferRef *bufferRef = av_buffer_create(data, size, NULL, opaque, flags);

        pkt.buf = bufferRef;
        pkt.data = data;
        pkt.size = size;
        pkt.stream_index  = st->index;
	pkt.pts = pkt.dts = frameTimeStamp;
	AVRational sampleRat = (AVRational){1, sampleRate};
	pkt.duration = av_rescale_q(samplesPerFrame, sampleRat, codec_ctx->time_base);

	/* rescale output packet timestamp values from codec to stream timebase */
    	//av_packet_rescale_ts(&pkt, codec_ctx->time_base, st->time_base);
	
	printf("@@@@@@@ DURATION : %lld\n",pkt.duration);

        //int ret = av_write_frame(ctx, &pkt);
        int ret = av_interleaved_write_frame(ctx, &pkt);
        return ret;
}

void addStream(OutputStream *out_st, AVFormatContext *ctx, AVCodec **cdc, enum AVCodecID codecId)
{
	printf("@@@@@@@@@@@@@@ ADD STREAM CALLED for codecID : %d\n",codecId);

	(*cdc) = avcodec_find_encoder(codecId);
	if(!(*cdc)) {
		printf("@@@@@ couldnt find codec \n");
		exit(1);
	}

	/*as we are passing a NULL AVCodec cdc, So AVCodecContext codec_ctx will not be allocated, we have to do it explicitly */ 
	AVStream *st = avformat_new_stream(ctx, *cdc);
	if(!st) {
		printf("@@@@@ couldnt init stream\n");
		exit(1);
	}
	printf("@@@@@@@@@@@@@@ STREAM INITED \n");

	out_st->st = st;
	st->id = ctx->nb_streams-1;
	/*AVCodecParameters *codecpars = st->codecpar;
	codecpars->codec_id = codecId;
	codecpars->codec_type = (*cdc)->type;
	*/

	AVCodecContext *codec_ctx = st->codec;
	if (!codec_ctx) {
		fprintf(stderr, "Could not alloc an encoding context\n");
		exit(1);
	}
	out_st->enc = codec_ctx;
  
  	//since opus is experimental codec
   	codec_ctx->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
	switch ((*cdc)->type) {
		case AVMEDIA_TYPE_AUDIO:
			codec_ctx->codec_id = codecId;
			codec_ctx->sample_fmt  = AV_SAMPLE_FMT_FLTP;//AV_SAMPLE_FMT_U8 or AV_SAMPLE_FMT_S16;
			codec_ctx->bit_rate    = 64000;
			codec_ctx->sample_rate = 48000;
			codec_ctx->channels    = 2;//1;
			codec_ctx->channel_layout = AV_CH_LAYOUT_STEREO; //AV_CH_LAYOUT_MONO;
			codec_ctx->codec_type = AVMEDIA_TYPE_AUDIO;

			/*codecpars->format = codec_ctx->sample_fmt;
			codecpars->channels = codec_ctx->channels;
			codecpars->sample_rate = codec_ctx->sample_rate;
			*/
			break;

		case AVMEDIA_TYPE_VIDEO:
			codec_ctx->codec_id = codecId;
			codec_ctx->bit_rate = 90000;
			codec_ctx->width    = 640;
			codec_ctx->height   = 480;
			
			//codec_ctx->time_base.den = STREAM_FRAME_RATE;
			//codec_ctx->time_base.num = 1;
			codec_ctx->time_base = (AVRational){1,STREAM_FRAME_RATE};
			codec_ctx->gop_size = 12;
			codec_ctx->pix_fmt = STREAM_PIX_FMT;
			codec_ctx->max_b_frames = 1;
			codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;

			/*codecpars->format = codec_ctx->pix_fmt;
			codecpars->width = codec_ctx->width;
			codecpars->height = codec_ctx->height;
			codecpars->sample_aspect_ratio = (AVRational){codec_ctx->width, codec_ctx->height};
			*/

			break;

		default:
			break;
	}	
	//codecpars->bit_rate = codec_ctx->bit_rate;

   	printf("@@@@@ CODECPARS %d\t%d\t%d\t%d\n",codecId,st->codecpar->codec_id,AVMEDIA_TYPE_VIDEO, (*cdc)->type);
	/* Some formats want stream headers to be separate. */
	if (ctx->oformat->flags & AVFMT_GLOBALHEADER)
		codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
}

void openVideo1(OutputStream *out_st, AVCodec *codec)
{
	AVFrame *frame;
	AVCodecContext *codec_ctx = out_st->enc;
	int ret = avcodec_open2(codec_ctx, codec, NULL);

	//frame = avcodec_alloc_frame();
	//if(!frame)
		//printf("@@@@@ Frame can't be allocated \n");

	/* copy the stream parameters to the muxer */
	ret = avcodec_parameters_from_context(out_st->st->codecpar, codec_ctx);
	if (ret < 0) {
		printf("Could not copy the stream parameters\n");
		exit(1);
	}

	printf("@@@@@ Video Opened : %d\n",ret);
}

void openAudio(OutputStream *out_st, AVCodec *codec)
{
	AVCodecContext *codec_ctx = out_st->enc;
	int ret = avcodec_open2(codec_ctx, codec, NULL);
	//int ret = ff_codec_open2_recursive(codec_ctx, codec, NULL);

	/* copy the stream parameters to the muxer */
	ret = avcodec_parameters_from_context(out_st->st->codecpar, codec_ctx);
	if (ret < 0) {
		printf("Could not copy the stream parameters\n");
		exit(1);
	}
	printf("@@@@@@ Audio Opened : %d\t%s\t%d\n",ret,codec->long_name,codec_ctx->codec_id);
}