/*******************************************************************
                    MPEG-2 VIDEO module
 *******************************************************************/
#include <stdio.h>

#include "gop_list.h"
#include "filename.h"
#include "registry.h"

#include "idct_int32.h"
#include "idct_mmx32.h"
#include "idct_double.h"
#include "idct_sse32.h"

#define MPEG2VIDEO_C
#include "mpeg2video.h"

/* grobal */
int open_mpeg2video(char *path, MPEG2VIDEO *out);
int close_mpeg2video(MPEG2VIDEO *p);
OUT_BUFFER_ELEMENT *read_frame(MPEG2VIDEO *in, __int64 frame);

/* local */
static void proc_sequence_header(MPEG2VIDEO *in);
static OUT_BUFFER_ELEMENT *proc_gop_header(MPEG2VIDEO *in, __int64 frame);
static OUT_BUFFER_ELEMENT *proc_picture_data(MPEG2VIDEO *in, __int64 frame);

static OUT_BUFFER_ELEMENT *store_forward_reference_frame(MPEG2VIDEO *in, __int64 frame);
static OUT_BUFFER_ELEMENT *store_backward_reference_frame(MPEG2VIDEO *in, __int64 frame);
static OUT_BUFFER_ELEMENT *store_current_decoding_frame(MPEG2VIDEO *in, __int64 frame);
static OUT_BUFFER_ELEMENT *rotate_reference_frame(MPEG2VIDEO *in, __int64 frame);

static int is_frame_in_current_gop(MPEG2VIDEO *p, __int64 frame);

static void sequence_header_to_decode_picture_parameter(SEQUENCE_HEADER *in, DECODE_PICTURE_PARAMETER *out);
static void picture_header_to_decode_picture_parameter(PICTURE_HEADER *in, DECODE_PICTURE_PARAMETER *out);

static void decode_2nd_field(MPEG2VIDEO *p);

static OUT_BUFFER_ELEMENT *add_frame_out_buffer_with_resize(MPEG2VIDEO *p, FRAME *data, OUTPUT_PARAMETER *prm);

static int is_registered_suffix(char *filepath);

static void clear_output_parameters(MPEG2VIDEO *p);

static void setup_m2v_config(M2V_CONFIG *p);

static void setup_chroma_upsampling_function(MPEG2VIDEO *p, int chroma_format, int simd);
static void setup_idct_function(DECODE_PICTURE_PARAMETER *p, M2V_CONFIG *prm);
static void setup_field_order(MPEG2VIDEO *p, M2V_CONFIG *prm);
static void resize_parameter_to_bgr_conversion_parameter(RESIZE_PARAMETER *in, BGR_CONVERSION_PARAMETER *out);

