#include <string.h>
#include <assert.h>
#include <stdint.h>
#include <stddef.h>

#define RTP_H2645_BITSTREAM_FORMAT_DETECT 1

static const uint8_t* h264_startcode(const uint8_t* data, int bytes)
{
	int i;
	for (i = 2; i + 1 < bytes; i++)
	{
		if (0x01 == data[i] && 0x00 == data[i - 1] && 0x00 == data[i - 2])
			return data + i + 1;
	}

	return NULL;
}

/// @return >0-ok, <=0-error
static inline int h264_avcc_length(const uint8_t* h264, int bytes, int avcc)
{
	int i;
	uint32_t n;

	n = 0;
	assert(3 <= avcc && avcc <= 4);
	for (i = 0; i < avcc && i < bytes; i++)
		n = (n << 8) | h264[i];
	return avcc >= bytes ? -1 : (int)n;
}

/// @return 1-true, 0-false
static int h264_avcc_bitstream_valid(const uint8_t* h264, int bytes, int avcc)
{
	int n;

	while (avcc + 1 < bytes)
	{
		n = h264_avcc_length(h264, bytes, avcc);
		if (n < 0 || n + avcc > bytes)
			return 0; // invalid

		h264 += n + avcc;
		bytes -= n + avcc;
	}

	return 0 == bytes ? 1 : 0;
}

/// @return 0-annexb, >0-avcc, <0-error
static int h264_bitstream_format(const uint8_t* h264, int bytes)
{
	uint32_t n;
	if (bytes < 4)
		return -1;

	n = ((uint32_t)h264[0]) << 16 | ((uint32_t)h264[1]) << 8 | ((uint32_t)h264[2]);
	if (0 == n && h264[3] <= 1)
	{
		return 0; // annexb
	}
	else if (1 == n)
	{
		// try avcc & annexb
		return h264_avcc_bitstream_valid(h264, bytes, 4) ? 4 : 0;
	}
	else
	{
		// try avcc 4/3 bytes
		return h264_avcc_bitstream_valid(h264, bytes, 4) ? 4 : (h264_avcc_bitstream_valid(h264, bytes, 3) ? 3 : -1);
	}
}

static int h264_avcc_nalu(const void* h264, int bytes, int avcc, int (*handler)(void* param, const uint8_t* nalu, int bytes, int last), void* param)
{
	int r;
	uint32_t n;
	const uint8_t* p, * end;

	r = 0;
	p = (const uint8_t*)h264;
	end = (const uint8_t*)h264 + bytes;
	for (n = h264_avcc_length(p, (int)(end - p), avcc); 0 == r && p + n + avcc <= end; n = h264_avcc_length(p, (int)(end - p), avcc))
	{
		assert(n > 0);
		if (n > 0)
		{
			r = handler(param, p + avcc, (int)n, p + avcc + n < end ? 0 : 1);
		}

		p += n + avcc;
	}

	return r;
}

///@param[in] h264 H.264 byte stream format data(A set of NAL units)
int rtp_h264_annexb_nalu(const void* h264, int bytes, int (*handler)(void* param, const uint8_t* nalu, int bytes, int last), void* param)
{
	int r;
	ptrdiff_t n;
	const uint8_t* p, * next, * end;

#if defined(RTP_H2645_BITSTREAM_FORMAT_DETECT)
	int avcc;
	avcc = h264_bitstream_format(h264, bytes);
	if (avcc > 0)
		return h264_avcc_nalu(h264, bytes, avcc, handler, param);
#endif

	end = (const uint8_t*)h264 + bytes;
	p = h264_startcode((const uint8_t*)h264, bytes);

	r = 0;
	while (p && 0 == r)
	{
		next = h264_startcode(p, (int)(end - p));
		if (next)
		{
			n = next - p - 3;
		}
		else
		{
			n = end - p;
		}

		while (n > 0 && 0 == p[n - 1]) n--; // filter tailing zero

		assert(n > 0);
		if (n > 0)
		{
			r = handler(param, p, (int)n, next ? 0 : 1);
		}

		p = next;
	}

	return r;
}