/* Spa A2DP AAC-ELD-A(irpods) codec */
/* SPDX-FileCopyrightText: Copyright © 2020 Wim Taymans */
/* SPDX-FileCopyrightText: Copyright © 2026 Pauli Virtanen */
/* SPDX-License-Identifier: MIT */

#include <unistd.h>
#include <stddef.h>
#include <errno.h>
#include <arpa/inet.h>

#include <spa/param/audio/format.h>
#include <spa/utils/dict.h>

#include <fdk-aac/aacenc_lib.h>
#include <fdk-aac/aacdecoder_lib.h>

#include "rtp.h"
#include "media-codecs.h"
#include "aac-bits.h"

static struct spa_log *global_log;

#define DEFAULT_AAC_BITRATE	256000

#define FRAMES_PER_PACKET	3
#define FRAME_SAMPLES		480
#define SAMPLE_SIZE		2 /* S16 */

struct frame_header {
	uint8_t data[4];
} __attribute__((packed));

struct impl {
	HANDLE_AACENCODER enc;
	HANDLE_AACDECODER dec;

	int codesize;

	uint32_t enc_delay;
	uint32_t dec_delay;

	uint16_t enc_seq;
};

static bool parse_frame_header(const struct frame_header *hdr, uint16_t *seq, uint16_t *size)
{
	const uint8_t *buf = hdr->data;

	if ((buf[0] & 0xf0) != 0xb0)
		return false;
	*seq = (buf[0] & 0x0f) << 8 | buf[1];

	if ((buf[2] & 0xf0) != 0x10)
		return false;
	*size = (buf[2] & 0x0f) << 8 | buf[3];

	return true;
}

static void write_frame_header(struct frame_header *hdr, uint16_t seq, uint16_t size)
{
	uint8_t *buf = hdr->data;

	buf[0] = 0xb0 | (seq & 0xf00) >> 8;
	buf[1] = seq & 0xff;
	buf[2] = 0x10 | (size & 0xf00) >> 8;
	buf[3] = size & 0xff;
}

static bool eld_supported(void)
{
	static bool supported = false, checked = false;
	HANDLE_AACENCODER enc = NULL;

	if (checked)
		return supported;

	if (aacEncOpen(&enc, 0, 2) != AACENC_OK)
		goto done;
	if (aacEncoder_SetParam(enc, AACENC_AOT, AOT_ER_AAC_ELD) != AACENC_OK)
		goto done;
	if (aacEncoder_SetParam(enc,  AACENC_SBR_MODE, 1) != AACENC_OK)
		goto done;

	supported = true;

done:
	if (enc)
		aacEncClose(&enc);
	checked = true;
	spa_log_debug(global_log, "FDK-AAC AAC-ELD-A support:%d", (int)supported);
	return supported;
}

static int codec_fill_caps(const struct media_codec *codec, uint32_t flags,
		const struct spa_dict *settings, uint8_t caps[A2DP_MAX_CAPS_SIZE])
{
	const a2dp_aac_eld_a_t a2dp_aac = {
		.info = codec->vendor,
		AAC_ELD_A_INIT_AOT(AAC_ELD_A_AOT_AAC_ELD),
		AAC_ELD_A_INIT_FREQ_CH(AAC_ELD_A_FREQ_48000, AAC_ELD_A_CH_MONO | AAC_ELD_A_CH_STEREO),
		AAC_ELD_A_INIT_FLAGS_BITRATE(AAC_ELD_A_FLAG_VBR, DEFAULT_AAC_BITRATE),
	};

	if (!eld_supported())
		return -ENOTSUP;

	memcpy(caps, &a2dp_aac, sizeof(a2dp_aac));
	return sizeof(a2dp_aac);
}

