#include <arpa/inet.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include "av_format_identify.h"


#define FIT_TIMES 2
#define TS_SYNC_BYTE 0x47
#define TS_PKT_SIZE 188
#define FLV_TAG_AUDIO 0x08
#define FLV_TAG_VIDEO 0x09
#define FLV_TAG_SCRIPT_DATA 0x12

//g_mp4_box_type_all: various box type in mp4
const char *g_mp4_box_type_all[] = { "ftyp","pdin","moov","mvhd","meta",
"trak","tkhd","tref","trgr","edts","elst","mdia","mdhd","hdlr","elng",
"minf","vmhd","smhd","hmhd","sthd","nmhd","dinf","dref","stbl","stsd",
"stts","ctts","cslg","stsc","stsz","stz2","stco","co64","stss","stsh",
"padb","stdp","sdtp","sbgp","sgpd","subs","saiz","saio","udta","mvex",
"mehd","trex","leva","moof","mfhd","traf","tfhd","trun","tfdt","mfra",
"tfra","mfro","mdat","free","skip","cprt","tsel","strk","stri","strd",
"iloc","ipro","sinf","frma","schm","schi","iinf","xml","bxml","pitm",
"fiin","paen","fire","fpar","fecr","segr","gitn","idat","iref","meco",
"mere","styp","sidx","ssix","prft", "uuid" };
//g_mp4_box_type_count:how many kinds of box type in mp4
const int g_mp4_box_type_count = sizeof(g_mp4_box_type_all) / sizeof(char *);

#pragma pack (1)
typedef struct ts_pkg_t
{
	uint8_t sync_byte;//sync_byte
	uint16_t con_pid;//3 bits+pid(13 bits)
	uint8_t con_ctt_cnt;//4 bits+continuity_counter(4 bits)
	char *pkt_data;//packet data
}ts_pkg;
#pragma pack ()
#define MIN_TS_PKT_LEN	(sizeof(ts_pkg) - 8)

#define MIN_TS_PKT_CNTS 4

/*To find that whether the initial size bytes pointed to by buff satisfies the rules of MPEG-TS AV Container
* by supposing that its first 1 byte is sync_byte in ts header.*/
int is_mpeg_ts_each(const char *buff, size_t size)
{
	int ret = 0;
	int fit_times = 0;
	const char *p = buff;
	const ts_pkg *pkg = NULL;

	uint8_t ctt_cnt_pre = 0, ctt_cnt_cur = 0;
	uint16_t pid_pre = 0, pid_cur = 0;

	//record the first PID and continuous_count
	if (p + MIN_TS_PKT_LEN - buff > size)//which is: while (p + (MIN_TS_PKT_LEN - 1) - buff >= size)
	{
		return ret;
	}
	pkg = (const ts_pkg *)p;
	pid_pre = htons(pkg->con_pid) & 0x1fff;
	ctt_cnt_pre = (pkg->con_ctt_cnt) & 0x0f;
	p += TS_PKT_SIZE;
	fit_times++;

	int pkt_cnt_divs[3] = { 1,2,4 };
	while (p + MIN_TS_PKT_LEN - buff <= size)//which is: while (p + (MIN_TS_PKT_LEN - 1) - buff < size)
	{
		pkg = (const ts_pkg *)p;
		if (pkg->sync_byte != TS_SYNC_BYTE)
		{
			break;
		}
		pid_cur = htons(pkg->con_pid) & 0x1fff;
		if (pid_cur == pid_pre && pid_cur != 8191)//no continuous_count rule for PID:8191
		{
			ctt_cnt_cur = (pkg->con_ctt_cnt) & 0x0f;
			if (ctt_cnt_cur != (ctt_cnt_pre + 1) % 16)
			{
				break;
			}
		}
		fit_times++;

		if (fit_times >= FIT_TIMES)
		{
			size_t whole_pkg_cnt = (size - 1) / TS_PKT_SIZE;
			if (whole_pkg_cnt < MIN_TS_PKT_CNTS)
			{
				break;
			}
			int i;
			for (i = 0; i < 3; i++)
			{
				const char *pkg_seg = buff + whole_pkg_cnt / pkt_cnt_divs[i] * TS_PKT_SIZE;
				if (*pkg_seg != TS_SYNC_BYTE)
				{
					return ret;
				}
			}
			ret = 1;
			break;
		}

		ctt_cnt_pre = ctt_cnt_cur;
		pid_pre = pid_cur;
        p += TS_PKT_SIZE;
	}
	return ret;
}


