#include "av_format_identify.h"
#include <arpa/inet.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <stdint.h>
#include <math.h>
#include <sys/types.h>  

#include <sys/io.h>
#include <arpa/inet.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

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" };
const int g_mp4_box_type_count = sizeof(g_mp4_box_type_all) / sizeof(char *);

/*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.*/

/*To identify whether the whole byte stream belongs to MPEG_TS container*/



/*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.*/

/*To identify whether the whole byte stream belongs to FLV container*/

/*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.*/

 /*To identify whether the whole byte stream belongs to MP4 container*/
#pragma pack (1)
typedef struct ts_pkg_t
{
	uint8_t sync_byte;
	uint16_t con_pid;
	uint8_t con_cnt;
	char *pyld_data;
}ts_pkg;
#pragma pack ()
#define MIN_TS_PKT_LEN	(sizeof(ts_pkg)-8)

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;
	
	//uint16_t  cont_cnt_pre=0, cont_cnt_cur=0;
	//cont_cnt_pre = (*(uint16_t *)(pkg->con_cnt)) & 0x0f;
	uint8_t cont_cnt_pre = 0, cont_cnt_cur = 0;
	//cont_cnt_pre = pkg->con_cnt;

	uint16_t pid_pre=0, pid_cur=0;
	//pid_pre = (htons(*(uint16_t *)(pkg->con_pid))) << 3 >> 3;

		
	int group[3] = { 1,2,4 };

	//p += TS_PKT_SIZE;
	/*when p == buff(first of while), record pid and continuous_count of package head and go to fit_times++ directly.
	special for mpeg-ts cause of the format requirement for two packages together.*/
	if (p - buff + MIN_TS_PKT_LEN > size)
	{
		return ret;
	}
	
	pkg = (const ts_pkg *)p;
	pid_pre = htons(pkg->con_pid) & 0x1fff;//(htons(pkg->con_pid)) << 3 >> 3;//?
	cont_cnt_pre = (pkg->con_cnt) & 0x0f;//(pkg->con_cnt) >> 4;
	p += TS_PKT_SIZE;
	fit_times++;

	while (p - buff + MIN_TS_PKT_LEN <= size)//while (p < (buff + size))
	{
		pkg = (const ts_pkg *)p;
		if (pkg->sync_byte != TS_SYNC_BYTE)
		{
			break;
		}

		pid_cur = htons(pkg->con_pid) & 0x1fff;//(htons(pkg->con_pid)) << 3 >> 3;
		//pid_cur = ((htonl(*(uint32_t *)(p+TS_PKT_SIZE)) & 0x000fff00) >> 8) + ((((htonl(*(uint32_t *)(p + TS_PKT_SIZE)) & 0x00f00000) >> 20) % 2) << 12);
		if (pid_cur == pid_pre && pid_cur != 8191)//
		{
			cont_cnt_cur = (pkg->con_cnt) & 0x0f;//(pkg->con_cnt) >> 4;// p + TS_PKT_SIZE;//cont_cnt_cur = (*(uint16_t *)(p + TS_PKT_SIZE)) & 0x0f;
			if (cont_cnt_cur != (cont_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 < 4)
			{
				break;
			}
			int i;
			for (i = 0; i < 3; i++)
			{
				const char *pkg_seg = buff + whole_pkg_cnt / group[i] * TS_PKT_SIZE;
				if (*pkg_seg != TS_SYNC_BYTE)
				{
					return ret;//break;
				}
			}
			ret = 1;
			break;
		}

		cont_cnt_pre = cont_cnt_cur;
		pid_pre = pid_cur;
        p += TS_PKT_SIZE;
	}
	return ret;
}

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)
		{
			printf("p-buff is %lx.\n",p-buff);
            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_len;
	uint8_t tag_type;
	uint32_t con_pyld_len;//uint8_t pyld_len[3];//payload length
}flv_tag;
#pragma pack ()
#define MIN_FLV_TAG_LEN	(sizeof(flv_tag))//(sizeof(flv_tag) + 1)//?

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_show = 0;
	//uint32_t tag_len_cnt = 0;// = 11 + (htonl(*(size_t *)p) & 0x00ffffff);
	//p += (tag_len_cnt + 4);
	uint32_t tag_len_cnt = 0;
	uint32_t tag_len_show = 0;
	if (p - buff + MIN_FLV_TAG_LEN - 4 > size)
	{
		return ret;
	}
	
    tag_len_cnt = (htonl(*(uint32_t *)p) & 0x00ffffff) + 11;//?? 11 + (htonl(tag->con_pyld_len) >> 8);//11 + (htonl(*(uint32_t *)tag->pyld_len) >> 8);
	p += tag_len_cnt;//p is the first previous_length after buff
	fit_times++;

	while (p - buff + MIN_FLV_TAG_LEN <= size)//while (p < (buff + size))
	{
		tag = (const flv_tag *)p;
		tag_len_show = htonl(tag->pre_len);

		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_pyld_len) >> 8) + 11;//11 + (htonl(*(uint32_t *)(tag->pyld_len)) >> 8);//(htonl(*(uint32_t *)tag)) & 0x00ffffff;

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

int is_flv(const char* buff, size_t size)
{
	int is_flv = 0;
	const char *p = buff;

	while (p - buff < size)//same as p - buff + 1 <= size//while (p < size + buff)
	{
		p = memchr_flv(p, size - (p - buff));
		if (p == NULL)
		{
			break;
		}
        is_flv = is_flv_each(p, size - (p - buff));
		if (is_flv == 1)
		{
			printf("p-buff is %lx.\n",p-buff);
			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)// , size_t size)
{
	int i=0,raw_filter=0;
	/*
	if (size<4)
	{
		return 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)
        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;
	char name[4];
	union
	{
		uint64_t large_size;
		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)))
int is_mp4_each(const char* buff, size_t size)
{
	int ret = 0;
	int fit_times = 0;
	const char *p = buff;
	const char *q = NULL;
	const mp4_box *box=NULL;
	uint64_t offset = 0;

	while (p - buff + MIN_MP4_BOX_SIZE <= size)//while(p-buff+MIN_MP4_BOX_SIZE<size)//at least contain atom_size and name of mp4 box for current p
	{
		box = (const mp4_box *)p;//?
		if(!is_mp4_boxname(box->name))//?,sizeof(box->name))
		{
			if (offset <= 0xffffffff)
			{
				p = q + 8;
				box = (const mp4_box *)p;
			}
			if (!is_mp4_boxname(box->name))
			{
				break;
			}
			//break;
		}		
		fit_times++;
		if (fit_times >= FIT_TIMES)
		{
			ret = 1;
			break;
		}//


		if(htonl(box->atom_size)==1)//?if(box->atom_size == 1)
		{
			if(size-(p-buff)>sizeof(mp4_box))//???
			{
				offset = htonll(box->large_size);//htonll(&(box->large_size));//p += htonll(&box->large_size);//p+=htonll(box->large_size);//TODO: Implement htonll
			}
			else
			{
				break;
			}
		}
		else
		{
			offset = htonl(box->atom_size);//p+=htonl(box->atom_size);
		}
		q = p;
		p += offset;

	}

	/*else
	{
		ret = 0;
	}//
	*/
	return ret;
}


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)//if there is no atom_size for the first box existing, the exit
		{
			p++;
			continue;
		}

		is_mp4 = is_mp4_each(p - 4, size - (p - buff) + 4);
		if (is_mp4 == 1)
		{
			printf("p-buff is %lx.\n",p-buff);
			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;
}