/*
 * Cantata
 *
 * Copyright (c) 2011-2022 Craig Drummond <craig.p.drummond@gmail.com>
 *
 */
/* This class is a C++/Qt version of input_ffmpeg.c from libebur128
 */

/* See LICENSE file for copyright and license details. */
#ifdef __cplusplus
#define __STDC_CONSTANT_MACROS
#ifdef _STDINT_H
#undef _STDINT_H
#endif
#include <stdint.h>
#endif
#ifdef __cplusplus
extern "C" {
#endif
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#if LIBAVFORMAT_VERSION_MAJOR >= 54
#include <libavutil/channel_layout.h>
#endif
#ifdef __cplusplus
}
#endif
#include "ebur128.h"
#include "ffmpeginput.h"
#include <QByteArray>
#include <QFile>
#include <QList>
#include <QMutex>
#include <QString>

static QMutex mutex;

#ifndef AVCODEC_MAX_AUDIO_FRAME_SIZE
#define AVCODEC_MAX_AUDIO_FRAME_SIZE 192000// 1 second of 48khz 32bit audio
#endif

#if defined FF_INPUT_BUFFER_PADDING_SIZE
#define FFMPEG_INPUT_BUFFER_PADDING_SIZE FF_INPUT_BUFFER_PADDING_SIZE
#elif defined AV_INPUT_BUFFER_PADDING_SIZE
#define FFMPEG_INPUT_BUFFER_PADDING_SIZE AV_INPUT_BUFFER_PADDING_SIZE
#else
#define FFMPEG_INPUT_BUFFER_PADDING_SIZE 32
#endif

#define BUFFER_SIZE ((((AVCODEC_MAX_AUDIO_FRAME_SIZE * 3) / 2) * sizeof(int16_t)) + FFMPEG_INPUT_BUFFER_PADDING_SIZE)

#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 89, 100)// Not 100% of version here!
#define GET_CODEC_TYPE(A) A->codecpar->codec_type
#else
#define GET_CODEC_TYPE(A) A->codec->codec_type
#endif

#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 24, 0)
#define AV_FREE(PKT) av_packet_unref(PKT)
#else
#define AV_FREE(PKT) av_free_packet(PKT)
#endif

struct FfmpegInput::Handle {
	Handle()
		: formatContext(0), codecContext(0), codec(0)
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(53, 35, 0)
		  ,
		  frame(0), gotFrame(0), packetLeft(false), flushing(false)
#endif
		  ,
		  audioStream(0), currentBytes(0)
	{
		packet = av_packet_alloc();
		origPacket = NULL;
		audioBuffer = (int16_t*)av_malloc(BUFFER_SIZE);
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(53, 35, 0)
#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(55, 39, 101)
		frame = avcodec_alloc_frame();
#else
		frame = av_frame_alloc();
#endif
#endif
	}
	~Handle()
	{
		if (origPacket != NULL) {
			av_packet_free(&origPacket);
		}
		if (packet != NULL) {
			av_packet_free(&packet);
		}
		if (audioBuffer) {
			av_free(audioBuffer);
		}
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(53, 35, 0)
		if (frame) {
#if LIBAVCODEC_VERSION_INT <= AV_VERSION_INT(54, 23, 100)
			av_free(&frame);
#elif LIBAVCODEC_VERSION_INT < AV_VERSION_INT(55, 39, 101)
			avcodec_free_frame(&frame);
#else
			av_frame_free(&frame);
#endif
		}
#endif
	}
	AVFormatContext* formatContext;
	AVCodecContext* codecContext;
	const AVCodec* codec;
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(53, 35, 0)
	AVFrame* frame;
	int gotFrame;
	bool packetLeft;
	bool flushing;
	AVPacket* origPacket;
#endif
	AVPacket* packet;
	int audioStream;
	int16_t* audioBuffer;
	float buffer[BUFFER_SIZE / 2 + 1];
	QList<QByteArray> bufferList;
	size_t currentBytes;
};