/******************************************************************************/
int open_mpeg2video(char *path, MPEG2VIDEO *out)
{
	int code;
	
	GOP g;
	READ_GOP_PARAMETER *gp;
	GOP_LIST *gl;

	char gl_path[FILENAME_MAX]; /* GOP LIST file path */

	/* initialize */
	memset(out, 0, sizeof(MPEG2VIDEO));

	if(!is_registered_suffix(path)){
		return 0;
	}

	if(!open_video_stream(path, &(out->bitstream))){
		/* can't open */
		return 0;
	}

	if(!vs_next_start_code(&(out->bitstream))){
		/* not MPEG-2 VIDEO file */
		close_video_stream(&(out->bitstream));
		return 0;
	}

	code = vs_get_bits(&(out->bitstream), 32);
	if(code != 0x1B3){
		/* not VIDEO stream file (Program Stream or Transport Stream) */
		close_video_stream(&(out->bitstream));
		return 0;
	}

	if(!read_sequence_header(&(out->bitstream), &(out->seq))){
		/* has invalid sequence header */
		close_video_stream(&(out->bitstream));
		return 0;
	}

	if(! out->seq.has_sequence_extension){
		/* not MPEG-2 VIDEO stream */
		close_video_stream(&(out->bitstream));
		return 0;
	}

	sequence_header_to_read_picture_header_option(&(out->seq), &(out->pic_opt));
	sequence_header_to_decode_picture_parameter(&(out->seq), &(out->dec_prm));

	while(vs_next_start_code(&(out->bitstream))){
		if(vs_get_bits(&(out->bitstream), 32) == 0x100){
			read_picture_header(&(out->bitstream), &(out->pic), &(out->pic_opt));
			if(out->pic.has_picture_coding_extension){
				out->orig_field_order = out->pic.pc.top_field_first;
			}else{
				out->orig_field_order = TOP_FIELD_FIRST;
			}
			video_stream_seek(&(out->bitstream), 0, SEEK_SET);
			break;
		}
	}

	gp = new_read_gop_parameter(&(out->bitstream), &(out->seq), &(out->pic_opt), out->field_order);
	if(gp == NULL){
		/* malloc failed */
		close_video_stream(&(out->bitstream));
		return 0;
	}

	out->rate = gp->rate;
	out->scale = gp->scale;

	if(read_gop(gp, &g)){
		/* gop timecode is sequential */
		out->fg.arg1 = (void *)gp;
		out->fg.func = find_gop_with_timecode;
		out->fg.release = delete_read_gop_parameter;
		
		out->total = count_frame(gp);
	}

	if(out->total <= 0){
		/* gop timecode is not sequential */
		delete_read_gop_parameter(gp);
		
		strcpy(gl_path, path);
		cut_suffix(gl_path);
		strcat(gl_path, ".gl");

		gl = load_gop_list(gl_path);
		if( (gl == NULL) || (gl->stream_length != out->bitstream.file_length) ){
			video_stream_seek(&(out->bitstream), 0, SEEK_SET);
			gl = new_gop_list(&(out->bitstream), &(out->pic_opt), out->field_order);
			if(gl == NULL){
				/* stream has something probrem */
				close_video_stream(&(out->bitstream));
				return 0;
			}
			store_gop_list(gl, gl_path);
		}
		
		out->fg.arg1 = (void *) gl;
		out->fg.func = find_gop_with_gop_list;
		out->fg.release = delete_gop_list;

		out->total = gl->num_of_frame;
	}

	setup_m2v_config(&(out->config));

	sequence_header_to_bgr_conversion_parameter(&(out->seq), &(out->bgr_prm), &(out->config));

	setup_field_order(out, &(out->config));

	if(out->config.simd & 1){
		out->yuv2bgr = yuv444_to_bgr_mmx;
		out->dec_prm.add_block_func = add_block_data_to_frame_mmx;
		out->dec_prm.mc_parameter.prediction_func = prediction_mmx;
	}else{
		out->yuv2bgr = yuv444_to_bgr;
		out->dec_prm.add_block_func = add_block_data_to_frame;
		out->dec_prm.mc_parameter.prediction_func = prediction;
	}

	setup_idct_function(&(out->dec_prm), &(out->config));

	out->rsz_prm = create_resize_parameter(&(out->seq), &(out->config));

	resize_parameter_to_bgr_conversion_parameter(out->rsz_prm, &(out->bgr_prm));
	
	if(out->seq.has_sequence_extension){
		setup_chroma_upsampling_function(out, out->seq.se.chroma_format, out->config.simd);
	}else{
		setup_chroma_upsampling_function(out, 1, out->config.simd);
	}

	video_stream_seek(&(out->bitstream), 0, SEEK_SET);

	out->dec_buf.forward = NULL;
	out->dec_buf.backward = NULL;
	out->dec_buf.current = NULL;

	out->current.frame_count = 0;
	out->current.start_frame = 0;

	clear_output_parameters(out);

	InitializeCriticalSection(&(out->lock));

	return 1;
}
	
int close_mpeg2video(MPEG2VIDEO *p)
{
	if(p == NULL){
		return 0;
	}

	clear_out_buffer(&(p->out_buf));

	if(p->dec_buf.forward){
		delete_frame(p->dec_buf.forward);
		p->dec_buf.forward = NULL;
	}
	if(p->dec_buf.backward){
		delete_frame(p->dec_buf.backward);
		p->dec_buf.backward = NULL;
	}
	
	if(p->fg.release){
		p->fg.release(p->fg.arg1);
	}

	release_resize_parameter(p->rsz_prm);
	
	close_video_stream(&(p->bitstream));

	DeleteCriticalSection(&(p->lock));
	
	return 1;
}