static int codec_select_config(const struct media_codec *codec, uint32_t flags,
		const void *caps, size_t caps_size,
		const struct media_codec_audio_info *info,
		const struct spa_dict *settings, uint8_t config[A2DP_MAX_CAPS_SIZE],
		void **config_data)
{
	a2dp_aac_eld_a_t conf;
	int aot, freq, ch, cflags, bitrate;

	if (!eld_supported())
		return -ENOTSUP;

	if (caps_size < sizeof(conf))
		return -EINVAL;

	memcpy(&conf, caps, sizeof(conf));

	if (codec->vendor.vendor_id != conf.info.vendor_id ||
	    codec->vendor.codec_id != conf.info.codec_id)
		return -ENOTSUP;

	if (AAC_ELD_A_GET_AOT(conf) & AAC_ELD_A_AOT_AAC_ELD)
		aot = AAC_ELD_A_AOT_AAC_ELD;
	else
		return -EINVAL;

	if (AAC_ELD_A_GET_FREQ(conf) & AAC_ELD_A_FREQ_48000)
		freq = AAC_ELD_A_FREQ_48000;
	else
		return -EINVAL;

	if (AAC_ELD_A_GET_CH(conf) & AAC_ELD_A_CH_STEREO)
		ch = AAC_ELD_A_CH_STEREO;
	else if (AAC_ELD_A_GET_CH(conf) & AAC_ELD_A_CH_MONO)
		ch = AAC_ELD_A_CH_MONO;
	else
		return -EINVAL;

	if (AAC_ELD_A_GET_FLAGS(conf) & AAC_ELD_A_FLAG_VBR)
		cflags = AAC_ELD_A_FLAG_VBR;
	else
		cflags = 0;

	bitrate = AAC_ELD_A_GET_BITRATE(conf);

	conf = (a2dp_aac_eld_a_t) {
		.info = conf.info,
		AAC_ELD_A_INIT_AOT(aot),
		AAC_ELD_A_INIT_FREQ_CH(freq, ch),
		AAC_ELD_A_INIT_FLAGS_BITRATE(cflags, bitrate),
	};

	memcpy(config, &conf, sizeof(conf));

	return sizeof(conf);
}

static int codec_validate_config(const struct media_codec *codec, uint32_t flags,
			const void *caps, size_t caps_size,
			struct spa_audio_info *info)
{
	a2dp_aac_eld_a_t conf;

	if (caps == NULL || caps_size < sizeof(conf))
		return -EINVAL;

	memcpy(&conf, caps, sizeof(conf));

	spa_zero(*info);
	info->media_type = SPA_MEDIA_TYPE_audio;
	info->media_subtype = SPA_MEDIA_SUBTYPE_raw;
	info->info.raw.format = SPA_AUDIO_FORMAT_S16;

	if (AAC_ELD_A_GET_AOT(conf) != AAC_ELD_A_AOT_AAC_ELD)
		return -EINVAL;

	switch (AAC_ELD_A_GET_FREQ(conf)) {
	case AAC_ELD_A_FREQ_48000:
		info->info.raw.rate = 48000;
		break;
	default:
		return -EINVAL;
	}

	switch (AAC_ELD_A_GET_CH(conf)) {
	case AAC_ELD_A_CH_STEREO:
		info->info.raw.channels = 2;
		info->info.raw.position[0] = SPA_AUDIO_CHANNEL_FL;
		info->info.raw.position[1] = SPA_AUDIO_CHANNEL_FR;
		break;
	case AAC_ELD_A_CH_MONO:
		info->info.raw.channels = 1;
		info->info.raw.position[0] = SPA_AUDIO_CHANNEL_MONO;
		break;
	default:
		return -EINVAL;
	}

	return 0;
}

static int codec_enum_config(const struct media_codec *codec, uint32_t flags,
		const void *caps, size_t caps_size, uint32_t id, uint32_t idx,
		struct spa_pod_builder *b, struct spa_pod **param)
{
	struct spa_audio_info info;
	struct spa_pod_frame f[1];
	int res;

	if ((res = codec_validate_config(codec, flags, caps, caps_size, &info)) < 0)
		return res;

	if (idx > 0)
		return 0;

	spa_pod_builder_push_object(b, &f[0], SPA_TYPE_OBJECT_Format, id);
	spa_pod_builder_add(b,
			SPA_FORMAT_mediaType,      SPA_POD_Id(info.media_type),
			SPA_FORMAT_mediaSubtype,   SPA_POD_Id(info.media_subtype),
			SPA_FORMAT_AUDIO_format,   SPA_POD_Id(info.info.raw.format),
			SPA_FORMAT_AUDIO_rate,     SPA_POD_Int(info.info.raw.rate),
			SPA_FORMAT_AUDIO_channels, SPA_POD_Int(info.info.raw.channels),
			SPA_FORMAT_AUDIO_position, SPA_POD_Array(sizeof(uint32_t),
					SPA_TYPE_Id, info.info.raw.channels, info.info.raw.position),
			0);

