#include <libavformat/avformat.h>
#ifdef FFCMP_HAVE_AVUTIL_FRAME_H
#include <libavutil/frame.h>
#else
#include <libavcodec/avcodec.h>
#endif
#include <getopt.h>
#include <stdbool.h>
#include <math.h>

#define MAX_COMPLETE_FRAMES 64

typedef struct
{
	AVFrame* frame;
	int stream_index;
} Frame;

typedef struct
{
	AVFormatContext* format_context;
	AVCodec* codec;
	AVPacket packet;
	AVFrame* current_frame;
	Frame complete_frames[MAX_COMPLETE_FRAMES];
	int n_complete_frames;
	int complete_frame_index;
} File;

static File
open_file(char* filename)
{
	File file;

	file.format_context = avformat_alloc_context();
	if (!file.format_context) {
		fprintf(stderr, "Could not create format context.\n");
		exit(EXIT_FAILURE);
	}
	int e = avformat_open_input(&file.format_context, filename, 0, 0);
	if (e < 0) {
		fprintf(stderr, "Failed to open %s\n", filename);
		exit(EXIT_FAILURE);
	}
	for (int i = 0; i < file.format_context->nb_streams; ++i) {
		file.codec = avcodec_find_decoder(file.format_context->streams[i]->codec->codec_id);
		if (!file.codec) {
			fprintf(stderr, "Could not find codec.\n");
			exit(EXIT_FAILURE);
		}
		if (avcodec_open2(file.format_context->streams[i]->codec, file.codec, 0) < 0) {
			fprintf(stderr, "Could not open codec.\n");
			exit(EXIT_FAILURE);
		}
	}

#ifdef FFCMP_HAVE_AVUTIL_FRAME_H
	file.current_frame = av_frame_alloc();
#else
	file.current_frame = avcodec_alloc_frame();
#endif
	if (!file.current_frame) {
		fprintf(stderr, "Could not allocate frame.\n");
		exit(EXIT_FAILURE);
	}

	file.n_complete_frames = 0;
	file.complete_frame_index = 0;

	return file;
}

bool
read_frame(File* file)
{
	static bool warned_about_video = false;

	int r = av_read_frame(file->format_context, &file->packet);
	if (r == AVERROR_EOF) {
		return true;
	}

	if (r < 0) {
		fprintf(stderr, "Failed to read frame.\n");
		exit(EXIT_FAILURE);
	}

	switch (file->format_context->streams[file->packet.stream_index]->codec->codec_type) {
	case AVMEDIA_TYPE_VIDEO:
		if (!warned_about_video) {
			fprintf(stderr, "Warning: ignoring video frames.\n");
			warned_about_video = true;
		}
		break;
	case AVMEDIA_TYPE_AUDIO:
	{
		AVPacket copy_packet = file->packet;
		while (copy_packet.size > 0) {
			int frame_finished;
			int decode_result = avcodec_decode_audio4(file->format_context->streams[file->packet.stream_index]->codec, file->current_frame, &frame_finished, &copy_packet);
			if (decode_result < 0) {
				fprintf(stderr, "Failed to decode audio.\n");
				exit(EXIT_FAILURE);
			}

			if (frame_finished) {
				file->complete_frames[file->n_complete_frames].frame = file->current_frame;
				file->complete_frames[file->n_complete_frames].stream_index = file->packet.stream_index;
				++file->n_complete_frames;
#ifdef FFCMP_HAVE_AVUTIL_FRAME_H
				file->current_frame = av_frame_alloc();
#else
				file->current_frame = avcodec_alloc_frame();
#endif
				if (!file->current_frame) {
					fprintf(stderr, "Could not allocate frame.\n");
					exit(EXIT_FAILURE);
				}
			}

			copy_packet.data -= decode_result;
			copy_packet.size -= decode_result;
		}
		break;
	default:
		fprintf(stderr, "Warning: ignoring other frame.\n");
		break;
	}
	}

	return false;
}

void help(char const * name)
{
	fprintf(stderr, "Syntax: %s [options] file1 file2\n", name);
	fprintf(stderr, "Options are:\n");
	fprintf(stderr, "\t--audio-sample-tolerance, -t  specify allowable difference in audio sample value, in bits (e.g. 1 means least significant bit can differ)\n");
}