OUT_BUFFER_ELEMENT *read_frame(MPEG2VIDEO *in, __int64 frame)
{
	int code;
	
	OUT_BUFFER_ELEMENT *r;
	
	EnterCriticalSection(&(in->lock));
	
	r = search_out_buffer(&(in->out_buf), frame);
	if(r){
		goto READ_FRAME_END;
	}

	if(!is_frame_in_current_gop(in, frame)){
		in->current = in->fg.func(in->fg.arg1, frame);
		video_stream_seek(&(in->bitstream), in->current.offset, SEEK_SET);
		if(in->dec_buf.forward){
			delete_frame(in->dec_buf.forward);
			in->dec_buf.forward = NULL;
		}
		if(in->dec_buf.backward){
			delete_frame(in->dec_buf.backward);
			in->dec_buf.backward = NULL;
		}
		
		clear_output_parameters(in);
		
		clear_out_buffer(&(in->out_buf));

		if(in->current.frame_count == 0){
			goto READ_FRAME_END;
		}
		if(in->current.start_frame > frame){
			goto READ_FRAME_END;
		}
	}

	while(vs_next_start_code(&(in->bitstream))){
		code = vs_get_bits(&(in->bitstream), 32);

		if(code == 0x1B3){
			proc_sequence_header(in);
		}else if(code == 0x1B8){
			r = proc_gop_header(in, frame);
		}else if(code == 0x100){
			r = proc_picture_data(in, frame);
		}
		
		if(r){
			goto READ_FRAME_END;
		}
	}
				
	if(in->dec_buf.forward){
		r = store_forward_reference_frame(in, frame);
	}
	
	if(in->dec_buf.backward){
		if(r){
			store_backward_reference_frame(in, frame);
		}else{
			r = store_backward_reference_frame(in, frame);
		}
	}

	in->fwd_prm.index = -1;
	in->bwd_prm.index = -1;

READ_FRAME_END:
	LeaveCriticalSection(&(in->lock));
	return r;
}

static void proc_sequence_header(MPEG2VIDEO *in)
{
	read_sequence_header(&(in->bitstream), &(in->seq));
	sequence_header_to_bgr_conversion_parameter(&(in->seq), &(in->bgr_prm), &(in->config));
	sequence_header_to_read_picture_header_option(&(in->seq), &(in->pic_opt));
	sequence_header_to_decode_picture_parameter(&(in->seq), &(in->dec_prm));
	resize_parameter_to_bgr_conversion_parameter(in->rsz_prm, &(in->bgr_prm)); 
}

static OUT_BUFFER_ELEMENT *proc_gop_header(MPEG2VIDEO *in, __int64 frame)
{
	OUT_BUFFER_ELEMENT *fwd;
	OUT_BUFFER_ELEMENT *bwd;
	
	fwd = NULL;
	bwd = NULL;
	
	vs_erase_bits(&(in->bitstream), 26); /* erase timecode & closed gop flag */
	
	if(vs_get_bits(&(in->bitstream), 1)){ /* broken link */

		if(in->dec_buf.forward){
			fwd = store_forward_reference_frame(in, frame);
		}
	
		if(in->dec_buf.backward){
			bwd = store_backward_reference_frame(in, frame);
		}
		if(fwd){
			return fwd;
		}else if(bwd){
			return bwd;
		}
	}

	return NULL;
}