void FfmpegInput::init()
{
	static int i = false;
	if (!i) {
#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(58, 18, 100)
		av_register_all();
#endif
		av_log_set_level(AV_LOG_ERROR);
		i = true;
	}
}

FfmpegInput::FfmpegInput(const QString& fileName)
{
	mutex.lock();
	handle = new Handle;

	bool ok = true;
#if LIBAVFORMAT_VERSION_MAJOR >= 54 || (LIBAVFORMAT_VERSION_MAJOR == 53 && LIBAVFORMAT_VERSION_INT >= AV_VERSION_INT(53, 2, 0)) || (LIBAVFORMAT_VERSION_MAJOR == 52 && LIBAVFORMAT_VERSION_INT >= AV_VERSION_INT(52, 110, 0))
	if (avformat_open_input(&handle->formatContext, QFile::encodeName(fileName).constData(), NULL, NULL) != 0)
#else
	if (av_open_input_file(&handle->formatContext, QFile::encodeName(fileName).constData(), NULL, 0, NULL) != 0)
#endif
	{
		mutex.unlock();
		delete handle;
		handle = 0;
		return;
	}
#if LIBAVFORMAT_VERSION_INT >= AV_VERSION_INT(53, 21, 0)
	if (ok && avformat_find_stream_info(handle->formatContext, 0) < 0) {
		ok = false;
	}
#else
	if (ok && av_find_stream_info(handle->formatContext) < 0) {
		ok = false;
	}
#endif

	// Find the first audio stream
	if (ok) {
		handle->audioStream = -1;
		for (size_t j = 0; j < handle->formatContext->nb_streams; ++j) {
			if (GET_CODEC_TYPE(handle->formatContext->streams[j]) ==
#if LIBAVCODEC_VERSION_MAJOR >= 53 || (LIBAVCODEC_VERSION_MAJOR == 52 && LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(52, 64, 0))
			    AVMEDIA_TYPE_AUDIO
#else
			    CODEC_TYPE_AUDIO
#endif
			) {
				handle->audioStream = (int)j;
				break;
			}
		}
		if (-1 == handle->audioStream) {
			ok = false;
		}
	}

	if (ok) {
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 89, 100)// Not 100% of version here!
		handle->codec = avcodec_find_decoder(handle->formatContext->streams[handle->audioStream]->codecpar->codec_id);
		handle->codecContext = avcodec_alloc_context3(handle->codec);
		if (avcodec_parameters_to_context(handle->codecContext, handle->formatContext->streams[handle->audioStream]->codecpar) < 0) {
			ok = false;
		}
#else
		// Get a pointer to the codec context for the audio stream
		handle->codecContext = handle->formatContext->streams[handle->audioStream]->codec;

// request float output if supported
#if LIBAVCODEC_VERSION_MAJOR >= 54
		handle->codecContext->request_sample_fmt = AV_SAMPLE_FMT_FLT;
#elif (LIBAVCODEC_VERSION_MAJOR == 53 && LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(53, 4, 0))
		handle->codecContext->request_sample_fmt = SAMPLE_FMT_FLT;
#endif
		// Find the decoder for the video stream, and open codec...
		handle->codec = avcodec_find_decoder(handle->codecContext->codec_id);
#endif

		if (ok) {
			QString floatCodec = QLatin1String(handle->codec->name) + QLatin1String("float");
			const AVCodec* possibleFloatCodec = avcodec_find_decoder_by_name(floatCodec.toLatin1().constData());
			if (possibleFloatCodec) {
				handle->codec = possibleFloatCodec;
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 89, 100)// Not 100% of version here!
				avcodec_free_context(&handle->codecContext);
				handle->codecContext = avcodec_alloc_context3(handle->codec);
				if (avcodec_parameters_to_context(handle->codecContext, handle->formatContext->streams[handle->audioStream]->codecpar) < 0) {
					ok = false;
				}
#endif
			}

			if (!handle->codec ||
#if LIBAVCODEC_VERSION_MAJOR >= 53
			    avcodec_open2(handle->codecContext, handle->codec, NULL) < 0)
#else
			    avcodec_open(handle->codecContext, handle->codec) < 0)
