#pragma once

#include "h264_sps.h"
#include "h264_pps.h"
#include "h264_slice.h"

//////////////////////////////////////////////////////////////////////////
struct h264_split_sys_t
{
	Buffer		h264_extra;
	Buffer		sei_nal;
	Buffer		sps_nal;
	Buffer		pps_nal;
	h264_sps_t	sps;
	h264_pps_t	pps;
	h264_slice_t slice;

	int		has_sei;
	int		has_sps;
	int		has_pps;

	int			i_nal_type;
	int			i_nal_ref_idc;

	h264_split_sys_t()
	{
		memset(&sps, 0, sizeof(sps));
		memset(&pps, 0, sizeof(pps));
		memset(&slice, 0, sizeof(slice));
		has_sei = FALSE;
		has_sps = FALSE;
		has_pps = FALSE;

		i_nal_type = 0;
		i_nal_ref_idc = 0;
	}
};


enum nal_unit_type_e
{
	NAL_UNKNOWN = 0,
	NAL_SLICE   = 1,
	NAL_SLICE_DPA   = 2,
	NAL_SLICE_DPB   = 3,
	NAL_SLICE_DPC   = 4,
	NAL_SLICE_IDR   = 5,    /* ref_idc != 0 */
	NAL_SEI         = 6,    /* ref_idc == 0 */
	NAL_SPS         = 7,
	NAL_PPS         = 8,
	NAL_AU_DELIMITER= 9
	/* ref_idc == 0 for 6,9,10,11,12 */
};

//////////////////////////////////////////////////////////////////////////

static void h264_decode_annexb( uint8_t *dst, int *dstlen,
							   const uint8_t *src, const int srclen )
{
	uint8_t *dst_sav = dst;
	const uint8_t *end = &src[srclen];

	while (src < end)
	{
		if (src < end - 3 && src[0] == 0x00 && src[1] == 0x00 && src[2] == 0x03)
		{
			*dst++ = 0x00;
			*dst++ = 0x00;

			src += 3;
			continue;
		}
		*dst++ = *src++;
	}

	*dstlen = dst - dst_sav;
}

static void h264_get_nal_type(int* p_nal_type, const uint8_t *p_nal)
{
	int i_nal_hdr;
	int i_offset;

	i_offset = (p_nal[2] == 0) ? 4:3;

	i_nal_hdr = p_nal[i_offset];
	*p_nal_type = i_nal_hdr&0x1f;
}

static void h264_find_frame_end(int* p_found_frame_start, int* p_new_frame, 
								const unsigned char *p_nal, const unsigned int nal_length, int i_nal_type)
{
	if (i_nal_type == 1/*NAL_SLICE*/ || i_nal_type == 2/*NAL_DPA*/ || i_nal_type == 5/*NAL_IDR_SLICE*/)
		*p_found_frame_start = TRUE;

	if (i_nal_type == 7/*NAL_SPS*/ || i_nal_type == 8/*NAL_PPS*/ || i_nal_type == 9/*NAL_AUD*/)
		*p_found_frame_start = FALSE;

	if( (*p_found_frame_start) )
	{
		for (uint32_t i = 3; i < nal_length; i++)
		{
			if (p_nal[i] & 0x80)
			{
				*p_found_frame_start = FALSE;
				*p_new_frame = TRUE;
				return;
			}
		}
	}

	*p_new_frame = FALSE;
}

//////////////////////////////////////////////////////////////////////////