static OUT_BUFFER_ELEMENT *proc_picture_data(MPEG2VIDEO *in, __int64 frame)
{
	OUT_BUFFER_ELEMENT *r;

	r = NULL;
	
	read_picture_header(&(in->bitstream), &(in->pic), &(in->pic_opt));
			
	picture_header_to_decode_picture_parameter(&(in->pic), &(in->dec_prm));

	in->dec_buf.current = new_frame(in->seq.h_size, in->seq.v_size);
	in->cur_prm.index = in->bwd_prm.index;

	picture_header_to_output_parameter(&(in->pic), &(in->cur_prm));

	if(in->pic.picture_coding_type != 3){
		r = rotate_reference_frame(in, frame);
	}

	if((in->pic.picture_coding_type == 3 ) && ( (in->dec_buf.forward == NULL) || (in->dec_buf.backward == NULL) )){
		delete_frame(in->dec_buf.current);
		in->dec_prm.mc_parameter.first_field = 0;
		return r;
	}else if((in->pic.picture_coding_type == 2) && (in->dec_buf.forward == NULL)){
		delete_frame(in->dec_buf.current);
		in->dec_prm.mc_parameter.first_field = 0;
		return r;
	}

	decode_picture(&(in->bitstream), &(in->dec_buf), &(in->dec_prm));

	if(in->pic.has_picture_coding_extension && in->pic.pc.picture_structure != 3){
		decode_2nd_field(in);
	}

	if(in->pic.picture_coding_type == 3){
		r = store_current_decoding_frame(in, frame);
	}

	return r;
}

static OUT_BUFFER_ELEMENT *store_forward_reference_frame(MPEG2VIDEO *in, __int64 frame)
{
	OUT_BUFFER_ELEMENT *r;
	
	r = add_frame_out_buffer_with_resize(in, in->dec_buf.forward, &(in->fwd_prm));
	in->dec_buf.forward = NULL;

	if(in->fwd_prm.index == frame){
		return r;
	}

	return NULL;
}

static OUT_BUFFER_ELEMENT *store_backward_reference_frame(MPEG2VIDEO *in, __int64 frame)
{
	OUT_BUFFER_ELEMENT *w;
	OUT_BUFFER_ELEMENT *r;

	FRAME *rff_buf;

	r = NULL;
	
	if(in->bwd_prm.repeat_first_field && (in->bwd_prm.top_field_first == in->field_order)){
		rff_buf = copy_frame(in->dec_buf.backward);
		w = add_frame_out_buffer_with_resize(in, rff_buf, &(in->bwd_prm));
		if(in->bwd_prm.index == frame){
			r = w;
		}
		in->bwd_prm.index += 1;
		in->bwd_prm.top_field_first = !(in->field_order);
		in->bwd_prm.repeat_first_field = 0;
		in->bwd_prm.picture_coding_type = 3;
	}
	
	w = add_frame_out_buffer_with_resize(in, in->dec_buf.backward, &(in->bwd_prm));
	if(in->bwd_prm.index == frame){
		r = w;
	}
	
	in->dec_buf.backward = NULL;

	return r;
}

static OUT_BUFFER_ELEMENT *store_current_decoding_frame(MPEG2VIDEO *in, __int64 frame)
{
	OUT_BUFFER_ELEMENT *w;
	OUT_BUFFER_ELEMENT *r;

	FRAME *rff_buf;

	r = NULL;
	
	if(in->cur_prm.repeat_first_field && (in->cur_prm.top_field_first == in->field_order)){
		rff_buf = copy_frame(in->dec_buf.current);
		
		w = add_frame_out_buffer_with_resize(in, rff_buf, &(in->cur_prm));
		
		if(in->cur_prm.index == frame){
			r = w;
		}

		in->cur_prm.index += 1;
		in->cur_prm.top_field_first = !(in->field_order);
		in->cur_prm.repeat_first_field = 0;
		in->cur_prm.picture_coding_type = 3;

		in->bwd_prm.index += 1;
	}

	w = add_frame_out_buffer_with_resize(in, in->dec_buf.current, &(in->cur_prm));
	
	if(in->cur_prm.index == frame){
		r = w;
	}
	
	in->bwd_prm.index += 1;

	return r;
}