#endif
			{
				ok = false;
			}
		}
	}

	if (ok) {
		mutex.unlock();
	}
	else {
#if LIBAVFORMAT_VERSION_INT >= AV_VERSION_INT(53, 21, 0)
		avformat_close_input(&handle->formatContext);
#else
		av_close_input_file(handle->formatContext);
#endif
		mutex.unlock();
		delete handle;
		handle = 0;
	}
}

FfmpegInput::~FfmpegInput()
{
	if (handle) {
		mutex.lock();
		avcodec_free_context(&handle->codecContext);
#if LIBAVFORMAT_VERSION_INT >= AV_VERSION_INT(53, 21, 0)
		avformat_close_input(&handle->formatContext);
#else
		av_close_input_file(handle->formatContext);
#endif
		mutex.unlock();
		delete handle;
		handle = 0;
	}
}

size_t FfmpegInput::totalFrames() const
{
	if (!handle) {
		return 0;
	}

	double tmp = (double)handle->formatContext->streams[handle->audioStream]->duration
			* (double)handle->formatContext->streams[handle->audioStream]->time_base.num
			/ (double)handle->formatContext->streams[handle->audioStream]->time_base.den
			* (double)handle->codecContext->sample_rate;

	return tmp <= 0.0 ? 0 : (size_t)(tmp + 0.5);
}

unsigned int FfmpegInput::channels() const
{
	return handle ? handle->codecContext->ch_layout.nb_channels : 0;
}

unsigned long FfmpegInput::sampleRate() const
{
	return handle ? handle->codecContext->sample_rate : 0;
}

float* FfmpegInput::buffer() const
{
	return handle ? handle->buffer : 0;
}

bool FfmpegInput::setChannelMap(int* st) const
{
	if (handle && handle->codecContext->ch_layout.u.mask) {
		unsigned int mapIndex = 0;
		int bitCounter = 0;
		while (mapIndex < (unsigned)handle->codecContext->ch_layout.nb_channels) {
			if (handle->codecContext->ch_layout.u.mask & (1 << bitCounter)) {
				switch (1 << bitCounter) {
#if LIBAVFORMAT_VERSION_MAJOR >= 54
				case AV_CH_FRONT_LEFT:
#else
				case CH_FRONT_LEFT:
#endif
					st[mapIndex] = EBUR128_LEFT;
					break;
#if LIBAVFORMAT_VERSION_MAJOR >= 54
				case AV_CH_FRONT_RIGHT:
#else
				case CH_FRONT_RIGHT:
#endif
					st[mapIndex] = EBUR128_RIGHT;
					break;
#if LIBAVFORMAT_VERSION_MAJOR >= 54
				case AV_CH_FRONT_CENTER:
#else
				case CH_FRONT_CENTER:
#endif
					st[mapIndex] = EBUR128_CENTER;
					break;
#if LIBAVFORMAT_VERSION_MAJOR >= 54
				case AV_CH_BACK_LEFT:
#else
				case CH_BACK_LEFT:
#endif
					st[mapIndex] = EBUR128_LEFT_SURROUND;
					break;
#if LIBAVFORMAT_VERSION_MAJOR >= 54
				case AV_CH_BACK_RIGHT:
#else
				case CH_BACK_RIGHT:
#endif
					st[mapIndex] = EBUR128_RIGHT_SURROUND;
					break;
				default:
					st[mapIndex] = EBUR128_UNUSED;
					break;
				}
				++mapIndex;
			}
			++bitCounter;
		}
		return true;
	}

	return false;
}

