// https://www.rfc-editor.org/rfc/rfc9328.html

#include "rtp-packet.h"
#include "rtp-payload-internal.h"
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <errno.h>

#define H266_RTP_AP 28
#define H266_RTP_FU 29

/*
+---------------+---------------+
|0|1|2|3|4|5|6|7|0|1|2|3|4|5|6|7|
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|F|Z| LayerID   |  Type   | TID |
+---------------+---------------+
*/

#define H266_TYPE(v) (((v) >> 3) & 0x1f)

/*
+---------------+
|0|1|2|3|4|5|6|7|
+-+-+-+-+-+-+-+-+
|S|E|P|  FuType |
+---------------+
*/
#define FU_START(v) ((v) & 0x80)
#define FU_END(v)	((v) & 0x40)
#define FU_MARK(v)	((v) & 0x20)
#define FU_NAL(v)	((v) & 0x1F)

struct rtp_decode_h266_t
{
	struct rtp_payload_t handler;
	void* cbparam;

	uint16_t seq; // rtp seq
	uint32_t timestamp;

	uint8_t* ptr;
	int size, capacity;

	int flags;
	int using_donl_field;
};

static void* rtp_h266_unpack_create(struct rtp_payload_t* handler, void* param)
{
	struct rtp_decode_h266_t* unpacker;
	unpacker = (struct rtp_decode_h266_t*)calloc(1, sizeof(*unpacker));
	if (!unpacker)
		return NULL;

	memcpy(&unpacker->handler, handler, sizeof(unpacker->handler));
	unpacker->cbparam = param;
	unpacker->flags = -1;
	return unpacker;
}

static void rtp_h266_unpack_destroy(void* p)
{
	struct rtp_decode_h266_t* unpacker;
	unpacker = (struct rtp_decode_h266_t*)p;

	if (unpacker->ptr)
		free(unpacker->ptr);
#if defined(_DEBUG) || defined(DEBUG)
	memset(unpacker, 0xCC, sizeof(*unpacker));
#endif
	free(unpacker);
}

// 4.3.2. Aggregation Packets (APs) 
/*
 0               1               2               3
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                          RTP Header                           |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|      PayloadHdr (Type=28)     |           NALU 1 DONL         |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|           NALU 1 Size         |            NALU 1 HDR         |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                                                               |
|                         NALU 1 Data . . .                     |
|                                                               |
+     . . .     +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|               |  NALU 2 DOND  |            NALU 2 Size        |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|          NALU 2 HDR           |                               |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+            NALU 2 Data        |
|                                                               |
|         . . .                 +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                               :    ...OPTIONAL RTP padding    |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
*/
static int rtp_h266_unpack_ap(struct rtp_decode_h266_t* unpacker, const uint8_t* ptr, int bytes, uint32_t timestamp)
{
	int r;
	int n;
	int len;
	//uint16_t donl;
	//uint16_t dond;

	//donl = unpacker->using_donl_field ? nbo_r16(ptr + 2) : 0;
	ptr += 2; // PayloadHdr
	n = 2 /*LEN*/ + (unpacker->using_donl_field ? 2 : 0);
	r = 0;

	for (bytes -= 2 /*PayloadHdr*/; 0 == r && bytes > n; bytes -= len + 2)
	{
		bytes -= n - 2; // skip DON
		ptr += n - 2; // skip DON
		len = nbo_r16(ptr);
		if (len + 2 > bytes || len < 3)
		{
			assert(0);
			unpacker->flags = RTP_PAYLOAD_FLAG_PACKET_LOST;
			unpacker->size = 0;
			return -EINVAL; // error
		}

		assert(H266_TYPE(ptr[3]) >= 0 && H266_TYPE(ptr[3]) < 31);
		r = unpacker->handler.packet(unpacker->cbparam, ptr + 2, len, timestamp, unpacker->flags);
		unpacker->flags = 0;
		unpacker->size = 0;

		ptr += len + 2; // next NALU
		n = 2 /*LEN*/ + (unpacker->using_donl_field ? 1 : 0);
	}

	return 0 == r ? 1 : r; // packet handled
}