static OUT_BUFFER_ELEMENT *rotate_reference_frame(MPEG2VIDEO *in, __int64 frame)
{
	OUT_BUFFER_ELEMENT *w;
	OUT_BUFFER_ELEMENT *r;
	
	FRAME *rff_buf;
	OUTPUT_PARAMETER rff_prm;

	r = NULL;
	
	if(in->dec_buf.forward){
		w = add_frame_out_buffer_with_resize(in, in->dec_buf.forward, &(in->fwd_prm));
		if(in->fwd_prm.index == frame){
			r = w;
		}
	}
	
	in->dec_buf.forward = in->dec_buf.backward;
	in->dec_buf.backward = in->dec_buf.current;
	in->fwd_prm = in->bwd_prm;
	in->bwd_prm = in->cur_prm;
	in->bwd_prm.index += 1;

	if((in->fwd_prm.picture_coding_type == 1) && (in->fwd_prm.index != in->current.start_frame) ){
		in->current.frame_count = in->fwd_prm.index - in->current.start_frame;
		in->current.start_frame = in->fwd_prm.index;
	}

	if(in->fwd_prm.repeat_first_field && (in->fwd_prm.top_field_first == in->field_order) ){
		rff_buf = copy_frame(in->dec_buf.forward);

		rff_prm = in->fwd_prm;
		rff_prm.index += 1;
		rff_prm.picture_coding_type = 3;
		rff_prm.top_field_first = !(in->field_order);
		rff_prm.repeat_first_field = 0;
		
		w = add_frame_out_buffer_with_resize(in, rff_buf, &rff_prm);
		if(rff_prm.index == frame){
			r = w;
		}

		in->bwd_prm.index += 1;
	}

	return r;
}

static int is_frame_in_current_gop(MPEG2VIDEO *p, __int64 frame)
{
	__int64 n;
	
	if(p->fwd_prm.index == frame){
		return 1;
	}

	if(p->bwd_prm.index == frame){
		return 1;
	}

	n = p->current.start_frame + p->current.frame_count;
	if(frame < n){
		if(frame < p->current.start_frame){
			return 0;
		}
		return 1;
	}
	
	return 0;
}

static void sequence_header_to_decode_picture_parameter(SEQUENCE_HEADER *in, DECODE_PICTURE_PARAMETER *out)
{
	sequence_header_to_read_slice_header_option(in, &(out->slice_option));
	sequence_header_to_read_macroblock_option(in, &(out->macroblock_option));
	sequence_header_to_read_block_option(in, &(out->block_option));
	sequence_header_to_mc_parameter(in, &(out->mc_parameter));
}

static void picture_header_to_decode_picture_parameter(PICTURE_HEADER *in, DECODE_PICTURE_PARAMETER *out)
{
	picture_header_to_read_macroblock_option(in, &(out->macroblock_option));
	picture_header_to_read_block_option(in, &(out->block_option));
	picture_header_to_mc_parameter(in, &(out->mc_parameter));
}

static void decode_2nd_field(MPEG2VIDEO *p)
{
	int code;

	__int64 offset;
	
	while(vs_next_start_code(&(p->bitstream))){
		code = vs_read_bits(&(p->bitstream), 32);
		if(code == 0x100){
			offset = video_stream_tell(&(p->bitstream));
			vs_erase_bits(&(p->bitstream), 32);
			
			read_picture_header(&(p->bitstream), &(p->pic), &(p->pic_opt));
			
			if(p->pic.has_picture_coding_extension && (p->pic.pc.picture_structure !=3) ){
				picture_header_to_decode_picture_parameter(&(p->pic), &(p->dec_prm));
				
				decode_picture(&(p->bitstream), &(p->dec_buf), &(p->dec_prm));

			}else{
				video_stream_seek(&(p->bitstream), offset, SEEK_SET);
			}

			return;
		}else{
			vs_erase_bits(&(p->bitstream), 32);
		}
	}
}

static OUT_BUFFER_ELEMENT *add_frame_out_buffer_with_resize(MPEG2VIDEO *p, FRAME *data, OUTPUT_PARAMETER *frame_prm)
{
	p->upsmp_c[frame_prm->progressive_frame](data);
	data = resize(data, p->rsz_prm);
	return add_frame_out_buffer(&(p->out_buf), data, frame_prm);
}

