/*
 * Copyright (c) 2010-2022 Belledonne Communications SARL.
 *
 * This file is part of mediastreamer2
 * (see https://gitlab.linphone.org/BC/public/mediastreamer2).
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include "android_mediacodec.h"

#include "media-codec-decoder.h"

using namespace mediastreamer;
using namespace std;

namespace mediastreamer {

MediaCodecDecoder::MediaCodecDecoder(const std::string &mime) : H26xDecoder(mime) {
	try {
		_impl = AMediaCodec_createDecoderByType(mime.c_str());
		if (_impl == nullptr) {
			ostringstream msg;
			msg << "could not create MediaCodec for '" << mime << "'";
			throw runtime_error(msg.str());
		}
		_format = createFormat(true);
		_bufAllocator = ms_yuv_buf_allocator_new();
		ms_yuv_buf_allocator_set_max_frames(_bufAllocator, 15); /* we need to accomodate decoding with 'B-frames' */
		_naluHeader.reset(H26xToolFactory::get(mime).createNaluHeader());
		_psStore.reset(H26xToolFactory::get(mime).createParameterSetsStore());
		startImpl();
	} catch (const runtime_error &e) {
		ms_error("MediaCodecDecoder: %s. Dumping decoding context", e.what());
		if (_impl) AMediaCodec_delete(_impl);
		_impl = nullptr;
	}
}

MediaCodecDecoder::~MediaCodecDecoder() {
	if (_impl) AMediaCodec_delete(_impl);
	if (_format) AMediaFormat_delete(_format);
	if (_bufAllocator) ms_yuv_buf_allocator_free(_bufAllocator);
}

bool MediaCodecDecoder::setParameterSets(MSQueue *parameterSets, uint64_t timestamp) {
	if (!feed(parameterSets, timestamp, true)) {
		ms_error("MediaCodecDecoder: parameter sets has been refused by the decoder.");
		return false;
	}
	_needParameters = false;
	return true;
}

bool MediaCodecDecoder::feed(MSQueue *encodedFrame, uint64_t timestamp) {
	bool status = false;

	if (_impl == nullptr) {
		ms_queue_flush(encodedFrame);
		return true;
	}

	_psStore->extractAllPs(encodedFrame);
	if (_psStore->hasNewParameters()) {
		ms_message("MediaCodecDecoder: new paramter sets received");
		_needParameters = true;
		_psStore->acknowlege();
	}

	if (_needParameters && _psStore->psGatheringCompleted()) {
		MSQueue parameters;
		ms_queue_init(&parameters);
		_psStore->fetchAllPs(&parameters);
		if (!setParameterSets(&parameters, timestamp)) {
			ms_error("MediaCodecDecoder: waiting for parameter sets.");
			goto clean;
		}
	}

	if (_needParameters) {
		ms_error("MediaCodecDecoder: missing parameter sets");
		goto clean;
	}

	if (_needKeyFrame) {
		if (!isKeyFrame(encodedFrame)) {
			ms_error("MediaCodecDecoder: waiting for key frame.");
			goto clean;
		}
		ms_message("MediaCodecDecoder: key frame received");
		_needKeyFrame = false;
	}

	if (!feed(encodedFrame, timestamp, false)) {
		goto clean;
	}

	_pendingFrames++;
	status = true;

clean:
	ms_queue_flush(encodedFrame);
	return status;
}