	*param = spa_pod_builder_pop(b, &f[0]);
	return *param == NULL ? -EIO : 1;
}

static void *codec_init(const struct media_codec *codec, uint32_t flags,
		void *config, size_t config_len, const struct spa_audio_info *info,
		void *props, size_t mtu)
{
	struct impl *this = NULL;
	a2dp_aac_eld_a_t conf;
	struct spa_audio_info config_info;
	uint8_t asc_buf[8];
	UCHAR *asc[1] = { (void *)asc_buf };
	UINT asc_len[1];
	UINT bitratemode;
	int bitrate, max_bitrate;
	int res;
	size_t rate, channels;

	if (config_len < sizeof(conf)) {
		res = -EINVAL;
		goto error;
	}
	memcpy(&conf, config, sizeof(conf));

	if (info->media_type != SPA_MEDIA_TYPE_audio ||
	    info->media_subtype != SPA_MEDIA_SUBTYPE_raw ||
	    info->info.raw.format != SPA_AUDIO_FORMAT_S16) {
		res = -EINVAL;
		goto error;
	}

	if ((res = codec_validate_config(codec, flags, config, config_len, &config_info)) < 0)
		goto error;
	if (config_info.info.raw.channels != info->info.raw.channels ||
			config_info.info.raw.rate != info->info.raw.rate) {
		res = -EINVAL;
		goto error;
	}

	this = calloc(1, sizeof(struct impl));
	if (this == NULL) {
		res = -errno;
		goto error;
	}

	mtu = mtu;
	rate = config_info.info.raw.rate;
	channels = config_info.info.raw.channels;

	/* No fragmentation: make sure it fits into MTU */
	if (mtu <= sizeof(struct rtp_header) + sizeof(struct frame_header) * FRAMES_PER_PACKET)
		goto error;

	max_bitrate = (mtu - sizeof(struct rtp_header) - sizeof(struct frame_header) * FRAMES_PER_PACKET)
		* 8 * rate / FRAME_SAMPLES / FRAMES_PER_PACKET;
	bitrate = SPA_MIN(max_bitrate, AAC_ELD_A_GET_BITRATE(conf));

	if (AAC_ELD_A_GET_FLAGS(conf) & AAC_ELD_A_FLAG_VBR)
		bitratemode = 5;
	else
		bitratemode = 0;

	/*
	 * Encoder
	 */
	res = aacEncOpen(&this->enc, 0, channels);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_AOT, AOT_ER_AAC_ELD);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc,  AACENC_SBR_RATIO, 1);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc,  AACENC_SBR_MODE, 1);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_CHANNELMODE, channels);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_SAMPLERATE, rate);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_GRANULE_LENGTH, FRAME_SAMPLES);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_PEAK_BITRATE, max_bitrate);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_BITRATE, bitrate);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_BITRATEMODE, bitratemode);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_TRANSMUX, TT_MP4_RAW);
	if (res != AACENC_OK)
		goto error;

	res = aacEncoder_SetParam(this->enc, AACENC_AFTERBURNER, 1);
	if (res != AACENC_OK)
		goto error;

	res = aacEncEncode(this->enc, NULL, NULL, NULL, NULL);
	if (res != AACENC_OK)
		goto error;

	AACENC_InfoStruct enc_info = {};
	res = aacEncInfo(this->enc, &enc_info);
	if (res != AACENC_OK)
		goto error;

	this->enc_delay = enc_info.nDelay;

	if (enc_info.frameLength != FRAME_SAMPLES)
		goto error;

	this->codesize = channels * FRAME_SAMPLES * SAMPLE_SIZE * FRAMES_PER_PACKET;

	/*
	 * Decoder
	 */
	this->dec = aacDecoder_Open(TT_MP4_RAW, 1);
	if (!this->dec) {
		res = -EINVAL;
		goto error;
	}
#ifdef AACDECODER_LIB_VL0
	res = aacDecoder_SetParam(this->dec, AAC_PCM_MIN_OUTPUT_CHANNELS, channels);
	if (res != AAC_DEC_OK)
		goto error;
	res = aacDecoder_SetParam(this->dec, AAC_PCM_MAX_OUTPUT_CHANNELS, channels);
	if (res != AAC_DEC_OK)
		goto error;
