#include <stdio.h>
#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <assert.h>
#include <limits.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/opt.h>
#include <libavutil/mem.h>
#include <libavutil/avutil.h>

enum
{
	width = 352,
	height = 288,
	pts_increment = 50
};

AVCodecContext* create_encoder_context(int request_mb_info)
{
	int res = 0;
	AVCodec* codec = NULL;
	AVCodecContext* av_ctx = avcodec_alloc_context3(NULL);
	assert(av_ctx != NULL);

	av_ctx->codec_type = AVMEDIA_TYPE_VIDEO;

	av_ctx->codec_id = CODEC_ID_H263;
	av_ctx->rtp_payload_size = 1400;
	av_ctx->thread_count = 1;

	av_ctx->qcompress = 0.5;
	av_ctx->max_qdiff = 8;
	av_ctx->max_b_frames = 0;

	av_ctx->flags = 0;
	av_ctx->rc_max_rate = 0;
	av_ctx->rc_min_rate = 0;
	av_ctx->rc_buffer_size = 0;
	av_ctx->slice_count = 0;
	av_ctx->profile = FF_PROFILE_UNKNOWN;

	av_ctx->level = FF_LEVEL_UNKNOWN;

	av_ctx->me_method = ME_EPZS;
	av_ctx->me_cmp |= FF_CMP_CHROMA;
	av_ctx->me_subpel_quality = 0;
	av_ctx->me_range = 16;
	av_ctx->scenechange_threshold = 0;
	av_ctx->i_quant_factor = 0.9;
	av_ctx->b_frame_strategy = 1;
	av_ctx->refs = 1;
	av_ctx->trellis = 0;
	av_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;

	av_ctx->width = width;
	av_ctx->height = height;
	av_ctx->pix_fmt = PIX_FMT_YUV420P;
	av_ctx->color_range = AVCOL_RANGE_JPEG;

	av_ctx->time_base.num = 1;
	av_ctx->time_base.den = 1000;

	av_ctx->gop_size = 50;
	av_ctx->keyint_min = 3;
	av_ctx->qmin = 2;
	av_ctx->qmax = 20;
	av_ctx->bit_rate = 800 * 1000;
	av_ctx->bit_rate_tolerance = av_ctx->bit_rate * 0.25f;

	codec = avcodec_find_encoder(av_ctx->codec_id);
	assert(codec != NULL);
	res = avcodec_open2(av_ctx, codec, NULL);
	assert(res == 0);

	if (request_mb_info)
	{
		res = av_opt_set_int(av_ctx, "mb_info", 1400, AV_OPT_SEARCH_CHILDREN);
		assert(res == 0);
	}

	return av_ctx;
}

void destroy_encoder_context(AVCodecContext* p)
{
	if (p)
	{
		if (p->internal)
			avcodec_close(p);
		if (p->extradata)
		{
			av_freep(&p->extradata);
			p->extradata_size = 0;
		}
		av_free(p);
	}
}

void init_frame(AVFrame* in_frame, size_t* raw_size)
{
	uint8_t* data = NULL;
	size_t size = 0;

	memset(in_frame, 0, sizeof(*in_frame));
	avcodec_get_frame_defaults(in_frame);

	*raw_size = size = avpicture_get_size(PIX_FMT_YUV420P, width, height);
	data = av_malloc(size);
	assert(data != NULL);
	memset(data, 0, size);

	avpicture_fill((AVPicture*)in_frame, data, PIX_FMT_YUV420P, width, height);
}

void generate_frame(AVFrame* in_frame, size_t raw_size, unsigned int pts)
{
	unsigned int i = 0, j, n = (width - 20) / 2, threshold = (pts / pts_increment) % n;
	unsigned char y, cb, cr;

	for (i = 0; i < width / 2; ++i)
	{
		y = i < threshold ? 0 : 255;
		if (threshold & 8u)
			y = ~y;
		in_frame->data[0][i] = in_frame->data[0][width - i - 1] = y;
		j = i / 2;
		if (i < threshold)
		{
			cb = cr = 127;
		}
		else
		{
			cb = 127.0f * (float)(i - threshold) / (float)(width / 2 - threshold);
			cr = -cb;
		}
		in_frame->data[1][j] = in_frame->data[1][(width / 2) - j - 1] = cb;
		in_frame->data[2][j] = in_frame->data[2][(width / 2) - j - 1] = cr;
	}

	for (i = 1; i < height; ++i)
	{
		memcpy(in_frame->data[0] + i * in_frame->linesize[0], in_frame->data[0], width);
	}

	for (i = 1; i < height / 2; ++i)
	{
		memcpy(in_frame->data[1] + i * in_frame->linesize[1], in_frame->data[1], width / 2);
		memcpy(in_frame->data[2] + i * in_frame->linesize[2], in_frame->data[2], width / 2);
	}
}

void init_packet(AVPacket* packet)
{
	memset(packet, 0, sizeof(*packet));
	av_init_packet(packet);
}