int main(int argc, char** argv)
{
	int audio_sample_tolerance = 0;

	int option_index = 0;
	while (true) {
		static struct option long_options[] = {
			{ "help", no_argument, 0, 'h' },
			{ "audio-sample-tolerance", required_argument, 0, 't' },
			{ 0, 0, 0, 0 }
		};

		int c = getopt_long(argc, argv, "ht:", long_options, &option_index);

		if (c == -1) {
			break;
		}

		switch (c) {
		case 'h':
			help(argv[0]);
			break;
		case 't':
			audio_sample_tolerance = atoi(optarg);
			break;
		}
	}

	if (argc - optind < 2 || argc - optind >= 3) {
		help(argv[0]);
		exit(EXIT_FAILURE);
	}

	av_register_all();

	File file[2] = {
		open_file(argv[optind]),
		open_file(argv[optind + 1])
	};

	if (file[0].format_context->nb_streams != file[1].format_context->nb_streams) {
		fprintf(stderr, "Files have different stream counts.\n");
		exit(EXIT_FAILURE);
	}

	for (int i = 0; i < file[0].format_context->nb_streams; ++i) {
		if (file[0].format_context->streams[i]->codec->codec_type != file[1].format_context->streams[i]->codec->codec_type) {
			fprintf(stderr, "Stream %d has different code type.\n", i);
			exit(EXIT_FAILURE);
		}
	}

	while (true) {
		bool done[2] = {
			read_frame(&file[0]),
			read_frame(&file[1])
		};

		if (done[0] != done[1]) {
			fprintf(stderr, "Files are different lengths.\n");
			exit(EXIT_FAILURE);
		}

		while (file[0].n_complete_frames > 0 && file[1].n_complete_frames > 0) {
			Frame frame = file[0].complete_frames[0];
			AVStream* stream = file[0].format_context->streams[frame.stream_index];

			if (
				file[0].format_context->streams[file[0].complete_frames[0].stream_index]->codec->sample_fmt !=
				file[1].format_context->streams[file[1].complete_frames[0].stream_index]->codec->sample_fmt) {
				fprintf(stderr, "Audio sample formats differ.\n");
				exit(EXIT_FAILURE);
			}

			if (
				file[0].format_context->streams[file[0].complete_frames[0].stream_index]->codec->channels !=
				file[1].format_context->streams[file[1].complete_frames[0].stream_index]->codec->channels) {
				fprintf(stderr, "Audio channel counts differ.\n");
				exit(EXIT_FAILURE);
			}

			if (
				file[0].complete_frames[0].frame->nb_samples !=
				file[1].complete_frames[0].frame->nb_samples) {
				fprintf(stderr, "Audio frame counts differ.\n");
				exit(EXIT_FAILURE);
			}

			int const size = av_samples_get_buffer_size(0, stream->codec->channels, frame.frame->nb_samples, stream->codec->sample_fmt, 1);
			int const check = av_sample_fmt_is_planar(stream->codec->sample_fmt) ? stream->codec->channels : 1;
			for (int i = 0; i < check; ++i) {
				if (memcmp(file[0].complete_frames[0].frame->data[i], file[1].complete_frames[0].frame->data[i], size) != 0) {

					int const channels = file[0].format_context->streams[file[0].complete_frames[0].stream_index]->codec->channels;
					int const frames = frame.frame->nb_samples;

					bool different = false;
					switch (stream->codec->sample_fmt) {
					case AV_SAMPLE_FMT_S16:
					{
						int const tol = pow(2, audio_sample_tolerance - 1);
						int16_t* p = (int16_t *) (file[0].complete_frames[0].frame->data[0]);
						int16_t* q = (int16_t *) (file[1].complete_frames[0].frame->data[0]);
						for (int i = 0; i < channels; ++i) {
							for (int j = 0; j < frames; ++j) {
								if (abs(*p - *q) > tol) {
									different = true;
									fprintf(stderr, "\tsamples %d vs %d at channel %d frame %d\n", *p, *q, i, j);
								}
								++p;
								++q;
							}
						}
						break;
					}
					case AV_SAMPLE_FMT_S32P:
					{
						int const tol = pow(2, audio_sample_tolerance - 1);
						int32_t** p = (int32_t **) (file[0].complete_frames[0].frame->data);
						int32_t** q = (int32_t **) (file[1].complete_frames[0].frame->data);
						for (int i = 0; i < channels; ++i) {
							for (int j = 0; j < frames; ++j) {
								if (abs(p[i][j] - q[i][j]) > tol) {
									different = true;
									fprintf(stderr, "\tsamples %d vs %d at channel %d frame %d\n", p[i][j], q[i][j], i, j);
								}
							}
						}
						break;
					}
					case AV_SAMPLE_FMT_FLTP:
					{
						float const tol = pow(2, audio_sample_tolerance - 1) / ((float) pow(2, 32));
						float** p = (float **) (file[0].complete_frames[0].frame->data);
						float** q = (float **) (file[1].complete_frames[0].frame->data);
						for (int i = 0; i < channels; ++i) {
							for (int j = 0; j < frames; ++j) {
								if (fabs(p[i][j] - q[i][j]) > tol) {
									different = true;
									fprintf(stderr, "\tsamples %f vs %f at channel %d frame %d\n", p[i][j], q[i][j], i, j);
								}
							}
						}
						break;
					}
					default:
						fprintf(stderr, "Audio frames differ and could not be compared in detail (sample format %d unsupported).\n", stream->codec->sample_fmt);
						break;
					}

					if (different) {
						fprintf(stderr, "Audio frames %d differ.\n", file[0].complete_frame_index);
						exit(EXIT_FAILURE);
					}
				}
			}

			memmove(file[0].complete_frames, file[0].complete_frames + 1, (MAX_COMPLETE_FRAMES - 1) * sizeof(Frame));
			memmove(file[1].complete_frames, file[1].complete_frames + 1, (MAX_COMPLETE_FRAMES - 1) * sizeof(Frame));
			--file[0].n_complete_frames;
			--file[1].n_complete_frames;
			++file[0].complete_frame_index;
			++file[1].complete_frame_index;
		}

		if (done[0]) {
			break;
		}
	}

	return 0;
}