#else
	res = aacDecoder_SetParam(this->dec, AAC_PCM_OUTPUT_CHANNELS, channels);
	if (res != AAC_DEC_OK)
		goto error;
#endif
	res = aac_make_asc(asc_buf, sizeof(asc_buf), AAC_AOT_ER_AAC_ELD,
			rate, rate, channels, true);
	if (res < 0)
		goto error;
	asc_len[0] = res;
	res = aacDecoder_ConfigRaw(this->dec, asc, asc_len);
	if (res != AAC_DEC_OK)
		goto error;

	this->dec_delay = 0;

	return this;

error:
	if (this && this->enc)
		aacEncClose(&this->enc);
	if (this && this->dec)
		aacDecoder_Close(this->dec);
	free(this);
	errno = -res;
	return NULL;
}

static void codec_deinit(void *data)
{
	struct impl *this = data;
	if (this->enc)
		aacEncClose(&this->enc);
	if (this->dec)
		aacDecoder_Close(this->dec);
	free(this);
}

static int codec_get_block_size(void *data)
{
	struct impl *this = data;
	return this->codesize;
}

static int codec_start_encode (void *data,
		void *dst, size_t dst_size, uint16_t seqnum, uint32_t timestamp)
{
	struct rtp_header *header;

	if (dst_size < sizeof(struct rtp_header))
		return -ENOSPC;

	header = (struct rtp_header *)dst;
	memset(header, 0, sizeof(struct rtp_header));

	header->v = 2;
	header->pt = 96;
	header->sequence_number = htons(seqnum);
	header->timestamp = htonl(timestamp);
	header->ssrc = htonl(1);

	return sizeof(struct rtp_header);
}

static int encode_frame(struct impl *this,
		const void *src, size_t src_size,
		void *dst, size_t dst_size,
		size_t *dst_out)
{
	struct frame_header *hdr = dst;
	int res;

	if (dst_size < sizeof(*hdr))
		return -ENOSPC;

	dst_size -= sizeof(*hdr);
	dst = SPA_PTROFF(dst, sizeof(*hdr), void);

	void *in_bufs[] = {(void *) src};
	int in_buf_ids[] = {IN_AUDIO_DATA};
	int in_buf_sizes[] = {src_size};
	int in_buf_el_sizes[] = {SAMPLE_SIZE};
	AACENC_BufDesc in_buf_desc = {
		.numBufs = 1,
		.bufs = in_bufs,
		.bufferIdentifiers = in_buf_ids,
		.bufSizes = in_buf_sizes,
		.bufElSizes = in_buf_el_sizes,
	};
	AACENC_InArgs in_args = {
		.numInSamples = src_size / SAMPLE_SIZE,
	};

	void *out_bufs[] = {dst};
	int out_buf_ids[] = {OUT_BITSTREAM_DATA};
	int out_buf_sizes[] = {dst_size};
	int out_buf_el_sizes[] = {SAMPLE_SIZE};
	AACENC_BufDesc out_buf_desc = {
		.numBufs = 1,
		.bufs = out_bufs,
		.bufferIdentifiers = out_buf_ids,
		.bufSizes = out_buf_sizes,
		.bufElSizes = out_buf_el_sizes,
	};
	AACENC_OutArgs out_args = {};

	res = aacEncEncode(this->enc, &in_buf_desc, &out_buf_desc, &in_args, &out_args);
	if (res != AACENC_OK)
		return -EINVAL;

	write_frame_header(hdr, this->enc_seq++, out_args.numOutBytes);

	*dst_out = out_args.numOutBytes + sizeof(*hdr);

	return out_args.numInSamples * SAMPLE_SIZE;
}

static int codec_encode(void *data,
		const void *src, size_t src_size,
		void *dst, size_t dst_size,
		size_t *dst_out, int *need_flush)
{
	struct impl *this = data;
	int consumed = 0;
	int i, res;

	*dst_out = 0;

	for (i = 0; i < FRAMES_PER_PACKET; ++i) {
		size_t out;

		res = encode_frame(this, src, src_size, dst, dst_size, &out);
		if (res < 0)
			return res;

		src = SPA_PTROFF(src, res, void);
		src_size -= res;
		consumed += res;

		dst = SPA_PTROFF(dst, out, void);
		dst_size -= out;
		*dst_out += out;
	}

	*need_flush = NEED_FLUSH_ALL;

	return consumed;
}