size_t FfmpegInput::readFrames()
{
	if (!handle || !channels()) {
		return 0;
	}

	size_t bufferPosition = 0;

	while (handle->currentBytes < BUFFER_SIZE) {
		size_t numberRead = readOnePacket();
		if (!numberRead) {
			break;
		}
		size_t bufferSize = numberRead * channels() * sizeof(float);
		handle->bufferList.append(QByteArray(reinterpret_cast<const char*>(handle->buffer), bufferSize));
		handle->currentBytes += bufferSize;
	}

	while (handle->bufferList.count() && handle->bufferList.first().size() + bufferPosition <= BUFFER_SIZE) {
		QByteArray b = handle->bufferList.takeAt(0);
		memcpy((char*)handle->buffer + bufferPosition, b.constData(), b.size());
		bufferPosition += b.size();
		handle->currentBytes -= b.size();
	}

	return bufferPosition / sizeof(float) / channels();
}

#if LIBAVCODEC_VERSION_MAJOR >= 54
static int decodePacket(FfmpegInput::Handle* handle)
{
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 89, 100)// Not 100% of version here!
	int ret = avcodec_receive_frame(handle->codecContext, handle->frame);
	if (0 == ret) {
		handle->gotFrame = true;
	}
	if (AVERROR(EAGAIN) == ret) {
		ret = 0;
	}
	if (0 == ret) {
		ret = avcodec_send_packet(handle->codecContext, handle->packet);
	}
	if (AVERROR(EAGAIN) == ret) {
		ret = 0;
	}
	return ret < 0 ? ret : handle->packet->size;
#else
	int ret = avcodec_decode_audio4(handle->codecContext, handle->frame, &handle->gotFrame, handle->packet);
	if (ret < 0) {
		return ret;
	}
	return FFMIN(ret, handle->packet->size);
#endif
}