static int is_registered_suffix(char *filepath)
{
	int i;
	
	static char *registered_suffix[] = {
		".mpeg",
		".mpg",
		".m2p",
		".vob",
		".m2v",
		".mpv",
		"", /* sentinel */
	};

	i = 0;
	while(registered_suffix[i][0]){
		if(check_suffix(filepath, registered_suffix[i])){
			return 1;
		}
		i += 1;
	}

	return 0;
}

static void setup_chroma_upsampling_function(MPEG2VIDEO *p, int chroma_format, int simd)
{
	if(simd & 1){
		switch(chroma_format){
		case 1: /* 420 */
			p->upsmp_c[0] = upsample_chroma_420i_mmx;
			p->upsmp_c[1] = upsample_chroma_420p_mmx;
			break;
		case 2: /* 422 */
			p->upsmp_c[0] = upsample_chroma_422_mmx;
			p->upsmp_c[1] = upsample_chroma_422_mmx;
			break;
		case 3:
			p->upsmp_c[0] = upsample_chroma_444;
			p->upsmp_c[1] = upsample_chroma_444;
			break;
		default:
			p->upsmp_c[0] = upsample_chroma_420i_mmx;
			p->upsmp_c[1] = upsample_chroma_420p_mmx;
			break;
		}
	}else{
		switch(chroma_format){
		case 1: /* 420 */
			p->upsmp_c[0] = upsample_chroma_420i;
			p->upsmp_c[1] = upsample_chroma_420p;
			break;
		case 2: /* 422 */
			p->upsmp_c[0] = upsample_chroma_422;
			p->upsmp_c[1] = upsample_chroma_422;
			break;
		case 3:
			p->upsmp_c[0] = upsample_chroma_444;
			p->upsmp_c[1] = upsample_chroma_444;
			break;
		default:
			p->upsmp_c[0] = upsample_chroma_420i;
			p->upsmp_c[1] = upsample_chroma_420p;
			break;
		}
	}
}

static void setup_idct_function(DECODE_PICTURE_PARAMETER *p, M2V_CONFIG *prm)
{
	if(prm->idct_type == M2V_CONFIG_IDCT_REFERENCE){
		if(prm->simd & M2V_CONFIG_USE_SSE){
			p->idct_func = idct_sse32;
		}else{
			p->idct_func = idct_double;
		}
	}else{
		if(prm->simd & M2V_CONFIG_USE_MMX){
			p->idct_func = idct_mmx32;
		}else{
			p->idct_func = idct_int32;
		}
	}
}

static void clear_output_parameters(MPEG2VIDEO *p)
{
	memset(&(p->fwd_prm), 0, sizeof(OUTPUT_PARAMETER));
	memset(&(p->cur_prm), 0, sizeof(OUTPUT_PARAMETER));
	memset(&(p->bwd_prm), 0, sizeof(OUTPUT_PARAMETER));
		
	p->fwd_prm.index = p->current.start_frame - 1;
	p->bwd_prm.index = p->current.start_frame - 1;
}

static void setup_m2v_config(M2V_CONFIG *p)
{
	p->simd = get_simd_mode();
	p->bt601 = get_color_conversion_type();
	p->aspect_ratio = get_resize_mode();
	p->field_mode = get_field_mode();
	p->idct_type = get_idct_type();
}

static void setup_field_order(MPEG2VIDEO *p, M2V_CONFIG *prm)
{
	switch(prm->field_mode){
	case 0:
		p->field_order = p->orig_field_order;
		break;
	case 1:
		p->field_order = TOP_FIELD_FIRST;
		break;
	case 2:
		p->field_order = BOTTOM_FIELD_FIRST;
		break;
	default:
		p->field_order = p->orig_field_order;
	}

	if(p->field_order != p->orig_field_order){
		p->total -= 1;
	}
}

static void resize_parameter_to_bgr_conversion_parameter(RESIZE_PARAMETER *in, BGR_CONVERSION_PARAMETER *out)
{
	if(in == NULL){
		return;
	}

	out->width = in->width;
	out->height = in->height;

	out->in_step = in->out_step;
}