MediaCodecDecoder::Status MediaCodecDecoder::fetch(mblk_t *&frame) {
	AMediaImage image = {0};
	int dst_pix_strides[4] = {1, 1, 1, 1};
	MSRect dst_roi = {0};
	AMediaCodecBufferInfo info;
	ssize_t oBufidx = -1;
	Status status = NoError;

	frame = nullptr;

	if (_impl == nullptr || _pendingFrames <= 0) {
		status = NoFrameAvailable;
		goto end;
	}
	{
		TimeReport dequeue("dequeueOutputBuffer");
		oBufidx = AMediaCodec_dequeueOutputBuffer(_impl, &info, _timeoutUs);
		while (oBufidx == AMEDIACODEC_INFO_OUTPUT_FORMAT_CHANGED ||
		       oBufidx == AMEDIACODEC_INFO_OUTPUT_BUFFERS_CHANGED) {
			ms_message("MediaCodecDecoder: %s", codecInfoToString(oBufidx).c_str());
			if (oBufidx == AMEDIACODEC_INFO_OUTPUT_FORMAT_CHANGED) {
				AMediaFormat *format = AMediaCodec_getOutputFormat(_impl);
				AMediaFormat_getInt32(format, "width", &_curWidth);
				AMediaFormat_getInt32(format, "height", &_curHeight);
				ms_message("MediaCodecDecoder: new format %ix%i :\n%s", _curWidth, _curHeight,
				           AMediaFormat_toString(format));
				AMediaFormat_delete(format);
			}
			oBufidx = AMediaCodec_dequeueOutputBuffer(_impl, &info, _timeoutUs);
		}
		dequeue.finished();
	}

	if (oBufidx < 0) {
		if (oBufidx == AMEDIACODEC_INFO_TRY_AGAIN_LATER) {
			return NoFrameAvailable;
		} else {
			ms_error("MediaCodecDecoder: error while dequeueing an output buffer: %s",
			         codecInfoToString(oBufidx).c_str());
			if (oBufidx == AMEDIA_ERROR_UNKNOWN) {
				try {
					resetImpl();
					startImpl();
				} catch (const runtime_error &e) {
					ms_error("MediaCodecDecoder: decoding session couldn't been started. The session is definitively "
					         "lost !");
				}
			}
			return DecodingFailure;
		}
		goto end;
	}

	_pendingFrames--;
	{
		TimeReport getOutputImage("AMediaCodec_getOutputImage");
		if (AMediaCodec_getOutputImage(_impl, oBufidx, &image) <= 0) {
			ms_error("MediaCodecDecoder: AMediaCodec_getOutputImage() failed");
			status = DecodingFailure;
			goto end;
		}
		getOutputImage.finished();
	}

	if (_curWidth && _curHeight && (_curWidth != image.crop_rect.w || _curHeight != image.crop_rect.h)) {
		/*
		 * Sometimes (not all devices), we observe an insconsistency between the width/height announced by the
		 * MediaCodec.getOutputFormat() and the width/height announced by the getOutputImage(). In some cases this also
		 * results in an garbled decoded image, but sometimes not. We print out the information, but no further action
		 * can be done here.
		 */
		ms_error("Mismatch between decoder new format and output image detected: %ix%i vs %ix%i", _curWidth, _curHeight,
		         image.crop_rect.w, image.crop_rect.h);
		_curWidth = image.crop_rect.w;
		_curHeight = image.crop_rect.h;
	}
	// ms_message("MediaCodecDecoder: got output image with presentation timestamp [%llu] ms.", (unsigned long long)
	// image.timestamp/1000000ULL);

	MSPicture pic;
	frame = ms_yuv_buf_allocator_get(_bufAllocator, &pic, image.crop_rect.w, image.crop_rect.h);
	// ms_message("image.crop_rect.w=%i, image.crop_rect.h=%i, image.row_strides=%i,%i,%i image.pixel_strides=%i,%i,%i",
	//	   image.crop_rect.w, image.crop_rect.h, image.row_strides[0], image.row_strides[1], image.row_strides[2],
	//		image.pixel_strides[0], image.pixel_strides[1], image.pixel_strides[2]);
	if (frame) {
		ms_yuv_buf_copy_with_pix_strides(image.buffers, image.row_strides, image.pixel_strides, image.crop_rect,
		                                 pic.planes, pic.strides, dst_pix_strides, dst_roi);
		mblk_set_timestamp_info(frame, (uint32_t)((image.timestamp / 1000000LL) * 90LL));
	} else status = VideoDecoder::Status::NoFrameAvailable;

	AMediaImage_close(&image);

end:
	if (oBufidx >= 0) AMediaCodec_releaseOutputBuffer(_impl, oBufidx, FALSE);
	return status;
}

AMediaFormat *MediaCodecDecoder::createFormat(bool withLowLatency) const {
	AMediaFormat *format = AMediaFormat_new();
	AMediaFormat_setString(format, "mime", getMime().c_str());
	AMediaFormat_setInt32(format, "color-format", 0x7f420888);
	AMediaFormat_setInt32(format, "max-width", 1920);
	AMediaFormat_setInt32(format, "max-height", 1920);
	AMediaFormat_setInt32(format, "priority", 0);
	if (withLowLatency) AMediaFormat_setInt32(format, "low-latency", 1); // MediaFormat.KEY_LOW_LATENCY
	return format;
}