size_t FfmpegInput::readOnePacket()
{
	if (!handle) {
		return 0;
	}

start:
	if (handle->flushing) {
		handle->packet->data = NULL;
		handle->packet->size = 0;
		decodePacket(handle);
		if (!handle->gotFrame) {
			return 0;
		}
		else {
			goto write_to_buffer;
		}
	}

	if (handle->packetLeft) {
		goto packetLeft;
	}

	//next_frame:
	for (;;) {
		if (av_read_frame(handle->formatContext, handle->packet) < 0) {
			handle->flushing = true;
			goto start;
		}
		if (handle->packet->stream_index != handle->audioStream) {
			AV_FREE(handle->packet);
			continue;
		}
		else {
			break;
		}
	}

	int ret;
	if (handle->origPacket != NULL) {
		AV_FREE(handle->origPacket);
	}
	handle->origPacket = av_packet_clone(handle->packet);
packetLeft:
	ret = decodePacket(handle);
	if (ret < 0) {
		goto free_packet;
	}
	handle->packet->data += ret;
	handle->packet->size -= ret;
	if (handle->packet->size > 0) {
		handle->packetLeft = true;
	}
	else {
	free_packet:
		AV_FREE(handle->origPacket);
		handle->origPacket = NULL;
		handle->packetLeft = false;
	}

	if (!handle->gotFrame) {
		goto start;
	}

write_to_buffer:;
	size_t numberRead = handle->frame->nb_samples;
	/* TODO: fix this */
	int numChannels = handle->codecContext->ch_layout.nb_channels;
	// channels = handle->frame->channels;

	if (handle->frame->nb_samples * numChannels > (int)sizeof handle->buffer) {
		return 0;
	}

	switch (handle->codecContext->sample_fmt) {
	case AV_SAMPLE_FMT_S16: {
		int16_t* dataShort = (int16_t*)handle->frame->extended_data[0];
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			handle->buffer[i] = ((float)dataShort[i]) / qMax(-(float)SHRT_MIN, (float)SHRT_MAX);
		}
		break;
	}
	case AV_SAMPLE_FMT_S32: {
		int32_t* dataInt = (int32_t*)handle->frame->extended_data[0];
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			handle->buffer[i] = ((float)dataInt[i]) / qMax(-(float)INT_MIN, (float)INT_MAX);
		}
		break;
	}
	case AV_SAMPLE_FMT_FLT: {
		float* dataFloat = (float*)handle->frame->extended_data[0];
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			handle->buffer[i] = dataFloat[i];
		}
		break;
	}
	case AV_SAMPLE_FMT_DBL: {
		double* dataDouble = (double*)handle->frame->extended_data[0];
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			handle->buffer[i] = (float)dataDouble[i];
		}
		break;
	}
	case AV_SAMPLE_FMT_S16P: {
		uint8_t** ed = handle->frame->extended_data;
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			int currentChannel = i / handle->frame->nb_samples;
			int currentSample = i % handle->frame->nb_samples;
			handle->buffer[currentSample * numChannels + currentChannel] = ((float)((int16_t*)ed[currentChannel])[currentSample]) / qMax(-(float)SHRT_MIN, (float)SHRT_MAX);
		}
		break;
	}
	case AV_SAMPLE_FMT_S32P: {
		uint8_t** ed = handle->frame->extended_data;
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			int currentChannel = i / handle->frame->nb_samples;
			int currentSample = i % handle->frame->nb_samples;
			handle->buffer[currentSample * numChannels + currentChannel] = ((float)((int32_t*)ed[currentChannel])[currentSample]) / qMax(-(float)INT_MIN, (float)INT_MAX);
		}
		break;
	}
	case AV_SAMPLE_FMT_FLTP: {
		uint8_t** ed = handle->frame->extended_data;
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			int currentChannel = i / handle->frame->nb_samples;
			int currentSample = i % handle->frame->nb_samples;
			handle->buffer[currentSample * numChannels + currentChannel] = ((float*)ed[currentChannel])[currentSample];
		}
		break;
	}
	case AV_SAMPLE_FMT_DBLP: {
		uint8_t** ed = handle->frame->extended_data;
		for (int i = 0; i < handle->frame->nb_samples * numChannels; ++i) {
			int currentChannel = i / handle->frame->nb_samples;
			int currentSample = i % handle->frame->nb_samples;
			handle->buffer[currentSample * numChannels + currentChannel] = ((double*)ed[currentChannel])[currentSample];
		}
		break;
	}
	case AV_SAMPLE_FMT_U8:
	case AV_SAMPLE_FMT_NONE:
	case AV_SAMPLE_FMT_NB:
	default:
		numberRead = 0;
		goto out;
	}

out:
	return numberRead;
}
#else
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(53, 35, 0)
static int decodeAudio(FfmpegInput::Handle* handle, int* frame_size_ptr)
{
	int ret, got_frame = 0;

	if (!handle->frame) {
		return AVERROR(ENOMEM);
	}

#if LIBAVFORMAT_VERSION_INT < AV_VERSION_INT(55, 0, 0)
	if (handle->codecContext->get_buffer != avcodec_default_get_buffer) {
		handle->codecContext->get_buffer = avcodec_default_get_buffer;
		handle->codecContext->release_buffer = avcodec_default_release_buffer;
	}
#endif

	ret = avcodec_decode_audio4(handle->codecContext, handle->frame, &got_frame, handle->packet);

	if (ret >= 0 && got_frame) {
		int ch, plane_size;
		int planar = av_sample_fmt_is_planar(handle->codecContext->sample_fmt);
		int data_size = av_samples_get_buffer_size(&plane_size, handle->codecContext->channels,
		                                           handle->frame->nb_samples, handle->codecContext->sample_fmt, 1);
		if (*frame_size_ptr < data_size) {
			return AVERROR(EINVAL);
		}

		memcpy(handle->audioBuffer, handle->frame->extended_data[0], plane_size);

		if (planar && handle->codecContext->channels > 1) {
			uint8_t* out = ((uint8_t*)(handle->audioBuffer)) + plane_size;
			for (ch = 1; ch < handle->codecContext->channels; ch++) {
				memcpy(out, handle->frame->extended_data[ch], plane_size);
				out += plane_size;
			}
		}
		*frame_size_ptr = data_size;
	}
	else {
		*frame_size_ptr = 0;
	}

	return ret;
}
#endif