static int codec_start_decode (void *data,
		const void *src, size_t src_size, uint16_t *seqnum, uint32_t *timestamp)
{
	const struct rtp_header *header = src;

	if (src_size <= sizeof(struct rtp_header))
		return -EINVAL;

	if (seqnum)
		*seqnum = ntohs(header->sequence_number);
	if (timestamp)
		*timestamp = ntohl(header->timestamp);

	return sizeof(struct rtp_header);
}

static int decode_frame(struct impl *this,
		const void *src, size_t src_size,
		void *dst, size_t dst_size,
		size_t *dst_out)
{
	uint16_t seq, size;
	const struct frame_header *hdr = src;

	if (src_size < sizeof(*hdr))
		return -EINVAL;

	if (!parse_frame_header(hdr, &seq, &size))
		return -EINVAL;

	src_size -= sizeof(*hdr);
	src = SPA_PTROFF(src, sizeof(*hdr), void);

	uint data_size = SPA_MIN(src_size, size);
	uint bytes_valid = data_size;
	CStreamInfo *aacinf;
	int res;

	res = aacDecoder_Fill(this->dec, (UCHAR **)&src, &data_size, &bytes_valid);
	if (res != AAC_DEC_OK) {
		spa_log_debug(global_log, "AAC buffer fill error: 0x%04X", res);
		return -EINVAL;
	}

	res = aacDecoder_DecodeFrame(this->dec, dst, dst_size, 0);
	if (res != AAC_DEC_OK) {
		spa_log_debug(global_log, "AAC decode frame error: 0x%04X", res);
		return -EINVAL;
	}

	aacinf = aacDecoder_GetStreamInfo(this->dec);
	if (!aacinf) {
		spa_log_debug(global_log, "AAC get stream info failed");
		return -EINVAL;
	}
	*dst_out = aacinf->frameSize * aacinf->numChannels * SAMPLE_SIZE;

	spa_assert_se(data_size >= bytes_valid);

	return data_size - bytes_valid + sizeof(*hdr);
}

static int codec_decode(void *data,
		const void *src, size_t src_size,
		void *dst, size_t dst_size,
		size_t *dst_out)
{
	struct impl *this = data;
	int consumed = 0;
	int res;

	*dst_out = 0;

	while (src_size) {
		size_t out;

		res = decode_frame(this, src, src_size, dst, dst_size, &out);
		if (res < 0)
			return res;

		src = SPA_PTROFF(src, res, void);
		src_size -= res;
		consumed += res;

		dst = SPA_PTROFF(dst, out, void);
		dst_size -= out;
		*dst_out += out;
	}

	return consumed;
}

static void codec_get_delay(void *data, uint32_t *encoder, uint32_t *decoder)
{
	struct impl *this = data;

	if (encoder)
		*encoder = this->enc_delay;

	if (decoder) {
		CStreamInfo *info = aacDecoder_GetStreamInfo(this->dec);
		if (info)
			this->dec_delay = info->outputDelay;
		*decoder = this->dec_delay;
	}
}

static void codec_set_log(struct spa_log *log_)
{
	global_log = log_;
	spa_log_topic_init(global_log, &codec_plugin_log_topic);
}

const struct media_codec a2dp_codec_aac_eld_a = {
	.id = SPA_BLUETOOTH_AUDIO_CODEC_AAC_ELD_A,
	.kind = MEDIA_CODEC_A2DP,
	.codec_id = A2DP_CODEC_VENDOR,
	.vendor = { .vendor_id = AAC_ELD_A_VENDOR_ID,
			.codec_id = AAC_ELD_A_CODEC_ID },
	.name = "aac_eld_a",
	.description = "AAC-ELD-A",
	.fill_caps = codec_fill_caps,
	.select_config = codec_select_config,
	.enum_config = codec_enum_config,
	.validate_config = codec_validate_config,
	.init = codec_init,
	.deinit = codec_deinit,
	.get_block_size = codec_get_block_size,
	.start_encode = codec_start_encode,
	.encode = codec_encode,
	.start_decode = codec_start_decode,
	.decode = codec_decode,
	.set_log = codec_set_log,
	.get_delay = codec_get_delay,
};

MEDIA_CODEC_EXPORT_DEF(
	"aac-eld-a",
	&a2dp_codec_aac_eld_a
);