/*To identify whether the whole bytestream belongs to MPEG_TS container*/
int is_mpeg_ts(const char* buff, size_t size)
{
	const char *p = buff;
	int is_ts=0;
	while (p - buff < size)
	{
		p = memchr(p, TS_SYNC_BYTE, size - (p - buff));
		if (p == NULL)
		{
			break;
		}
		is_ts = is_mpeg_ts_each(p, size - (p - buff));
		if (is_ts == 1)
		{
            break;
		}
		p++;		
	}
	return is_ts;
}


void *memchr_flv(const char *buff, size_t size)
{
	void *ret_memflv = NULL;
	const char *p = buff;
	while (p - buff < size)
	{
		if (*p == FLV_TAG_AUDIO || *p == FLV_TAG_VIDEO || *p == FLV_TAG_SCRIPT_DATA)
		{
			ret_memflv = (void *)p;
			break;
		}
		p++;
	}
	return ret_memflv;
}


#pragma pack (1)
typedef struct flv_tag_t
{
	uint32_t pre_tag_size;//Previous Tag Size
	uint8_t tag_type;//Tag Type, in Tag Header
	uint32_t con_data_size;//Data Size(3 Bytes) + 1 Byte, in Tag Header
}flv_tag;
#pragma pack ()
#define MIN_FLV_TAG_LEN	(sizeof(flv_tag))

/*To find that whether the initial size bytes pointed to by buff satisfies the rules of FLV AV Container
* by supposing that its first 1 byte is TagType in FLV Tag Header.*/
int is_flv_each(const char* buff, size_t size)
{
	int ret = 0;
	int fit_times = 0;
	const char *p = buff;
	const flv_tag *tag = NULL;
	
	uint32_t tag_len_cnt = 0;
	uint32_t tag_len_show = 0;

	//record the first Tag Size and Data Size without Previous Tag Size.
	if (p + MIN_FLV_TAG_LEN - 4 - buff > size)//which is (p + (MIN_FLV_TAG_LEN - 4 - 1) - buff >= size), p is located in tag_type;
	{
		return ret;
	}
    tag_len_cnt = (htonl(*(uint32_t *)p) & 0x00ffffff) + 11;//?? 11 + (htonl(tag->con_data_size) >> 8)
	p += tag_len_cnt;//p is the first previous_length after buff
	fit_times++;

	while (p + MIN_FLV_TAG_LEN - buff <= size)//which is: while (p + (MIN_FLV_TAG_LEN - 1) - buff < size), during this while loop, p is located in pre_tag_size
	{
		tag = (const flv_tag *)p;
		tag_len_show = htonl(tag->pre_tag_size);

		if ((tag->tag_type) != FLV_TAG_AUDIO && (tag->tag_type) != FLV_TAG_VIDEO && (tag->tag_type) != FLV_TAG_SCRIPT_DATA)
		{
			break;
		}
		if (tag_len_cnt != tag_len_show)
		{
			break;
		}

		uint32_t tag_len_cnt = (htonl(tag->con_data_size) >> 8) + 11;

		fit_times++;
		if (fit_times >= FIT_TIMES)
		{
			ret = 1;
			break;
		}
		p += (4 + tag_len_cnt);
	}
	return ret;
}


/*To identify whether the whole bytestream belongs to FLV container*/
int is_flv(const char* buff, size_t size)
{
	int is_flv = 0;
	const char *p = buff;

	while (p - buff < size)
	{
		p = memchr_flv(p, size - (p - buff));
		if (p == NULL)
		{
			break;
		}
        is_flv = is_flv_each(p, size - (p - buff));
		if (is_flv == 1)
		{
			break;
		}
		p++;
	}
	return is_flv;
}