UINT h264_parse_nal(h264_split_sys_t* p_sys, Buffer* p_buff, VBuffer* p_buff_list)
{
	UINT i_split_size = 0;

	/*const*/ uint8_t *p_buffer = p_buff->m_pData;
	int i_buffer = p_buff->m_nDataSize;

	while( i_buffer > 4 && ( p_buffer[0] != 0 || p_buffer[1] != 0 || p_buffer[2] != 1 ) )
	{
		i_buffer--;
		p_buffer++;
	}

	/* Split nal units */
	while( i_buffer > 4 )
	{
		int i_offset;
		int i_size = i_buffer;
		int i_skip = i_buffer;

		/* search nal end */
		for( i_offset = 4; i_offset+2 < i_buffer ; i_offset++)
		{
			if( p_buffer[i_offset] == 0 && p_buffer[i_offset+1] == 0 && p_buffer[i_offset+2] == 1 )
			{
				/* we found another startcode */
				i_size = i_offset - ( p_buffer[i_offset-1] == 0 ? 1 : 0);
				i_skip = i_offset;
				break;
			}
		}

		/* TODO add STAP-A to remove a lot of overhead with small slice/sei/... */

		// 获取一个空缓存
		Buffer* p_nal = p_buff_list->GetEmptyBuffer();
		if (p_nal == NULL)
			p_nal = new Buffer();
		p_nal->FillData((PBYTE)p_buffer, i_size);

		// 获取当前nal type
		int i_nal_type;
		h264_get_nal_type(&i_nal_type, p_nal->m_pData);

		// 找到sei nal
		if (!p_sys->has_sei && i_nal_type == 6)
		{
			p_sys->has_sei = TRUE;
			p_sys->h264_extra.AppendData(p_nal->m_pData, p_nal->m_nDataSize);
			p_sys->sei_nal.FillData(p_nal->m_pData, p_nal->m_nDataSize);
		}

		// 找到sps nal
		if (!p_sys->has_sps && i_nal_type == 7)
		{
			p_sys->has_sps = TRUE;
			p_sys->h264_extra.AppendData(p_nal->m_pData, p_nal->m_nDataSize);
			p_sys->sps_nal.FillData(p_nal->m_pData, p_nal->m_nDataSize);

			Buffer sps_data;
			sps_data.AllocateBuffer(p_nal->m_nDataSize);

			h264_decode_annexb( sps_data.m_pBuffer, (INT *)&sps_data.m_nDataSize, p_nal->m_pData+4, p_nal->m_nDataSize-4);
			h264_decode_seq_parameter_set(sps_data.m_pData, sps_data.m_nDataSize, &p_sys->sps);
			if (!p_sys->sps.num_units_in_tick || !p_sys->sps.time_scale)
			{
				p_sys->sps.time_scale = (p_sys->sps.offset_for_top_to_bottom_field == 0) ? 50: 25;
				p_sys->sps.num_units_in_tick = 1;
			}
		}

		// 找到pps nal
		if ( !p_sys->has_pps && i_nal_type == 8)
		{
			p_sys->has_pps = TRUE;

			Buffer pps_data;
			pps_data.AllocateBuffer(p_nal->m_nDataSize);

			h264_decode_annexb( pps_data.m_pBuffer, (INT *)&pps_data.m_nDataSize, p_nal->m_pData+4, p_nal->m_nDataSize-4);
			h264_decode_picture_parameter_set(&p_sys->pps, pps_data.m_pData, pps_data.m_nDataSize);

			p_sys->h264_extra.AppendData(p_nal->m_pData, p_nal->m_nDataSize);
			p_sys->pps_nal.FillData(p_nal->m_pData, p_nal->m_nDataSize);
		}

		// 将nal加入缓存链表
		p_buff_list->AddFullBuffer(p_nal);

		i_buffer -= i_skip;
		p_buffer += i_skip;

		i_split_size += i_size;
	}

	p_buff->m_pData = p_buffer;
	p_buff->m_nDataSize = i_buffer;

	return i_split_size;
}