// 4.3.3. Fragmentation Units
/*
 0               1               2               3
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|     PayloadHdr (Type=29)      |    FU header  |  DONL (cond)  |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-|
|  DONL (cond)  |                                               |
|-+-+-+-+-+-+-+-+                                               |
|                           FU payload                          |
|                                                               |
|                               +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                               :    ...OPTIONAL RTP padding    |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

+---------------+
|0|1|2|3|4|5|6|7|
+-+-+-+-+-+-+-+-+
|S|E|   FuType  |
+---------------+
*/
static int rtp_h266_unpack_fu(struct rtp_decode_h266_t* unpacker, const uint8_t* ptr, int bytes, uint32_t timestamp)
{
	int r, n;
	uint8_t fuheader;

	r = 0;
	n = 1 /*FU header*/ + (unpacker->using_donl_field ? 4 : 2);
	if (bytes < n || unpacker->size + bytes - n > RTP_PAYLOAD_MAX_SIZE)
	{
		assert(0);
		return -EINVAL;
	}

	if (unpacker->size + bytes - n + 2 /*NALU*/ > unpacker->capacity)
	{
		void* p = NULL;
		int size = unpacker->size + bytes + 2;
		size += size / 4 > 128000 ? size / 4 : 128000;
		p = realloc(unpacker->ptr, size);
		if (!p)
		{
			// set packet lost flag
			unpacker->flags = RTP_PAYLOAD_FLAG_PACKET_LOST;
			unpacker->size = 0;
			return -ENOMEM;
		}
		unpacker->ptr = (uint8_t*)p;
		unpacker->capacity = size;
	}

	fuheader = ptr[2];
	if (FU_START(fuheader))
	{
#if 0
		if (unpacker->size > 0)
		{
			unpacker->flags |= RTP_PAYLOAD_FLAG_PACKET_CORRUPT;
			unpacker->handler.packet(unpacker->cbparam, unpacker->ptr, unpacker->size, unpacker->timestamp, unpacker->flags);
			unpacker->flags = 0;
			unpacker->size = 0; // reset
		}
#endif

		assert(unpacker->capacity > 2);
		unpacker->size = 2; // NAL unit type byte
		unpacker->ptr[0] = ptr[0];
		unpacker->ptr[1] = (FU_NAL(fuheader) << 3) | (ptr[1] & 0x07); // replace NAL Unit Type Bits
		assert(H266_TYPE(unpacker->ptr[1]) >= 0 && H266_TYPE(unpacker->ptr[1]) <= 31);
	}
	else
	{
		if (0 == unpacker->size)
		{
			unpacker->flags = RTP_PAYLOAD_FLAG_PACKET_LOST;
			return 0; // packet discard
		}
		assert(unpacker->size > 0);
	}

	unpacker->timestamp = timestamp;
	if (bytes > n)
	{
		assert(unpacker->capacity >= unpacker->size + bytes - n);
		memmove(unpacker->ptr + unpacker->size, ptr + n, bytes - n);
		unpacker->size += bytes - n;
	}

	if (FU_END(fuheader)) // FU_MARK(fuheader) ??
	{
		r = unpacker->handler.packet(unpacker->cbparam, unpacker->ptr, unpacker->size, timestamp, unpacker->flags);
		unpacker->flags = 0;
		unpacker->size = 0;
	}

	return 0 == r ? 1 : r; // packet handled
}

static int rtp_h266_unpack_input(void* p, const void* packet, int bytes)
{
	int r, nal;
	const uint8_t* ptr;
	struct rtp_packet_t pkt;
	struct rtp_decode_h266_t* unpacker;

	unpacker = (struct rtp_decode_h266_t*)p;
	if (!unpacker || 0 != rtp_packet_deserialize(&pkt, packet, bytes) || pkt.payloadlen < (unpacker->using_donl_field ? 5 : 3))
		return -EINVAL;

	if (-1 == unpacker->flags)
	{
		unpacker->flags = 0;
		unpacker->seq = (uint16_t)(pkt.rtp.seq - 1); // disable packet lost
	}

	if ((uint16_t)pkt.rtp.seq != (uint16_t)(unpacker->seq + 1))
	{
		unpacker->flags = RTP_PAYLOAD_FLAG_PACKET_LOST;
		unpacker->size = 0; // discard previous packets
	}
	unpacker->seq = (uint16_t)pkt.rtp.seq;

	assert(pkt.payloadlen > 2);
	ptr = (const uint8_t*)pkt.payload;
	nal = H266_TYPE(ptr[1]);

	if (nal > 31)
		return 0; // packet discard, Unsupported (VVC) NAL type

	switch (nal)
	{
	case H266_RTP_AP: // aggregated packet (AP) - with two or more NAL units
		return rtp_h266_unpack_ap(unpacker, ptr, pkt.payloadlen, pkt.rtp.timestamp);

	case H266_RTP_FU: // fragmentation unit (FU)
		return rtp_h266_unpack_fu(unpacker, ptr, pkt.payloadlen, pkt.rtp.timestamp);

	default: // 4.3.1. Single NAL Unit Packets
		r = unpacker->handler.packet(unpacker->cbparam, ptr, pkt.payloadlen, pkt.rtp.timestamp, unpacker->flags);
		unpacker->flags = 0;
		unpacker->size = 0;
		return 0 == r ? 1 : r; // packet handled
	}
}

struct rtp_payload_decode_t* rtp_h266_decode()
{
	static struct rtp_payload_decode_t unpacker = {
		rtp_h266_unpack_create,
		rtp_h266_unpack_destroy,
		rtp_h266_unpack_input,
	};

	return &unpacker;
}