AVFormatContext* create_format_context()
{
	AVCodec* video_codec = NULL;
	AVCodecContext* av_video_codec_ctx = NULL;
	AVStream* av_stream = NULL;
	int res = 0;
	AVFormatContext* av_format_ctx = avformat_alloc_context();
	assert(av_format_ctx != NULL);

	av_format_ctx->oformat = av_guess_format("avi", "output.avi", NULL);
	assert(av_format_ctx->oformat != NULL);

	av_format_ctx->oformat->video_codec = CODEC_ID_NONE;
	av_format_ctx->video_codec_id = CODEC_ID_NONE;
	av_format_ctx->oformat->audio_codec = CODEC_ID_NONE;
	av_format_ctx->audio_codec_id = CODEC_ID_NONE;

	res = avio_open(&av_format_ctx->pb, "output.avi", AVIO_FLAG_WRITE);
	assert(res >= 0);

	video_codec = avcodec_find_encoder(CODEC_ID_H263);
	assert(video_codec != NULL);

	av_format_ctx->oformat->video_codec = CODEC_ID_H263;
	av_format_ctx->video_codec_id = CODEC_ID_H263;

	av_stream = avformat_new_stream(av_format_ctx, video_codec);
	assert(av_stream != NULL);
	av_stream->id = 0;

	av_stream->r_frame_rate = av_d2q(1000.0 / pts_increment, USHRT_MAX);
	av_stream->time_base.num = av_stream->r_frame_rate.den;
	av_stream->time_base.den = av_stream->r_frame_rate.num;

	av_video_codec_ctx = av_stream->codec;

	av_video_codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
	av_video_codec_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;
	av_video_codec_ctx->time_base.num = 1;
	av_video_codec_ctx->time_base.den = 1000;
	av_video_codec_ctx->width = width;
	av_video_codec_ctx->height = height;
	av_video_codec_ctx->pix_fmt = PIX_FMT_YUV420P;

	res = avformat_write_header(av_format_ctx, NULL);
	assert(res >= 0);

	return av_format_ctx;
}

void destroy_format_context(AVFormatContext* p)
{
	unsigned int i = 0;
	if (p)
	{
		av_write_trailer(p);
		if (p->pb)
		{
			avio_sync(p->pb);
			avio_close(p->pb);
			p->pb = NULL;
		}

		for (i = 0; i < p->nb_streams; ++i)
		{
			if (p->streams[i]->codec->internal) /* this field is only set when the codec is opened */
				avcodec_close(p->streams[i]->codec);
		}

		if (p->iformat)
		{
			avformat_close_input(&p);
		}
		else
		{
			if (p->oformat && (p->oformat->flags & AVFMT_NOFILE) == 0 && (p->oformat->flags & AVFMT_FLAG_CUSTOM_IO) == 0 && p->pb)
			{
				avio_close(p->pb);
				p->pb = NULL;
			}
			avformat_free_context(p);
		}
	}
}

void write_video(AVFormatContext* av_format_ctx, uint8_t* data, unsigned int size, uint64_t pts, uint64_t dts, unsigned int keyframe)
{
	int res = 0;

	/*
	 * Note: Reconstructing AVPacket is essential because in real life the encoded data is not passed directly to libavformat/libavcodec.
	 * It may be sent over network, for example.
	 */
	AVPacket packet;

	printf("Encoded video, pts: %u, dts: %u, size: %u\n", (unsigned int)pts, (unsigned int)dts, size);

	init_packet(&packet);
	packet.stream_index = 0;
	packet.data = data;
	packet.size = size;
	packet.pts = pts;
	packet.dts = dts;
	packet.flags = keyframe ? AV_PKT_FLAG_KEY : 0;

	res = av_interleaved_write_frame(av_format_ctx, &packet);
	assert(res == 0);
}

int main(int argc, char* argv[])
{
	AVCodecContext* av_ctx = NULL;
	AVFrame in_frame;
	size_t raw_size = 0;
	AVPacket out_packet;
	int got_packet = 0, res = 0;
	AVFormatContext* av_format_ctx = NULL;
	unsigned int pts = 0;
	int request_mb_info = 1;
	uint8_t* out_buffer = NULL;

	if (argc >= 2 && strcmp(argv[1], "--without_side_data") == 0)
		request_mb_info = 0;

	av_register_all();

	av_ctx = create_encoder_context(request_mb_info);

	init_frame(&in_frame, &raw_size);
	init_packet(&out_packet);

	out_buffer = av_malloc(raw_size * 4u);
	assert(out_buffer != NULL);

	av_format_ctx = create_format_context();

	for (pts = 0; pts < 5000; pts += pts_increment)
	{
		in_frame.pts = pts;
		in_frame.pict_type = (enum AVPictureType)0; /* AV_PICTURE_TYPE_NONE */
		generate_frame(&in_frame, raw_size, pts);

		out_packet.data = out_buffer;
		out_packet.size = raw_size * 4u;

		got_packet = 0;
		res = avcodec_encode_video2(av_ctx, &out_packet, &in_frame, &got_packet);
		assert(res == 0);

		if (!got_packet)
			continue;

/*		This assert fails with ffmpeg and not with libav. av_packet_get_side_data doesn't find the merged side data.
		if (request_mb_info)
			assert(av_packet_get_side_data(&out_packet, AV_PKT_DATA_H263_MB_INFO, NULL));
*/
		/*
		 * Note: since we provided our own buffer in AVPacket to avcodec_encode_video2, out_packet.data must be equal to out_buffer here.
		 * This is according to the avcodec_encode_video2 documentation. In reality it is not, av_packet_merge_side_data() allocates a new buffer and
		 * writes the whole frame and side data to the new buffer. The old buffer gets released and its memory only contain the encoded frame + junk at the end.
		 *
		 * This works as intended with libav, which uses the buffer supplied in AVPacket.
		 */
		write_video(av_format_ctx, out_buffer, out_packet.size, out_packet.pts, out_packet.dts, (out_packet.flags & AV_PKT_FLAG_KEY) != 0);

		out_packet.data = NULL;
		out_packet.size = 0;
		av_free_packet(&out_packet);
	}

	destroy_format_context(av_format_ctx);
	destroy_encoder_context(av_ctx);

	return 0;
}