int h264_parse_frame_info(h264_split_sys_t* p_sys, 
					uint8_t* p_nal, 
					int32_t i_nal_size,
					int* p_found_frame_start,
					int64_t* p_frame_count, 
					int64_t* p_group_offset, 
					int64_t* p_pts,
					int64_t* p_dts,
					int* p_flags)
{
	int i_nal_type = 0;
	int b_new_frame = FALSE;

	do 
	{
		if (i_nal_size < 3)
			break;

		h264_get_nal_type(&i_nal_type, p_nal);

		// 检查是否是slice
		if ( i_nal_type < 1/*NAL_SLICE*/ || i_nal_type > 5/*NAL_SLICE_IDR*/ )
			break;

		// 判断是否帧结束
		h264_find_frame_end(p_found_frame_start, &b_new_frame, p_nal, i_nal_size, i_nal_type);
		//const int i_nal_ref_idc = (p_nal[3] >> 5)&0x03;

		// 解析slice类型
		Buffer slice_data;
		slice_data.AllocateBuffer(i_nal_size-4);

		h264_decode_annexb(slice_data.m_pData, 
			(int*)&slice_data.m_nDataSize, 
			p_nal+4, 
			__min(i_nal_size-4, 60));

		h264_slice_t slice;
		h264_decode_slice(&slice, 
			slice_data.m_pData, slice_data.m_nDataSize, i_nal_type, &p_sys->sps, &p_sys->pps);

		// 根据slice类型判断帧类型
		switch(slice.i_slice_type)
		{
		case 2:	case 7:	
		case 4:	case 9:
			*p_flags = 0x0002/*BLOCK_FLAG_TYPE_I*/;
			break;
		case 0:	case 5:
		case 3:	case 8:
			*p_flags = 0x0004/*BLOCK_FLAG_TYPE_P*/;
			break;
		case 1:
		case 6:
			*p_flags = 0x0008/*BLOCK_FLAG_TYPE_B*/;
			break;
		default:
			*p_flags = 0;
			break;
		}

		//if( slice.i_frame_num != p_sys->slice.i_frame_num ||
		//	slice.i_pic_parameter_set_id != p_sys->slice.i_pic_parameter_set_id ||
		//	slice.i_field_pic_flag != p_sys->slice.i_field_pic_flag ||
		//	i_nal_ref_idc != p_sys->i_nal_ref_idc )
		//	b_new_frame = true;

		//if( (slice.i_bottom_field_flag != -1) &&
		//	(p_sys->slice.i_bottom_field_flag != -1) &&
		//	(slice.i_bottom_field_flag != p_sys->slice.i_bottom_field_flag) )
		//	b_new_frame = true;

		//if( p_sys->sps.i_pic_order_cnt_type == 0 &&
		//	( slice.i_pic_order_cnt_lsb != p_sys->slice.i_pic_order_cnt_lsb ||
		//	slice.i_delta_pic_order_cnt_bottom != p_sys->slice.i_delta_pic_order_cnt_bottom ) )
		//	b_new_frame = true;
		//else if( p_sys->sps.i_pic_order_cnt_type == 1 &&
		//	( slice.i_delta_pic_order_cnt0 != p_sys->slice.i_delta_pic_order_cnt0 ||
		//	slice.i_delta_pic_order_cnt1 != p_sys->slice.i_delta_pic_order_cnt1 ) )
		//	b_new_frame = true;

		//if( ( i_nal_type == NAL_SLICE_IDR || p_sys->i_nal_type == NAL_SLICE_IDR ) &&
		//	( i_nal_type != p_sys->i_nal_type || slice.i_idr_pic_id != p_sys->slice.i_idr_pic_id ) )
		//	b_new_frame = true;

		//p_sys->i_nal_ref_idc = i_nal_ref_idc;
		//p_sys->i_nal_type = i_nal_type;
		//p_sys->slice = slice;

		if (!b_new_frame)
			break;
		
		(*p_frame_count)++;

		// 根据计算slice中的i_pic_order_cnt_lsb、i_frame_num信息，以及I帧位置计算pts和dts
		if (p_sys->sps.i_pic_order_cnt_type == 0)
		{
			if (slice.i_pic_order_cnt_lsb == 0)
			{
				*p_group_offset = (*p_frame_count);
				*p_pts = *p_dts = (*p_frame_count) * ((p_sys->sps.offset_for_top_to_bottom_field == 0)?2:1) * 1000LL * p_sys->sps.num_units_in_tick / p_sys->sps.time_scale;
			}
			else
			{
				*p_pts = ( slice.i_pic_order_cnt_lsb + (*p_group_offset)*2)* 1000LL * p_sys->sps.num_units_in_tick / p_sys->sps.time_scale;
				*p_dts = (*p_frame_count) * ((p_sys->sps.offset_for_top_to_bottom_field == 0)?2:1) * 1000LL * p_sys->sps.num_units_in_tick / p_sys->sps.time_scale;
			}
		}
		else
		{
			*p_pts = *p_dts = (*p_frame_count) * ((p_sys->sps.offset_for_top_to_bottom_field == 0)?2:1) * 1000LL * p_sys->sps.num_units_in_tick / p_sys->sps.time_scale;
		}

		//printf("order(%d) IPB(%s) pts(%I64d) frame count(%I64d) I offset(%I64d) icount(%I64d)\n", slice.i_pic_order_cnt_lsb, sssss, *p_pts, *p_frame_count, *p_group_offset, test_count);

	} while (FALSE);

	return b_new_frame;
}