size_t FfmpegInput::readOnePacket()
{
	if (!handle) {
		return 0;
	}

next_frame:
	for (;;) {
		if (av_read_frame(handle->formatContext, handle->packet) < 0) {
			return 0;
		}
		if (handle->packet->stream_index == handle->audioStream) {
			break;
		}
		AV_FREE(handle->packet);
	}

	size_t numberRead = 0;
	int dataSize = BUFFER_SIZE;
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(53, 35, 0)
	int len = decodeAudio(handle, &dataSize);
#elif LIBAVCODEC_VERSION_MAJOR >= 53 || (LIBAVCODEC_VERSION_MAJOR == 52 && LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(52, 23, 0))
	int len = avcodec_decode_audio3(handle->codecContext, handle->audioBuffer, &dataSize, handle->packet);
#else
	int len = avcodec_decode_audio2(handle->codecContext, handle->audioBuffer, &dataSize, handle->packet->data, handle->packet->size);
#endif
	if (len < 0) {
		goto out;
	}

	/* No data used, (happens with metadata frames for example) */
	if (len <= 0) {
		AV_FREE(handle->packet);
		goto next_frame;
	}
	switch (handle->codecContext->sample_fmt) {
	case SAMPLE_FMT_S16: {
		int16_t* dataShort = (int16_t*)handle->audioBuffer;
		numberRead = (size_t)dataSize / sizeof(int16_t) / (size_t)handle->codecContext->channels;
		for (unsigned int i = 0; i < (size_t)dataSize / sizeof(int16_t); ++i) {
			handle->buffer[i] = ((float)dataShort[i]) / qMax(-(float)SHRT_MIN, (float)SHRT_MAX);
		}
		break;
	}
	case SAMPLE_FMT_S32: {
		int32_t* dataInt = (int32_t*)handle->audioBuffer;
		numberRead = (size_t)dataSize / sizeof(int32_t) / (size_t)handle->codecContext->channels;
		for (unsigned int i = 0; i < (size_t)dataSize / sizeof(int32_t); ++i) {
			handle->buffer[i] = ((float)dataInt[i]) / qMax(-(float)INT_MIN, (float)INT_MAX);
		}
		break;
	}
	case SAMPLE_FMT_FLT: {
		float* dataFloat = (float*)handle->audioBuffer;
		numberRead = (size_t)dataSize / sizeof(float) / (size_t)handle->codecContext->channels;
		for (unsigned int i = 0; i < (size_t)dataSize / sizeof(float); ++i) {
			handle->buffer[i] = dataFloat[i];
		}
		break;
	}
	case SAMPLE_FMT_DBL: {
		double* dataDouble = (double*)handle->audioBuffer;
		numberRead = (size_t)dataSize / sizeof(double) / (size_t)handle->codecContext->channels;
		for (unsigned int i = 0; i < (size_t)dataSize / sizeof(double); ++i) {
			handle->buffer[i] = (float)dataDouble[i];
		}
		break;
	}
	case SAMPLE_FMT_U8:
	case SAMPLE_FMT_NONE:
	case SAMPLE_FMT_NB:
	default:
		numberRead = 0;
		goto out;
	}

out:
	AV_FREE(handle->packet);
	return numberRead;
}
#endif

bool FfmpegInput::isFloatCodec() const
{
	return handle && handle->codec && QString(QLatin1String(handle->codec->name)).endsWith(QLatin1String("float"));
}