void MediaCodecDecoder::startImpl() {
	media_status_t status = AMEDIA_OK;
	ostringstream errMsg;
	ms_message("MediaCodecDecoder: starting decoder with following parameters:\n%s", AMediaFormat_toString(_format));
	if ((status = AMediaCodec_configure(_impl, _format, nullptr, nullptr, 0)) != AMEDIA_OK) {
		ms_warning("MediaCodecDecoder: configure failed, retrying without low-latency parameter.");
		AMediaFormat_delete(_format);
		_format = createFormat(false);
		if ((status = AMediaCodec_configure(_impl, _format, nullptr, nullptr, 0)) != AMEDIA_OK) {
			errMsg << "configuration failure: " << int(status);
			throw runtime_error(errMsg.str());
		}
	}

	if ((status = AMediaCodec_start(_impl)) != AMEDIA_OK) {
		errMsg << "starting failure: " << int(status);
		throw runtime_error(errMsg.str());
	}

	ms_message("MediaCodecDecoder: decoder successfully started. In-force parameters:\n%s",
	           AMediaFormat_toString(_format));
}

void MediaCodecDecoder::stopImpl() noexcept {
	ms_message("MediaCodecDecoder: stopping decoder");
	AMediaCodec_stop(_impl);
}

void MediaCodecDecoder::resetImpl() noexcept {
	ms_message("MediaCodecDecoder: reseting decoder");
	if (AMediaCodec_reset(_impl) != AMEDIA_OK) {
		ms_error("MediaCodecDecoder: decoder couldn't been reset. Throwing decoding session out");
		AMediaCodec_delete(_impl);
		_impl = AMediaCodec_createDecoderByType(_mime.c_str());
		if (_impl == nullptr)
			ms_error(
			    "MediaCodecDecoder: couldn't recreate decoding session. The decoding session is definitively lost !");
	}
}

bool MediaCodecDecoder::feed(MSQueue *encodedFrame, uint64_t timestamp, bool isPs) {
	if (_impl == nullptr) return false;

	ssize_t iBufidx = AMediaCodec_dequeueInputBuffer(_impl, _timeoutUs);
	if (iBufidx < 0) {
		ms_error("MediaCodecDecoder: %s.", iBufidx == -1
		                                       ? "no buffer available for queuing this frame ! Decoder is too slow"
		                                       : "AMediaCodec_dequeueInputBuffer() had an exception");
		return false;
	}

	size_t bufsize;
	size_t streamSize;
	uint8_t *buf = AMediaCodec_getInputBuffer(_impl, iBufidx, &bufsize);
	if (buf == nullptr) {
		ms_error("MediaCodecDecoder: AMediaCodec_getInputBuffer() returned NULL");
		return false;
	}

	try {
		streamSize = H26xUtils::nalusToByteStream(encodedFrame, buf, bufsize);
	} catch (const invalid_argument &e) {
		ms_error(
		    "MediaCodecDecoder: cannot convert the all the nal units content into the input buffer size : bufsize=%zu",
		    bufsize);
	}

	uint32_t flags = isPs ? BufferFlag::CodecConfig : BufferFlag::None;
	if (AMediaCodec_queueInputBuffer(_impl, iBufidx, 0, streamSize, timestamp * 1000ULL, flags) != 0) {
		ms_error("MediaCodecDecoder: AMediaCodec_queueInputBuffer() had an exception");
		return false;
	}

	return true;
}

bool MediaCodecDecoder::isKeyFrame(const MSQueue *frame) const {
	for (const mblk_t *nalu = ms_queue_peek_first(frame); !ms_queue_end(frame, nalu);
	     nalu = ms_queue_next(frame, nalu)) {
		_naluHeader->parse(nalu->b_rptr);
		if (_naluHeader->getAbsType().isKeyFramePart()) return true;
	}
	return false;
}

std::string MediaCodecDecoder::codecInfoToString(ssize_t codecStatusCode) {
	switch (codecStatusCode) {
		case AMEDIA_ERROR_UNKNOWN:
			return "MediaCodec had an exception";
		case AMEDIACODEC_INFO_OUTPUT_BUFFERS_CHANGED:
			return "output buffers has changed";
		case AMEDIACODEC_INFO_OUTPUT_FORMAT_CHANGED:
			return "output format has changed";
		case AMEDIACODEC_INFO_TRY_AGAIN_LATER:
			return "no output buffer available";
		default:
			break;
	}
	ostringstream os;
	if (codecStatusCode >= 0) {
		os << "unqueued buffer (index=" << codecStatusCode << ")";
	} else {
		os << "unknown error (" << codecStatusCode << ")";
	}
	return os.str();
}

} // namespace mediastreamer