int is_mp4_boxname_deep(const char *buff)
{
	int i;
	size_t size_type_h;
	for (i = 0; i < g_mp4_box_type_count; i++)
	{
		size_type_h = strlen(g_mp4_box_type_all[i]);
		if (memcmp(buff, g_mp4_box_type_all[i], size_type_h) == 0)
			return 1;
	}
	return 0;
}


int is_mp4_boxname(const char *buff)
{
	int i = 0, raw_filter = 0;

	for (i = 0; i < 4; i++)
	{
		if (((*(buff + i) >= '0') && (*(buff + i) <= '9')) || ((*(buff + i) >= 'a') && (*(buff + i) <= 'z')))
		{
			
			raw_filter++;
		}
		else
		{
			break;
		}
	}
	if (raw_filter == 4)//box type:xmf
	{
		return 1;
	}
	else if (raw_filter == 3)
	{
		if (memcmp(buff, "xmf", 3) == 0)
		{
			return 1;
		}
	}
	return 0;
}


void *memmem_mp4(const char *buff, size_t size)
{
	void *ret_memmp4 = NULL;
	const char *p = buff;
	while (p + 3 - buff < size)
	{
        if ((is_mp4_boxname(p) == 1) && (is_mp4_boxname_deep(p) == 1))
		{
			ret_memmp4 = (void *)p;
			break;
		}
		p++;
	}
	return ret_memmp4;
}


#pragma pack (1)
typedef struct mp4_box_t
{
	uint32_t atom_size;//Box Size
	char name[4];//Box Type, in BoxHeader
	union
	{
		uint64_t large_size;//when the Box Size is larger than 0xffffffff, using Largesize to save the size of atom
		char* box_data;
	};
}mp4_box;
#pragma pack ()
#define MIN_MP4_BOX_SIZE	(sizeof(mp4_box)-8)
#define htonll(buff)  (((uint64_t) htonl(buff)) << 32) + htonl(buff >> 32)//(((uint64_t)(htonl(*(uint32_t *)buff))<<32)+htonl(*(uint32_t *)(buff+4)))

/*To find that whether the initial size bytes pointed to by buff satisfies the rules of MP4 AV Container
* by supposing that its first 4 byte is Boxtype in Box of MP4 file.*/
int is_mp4_each(const char* buff, size_t size)
{
	int ret = 0;
	int fit_times = 0;
	const char *p = buff;
	const mp4_box *box=NULL;

	while (p + MIN_MP4_BOX_SIZE - buff <= size)//which is while (p + (MIN_MP4_BOX_SIZE - 1) - buff <= size), p is located in atom_size
	{
		box = (const mp4_box *)p;
		if(!is_mp4_boxname(box->name))
		{
			break;
		}		
		fit_times++;
		if (fit_times >= FIT_TIMES)
		{
			ret = 1;
			break;
		}//


		if(htonl(box->atom_size)==1)
		{
			if (p + sizeof(mp4_box) - buff < size)//which is (p + (sizeof(mp4_box) - 1) - buff < size)
			{
				p += htonll(box->large_size);
			}
			else
			{
				break;
			}
		}
		else
		{
			p += htonl(box->atom_size);
		}
	}

	return ret;
}


/*To identify whether the whole bytestream belongs to MP4 container*/
int is_mp4(const char* buff, size_t size)
{
	int is_mp4 = 0;
	const char *p = buff;

	while (p - buff< size)
	{
		p = (const char *)memmem_mp4(p, size - (p - buff));
		if (p == NULL)
		{
			break;
		}
		if (p - 4 - buff < 0)//if there is no atom_size for the first box existing, forcing the next iteration of the loop
		{
			p++;
			continue;
		}

		is_mp4 = is_mp4_each(p - 4, size - (p - buff) + 4);
		if (is_mp4 == 1)
		{
			break;
		}
		p++;
	}
	return is_mp4;
}


int AV_container_identify(const char* buff, size_t size)
{
	if (1 == is_mpeg_ts(buff, size))
	{
		return AV_CONTAINER_MPEGTS;
	}
	if (1 == is_flv(buff, size))
	{
		return AV_CONTAINER_FLV;
	}
	if (1 == is_mp4(buff, size))
	{
		return AV_CONTAINER_MP4;
	}
	return AV_CONTAINER_UNKNOWN;
}