/*******************************************************************
                        PES treating module
 *******************************************************************/

#include <string.h>

#define PES_C
#include "pes.h"

typedef struct {
	int pes_scrambling_control;
	int pes_priority;
	int data_alignment_indicator;
	int copyright;
	int original_or_copy;
	int pts_dts_flag;
	int escr_flag;
	int es_rate_flag;
	int dsm_trick_mode_flag;
	int additional_copy_info_flag;
	int pes_crc_flag;
	int pes_extention_flag;
	int pes_header_data_length;
}STANDARD_PES_HEADER;

int read_pes_packet(BITSTREAM *in, PES_PACKET *out);
unsigned int get_pes_packet_data_length(PES_PACKET *p);
int extract_pes_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type);
int extract_pes_packet_data(PES_PACKET *p, unsigned char *data, unsigned int *length);

static unsigned int private_1_stream_data_length(PES_PACKET *p);
static unsigned int audio_stream_data_length(PES_PACKET *p);
static unsigned int video_stream_data_length(PES_PACKET *p);

static int private_1_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type);
static int audio_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type);
static int video_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type);

static int extract_standard_pes_header(PES_PACKET *p, STANDARD_PES_HEADER *sph);

int read_pes_packet(BITSTREAM *in, PES_PACKET *out)
{
	int code, n, m;

	do{
		if(! bs_next_packet_prefix(in)){
			return 0;
		}

		bs_erase_bits(in, 24);
		code = bs_get_bits(in, 8);
		
	}while(code < 0xbb);

	out->stream_id = code;

	n = bs_get_bits(in, 16);
	out->data_length = 0;
	while(n){
		m = bs_read(in, out->data+out->data_length, n);
		n -= m;
		out->data_length += m;
		if(m == 0){
			return 0;
		}
	}

	return out->data_length;
}

unsigned int get_pes_packet_data_length(PES_PACKET *p)
{
	if(p->stream_id == 0xbd){
		return private_1_stream_data_length(p);
	}else if(p->stream_id >= 0xc0 && p->stream_id <= 0xdf){
		return audio_stream_data_length(p);
	}else if(p->stream_id >= 0xe0 && p->stream_id <= 0xef){
		return video_stream_data_length(p);
	}else{
		return 0;
	}
}

int extract_pes_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type)
{
	if(p->stream_id == 0xbd){
		return private_1_stream_type(p, type);
	}else if(p->stream_id >= 0xc0 && p->stream_id <= 0xdf){
		return audio_stream_type(p, type);
	}else if(p->stream_id >= 0xe0 && p->stream_id <= 0xef){
		return video_stream_type(p, type);
	}else{
		return 0;
	}
}

int extract_pes_packet_data(PES_PACKET *p, unsigned char *data, unsigned int *length)
{
	if((data == NULL) || (length == NULL) ){
		return 0;
	}

	*length = get_pes_packet_data_length(p);

	memcpy(data, p->data+(p->data_length-(*length)), *length);

	return *length;
}

static unsigned int private_1_stream_data_length(PES_PACKET *p)
{
	unsigned int r;
	unsigned char *w;
	STANDARD_PES_HEADER sph;

	extract_standard_pes_header(p, &sph);
	w = p->data + sph.pes_header_data_length;

	r = 0;
	if( (w[0] >= 0x80) && (w[0] <= 0x8f) ){
		/* AC3 stream */
		r = p->data_length - sph.pes_header_data_length - 4;
	}else{
		/* unknown stream */
	}

	return r;
}

static unsigned int audio_stream_data_length(PES_PACKET *p)
{
	STANDARD_PES_HEADER sph;

	extract_standard_pes_header(p, &sph);

	return p->data_length - sph.pes_header_data_length;
}

static unsigned int video_stream_data_length(PES_PACKET *p)
{
	STANDARD_PES_HEADER sph;

	extract_standard_pes_header(p, &sph);

	return p->data_length - sph.pes_header_data_length;
}

static int private_1_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type)
{
	unsigned char *w;
	STANDARD_PES_HEADER sph;

	extract_standard_pes_header(p, &sph);
	w = p->data + sph.pes_header_data_length;

	if( (w[0] >= 0x80) && (w[0] <= 0x8f) ){
		type->type = PES_STREAM_TYPE_AC3;
		type->id = w[0] & 0x0f;
		return 1;
	}else{
		type->type = PES_STREAM_TYPE_UNKNOWN;
		type->id = 0;
		return 0;
	}
}

static int audio_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type)
{
	type->type = PES_STREAM_TYPE_AUDIO;
	type->id = p->stream_id & 0x1f;

	return 1;
}

static int video_stream_type(PES_PACKET *p, PES_STREAM_TYPE *type)
{
	type->type = PES_STREAM_TYPE_VIDEO;
	type->id = p->stream_id & 0x0f;
	
	return 1;
}

static int extract_standard_pes_header(PES_PACKET *p, STANDARD_PES_HEADER *sph)
{
	if( (p == NULL) || (sph == NULL) ){
		return 0;
	}
	
	sph->pes_scrambling_control = (p->data[0] & 0x30) >> 4;
	sph->pes_priority = (p->data[0] & 0x08) >> 3;
	sph->data_alignment_indicator = (p->data[0] & 0x04) >> 2;
	sph->copyright = (p->data[0] & 0x02) >> 1;
	sph->original_or_copy = (p->data[0] & 0x01);
	sph->pts_dts_flag = (p->data[1] & 0xC0) >> 6;
	sph->escr_flag = (p->data[1] & 0x20) >> 5;
	sph->es_rate_flag = (p->data[1] & 0x10) >> 4;
	sph->dsm_trick_mode_flag = (p->data[1] & 0x08) >> 3;
	sph->additional_copy_info_flag = (p->data[1] & 0x04) >> 2;
	sph->pes_crc_flag = (p->data[1] & 0x02) >> 1;
	sph->pes_extention_flag = (p->data[1] & 0x01);
	sph->pes_header_data_length = p->data[2] + 3;

	return 1;
}