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

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

#define MPEG2VIDEO_C
#include "mpeg2video.h"

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

static int get_color_conversion_type();
static int get_simd_mode();
static void select_idct_function(DECODE_PICTURE_PARAMETER *p);

static void delete_data_bgr_buffer(BGR_BUFFER *p);
static void add_data_bgr_buffer(BGR_BUFFER *p, unsigned char *data, __int64 index, int i_frame);
static BGR_BUFFER_ELEMENT *search_bgr_buffer(BGR_BUFFER *p, __int64 index);

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);

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

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

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

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

	code = 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;
	}

	out->remap = get_color_conversion_type();
	sequence_header_to_bgr_conversion_parameter(&(out->seq), &(out->cur_prm.prm), out->remap);
	out->fwd_prm.prm = out->cur_prm.prm;
	out->bwd_prm.prm = out->cur_prm.prm;
	out->fwd_prm.index = -1;
	out->bwd_prm.index = -1;

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

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

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

	gp = new_read_gop_parameter(&(out->bitstream), &(out->seq), &(out->pic_opt));
	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);
		
		video_stream_seek(&(out->bitstream), 0, SEEK_SET);
		gl = new_gop_list(&(out->bitstream), &(out->pic_opt));
		if(gl == NULL){
			/* stream has no gop layer */
			close_video_stream(&(out->bitstream));
			return 0;
		}
		
		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;
	}
	
	video_stream_seek(&(out->bitstream), 0, SEEK_SET);

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

	select_idct_function(&(out->dec_prm));

	simd = get_simd_mode();
	if(simd & 1){
		out->frm2bgr = frame_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->frm2bgr = frame_to_bgr;
		out->dec_prm.add_block_func = add_block_data_to_frame;
		out->dec_prm.mc_parameter.prediction_func = prediction;
	}

	InitializeCriticalSection(&(out->lock));

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

	while(p->out_buf.head){
		delete_data_bgr_buffer(&(p->out_buf));
	}

	if(p->fg.release){
		p->fg.release(p->fg.arg1);
	}

	close_video_stream(&(p->bitstream));

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

int read_frame(MPEG2VIDEO *in, unsigned char **out, __int64 frame)
{
	int code;
	
	BGR_BUFFER_ELEMENT *w;

	unsigned char *buffer;

	EnterCriticalSection(&(in->lock));

	w = search_bgr_buffer(&(in->out_buf), frame);
	if(w){
		*out = w->data;
		goto READ_FRAME_NORMAL_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;
		}
		in->fwd_prm.index = in->current.start_frame - 1;
		in->bwd_prm.index = in->current.start_frame - 1;
		in->fwd_prm.i_frame = 0;
		in->bwd_prm.i_frame = 0;
		while(in->out_buf.head){
			delete_data_bgr_buffer(&(in->out_buf));
		}
		if(in->current.frame_count == 0){
			goto READ_FRAME_ABNORMAL_END;
		}
		if(in->current.start_frame > frame){
			goto READ_FRAME_ABNORMAL_END;
		}
	}
		
	*out = NULL;
	while(next_start_code(&(in->bitstream))){
		code = get_bits(&(in->bitstream), 32);

		if(code == 0x1B3){ /* sequence */
			
			read_sequence_header(&(in->bitstream), &(in->seq));
			
			sequence_header_to_bgr_conversion_parameter(&(in->seq), &(in->cur_prm.prm), in->remap);
			in->fwd_prm.prm = in->cur_prm.prm;
			in->bwd_prm.prm = in->cur_prm.prm;

			sequence_header_to_read_picture_header_option(&(in->seq), &(in->pic_opt));

			sequence_header_to_decode_picture_parameter(&(in->seq), &(in->dec_prm));
			
		}else if(code == 0x1B8){ /* GOP */
			
			erase_bits(&(in->bitstream), 26);
			if(get_bits(&(in->bitstream), 1)){ /* broken link */

				if(in->dec_buf.forward){
					buffer = (unsigned char *)malloc(in->seq.h_size * in->seq.v_size * 3);
					if(buffer == NULL){
						goto READ_FRAME_ABNORMAL_END;
					}
					in->frm2bgr(in->dec_buf.forward, buffer, &(in->fwd_prm.prm));
					add_data_bgr_buffer(&(in->out_buf), buffer, in->fwd_prm.index, in->fwd_prm.i_frame);
					delete_frame(in->dec_buf.forward);
					in->dec_buf.forward = NULL;
					if(in->fwd_prm.index == frame){
						*out = buffer;
					}
				}
	
				if(in->dec_buf.backward){
					buffer = (unsigned char *)malloc(in->seq.h_size * in->seq.v_size * 3);
					if(buffer == NULL){
						goto READ_FRAME_ABNORMAL_END;
					}
					in->frm2bgr(in->dec_buf.backward, buffer, &(in->bwd_prm.prm));
					add_data_bgr_buffer(&(in->out_buf), buffer, in->bwd_prm.index, in->bwd_prm.i_frame);
					delete_frame(in->dec_buf.backward);
					in->dec_buf.backward = NULL;
					if(in->bwd_prm.index == frame){
						*out = buffer;
					}
				}

				if(out){
					goto READ_FRAME_NORMAL_END;
				}
			}
			
		}else if(code == 0x100){ /* picture */

			read_picture_header(&(in->bitstream), &(in->pic), &(in->pic_opt));
			
			picture_header_to_decode_picture_parameter(&(in->pic), &(in->dec_prm));
			picture_header_to_bgr_conversion_parameter(&(in->pic), &(in->cur_prm.prm));

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

			switch(in->pic.picture_coding_type){
			case 1:
				in->cur_prm.i_frame = 1;
			case 2:
				if(in->dec_buf.forward){
					buffer = (unsigned char *)malloc(in->seq.h_size * in->seq.v_size * 3);
					if(buffer == NULL){
						goto READ_FRAME_ABNORMAL_END;
					}
					in->frm2bgr(in->dec_buf.forward, buffer, &(in->fwd_prm.prm));
					add_data_bgr_buffer(&(in->out_buf), buffer, in->fwd_prm.index, in->fwd_prm.i_frame);
					delete_frame(in->dec_buf.forward);
					if(in->fwd_prm.index == frame){
						*out = buffer;
					}
				}
				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.i_frame && (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->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;
				continue;
			}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;
				continue;
			}

			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){
				buffer = (unsigned char *)malloc(in->seq.h_size * in->seq.v_size * 3);
				if(buffer == NULL){
					goto READ_FRAME_ABNORMAL_END;
				}
				in->frm2bgr(in->dec_buf.current, buffer, &(in->cur_prm.prm));
				add_data_bgr_buffer(&(in->out_buf), buffer, in->cur_prm.index, 0);
				delete_frame(in->dec_buf.current);
				
				in->bwd_prm.index += 1;
				
				if(in->cur_prm.index == frame){
					*out = buffer;
					goto READ_FRAME_NORMAL_END;
				}
			}

			if(*out){
				goto READ_FRAME_NORMAL_END;
			}
		}
	}
				
	if(in->dec_buf.forward){
		buffer = (unsigned char *)malloc(in->seq.h_size * in->seq.v_size * 3);
		if(buffer == NULL){
			goto READ_FRAME_ABNORMAL_END;
		}
		in->frm2bgr(in->dec_buf.forward, buffer, &(in->fwd_prm.prm));
		add_data_bgr_buffer(&(in->out_buf), buffer, in->fwd_prm.index, in->fwd_prm.i_frame);
		delete_frame(in->dec_buf.forward);
		in->dec_buf.forward = NULL;
		if(in->fwd_prm.index == frame){
			*out = buffer;
		}
	}
	
	if(in->dec_buf.backward){
		buffer = (unsigned char *)malloc(in->seq.h_size * in->seq.v_size * 3);
		if(buffer == NULL){
			goto READ_FRAME_ABNORMAL_END;
		}
		in->frm2bgr(in->dec_buf.backward, buffer, &(in->bwd_prm.prm));
		add_data_bgr_buffer(&(in->out_buf), buffer, in->bwd_prm.index, in->bwd_prm.i_frame);
		delete_frame(in->dec_buf.backward);
		in->dec_buf.backward = NULL;
		if(in->bwd_prm.index == frame){
			*out = buffer;
		}
	}

	if(*out){
		goto READ_FRAME_NORMAL_END;
	}

READ_FRAME_ABNORMAL_END:
	LeaveCriticalSection(&(in->lock));
	__asm{finit}
	return 0;
	
READ_FRAME_NORMAL_END:
	LeaveCriticalSection(&(in->lock));
	__asm{finit}
	return 1;
}

static int get_color_conversion_type()
{
	HKEY key;
	DWORD type;
	DWORD size;
	DWORD value;

	if(RegOpenKeyEx(HKEY_CURRENT_USER, "Software\\marumo\\mpeg2vid_vfp", 0, KEY_READ, &key) != ERROR_SUCCESS){
		return 1;
	}

	type = REG_DWORD;
	size = sizeof(DWORD);
	
	if(RegQueryValueEx(key, "re_map", NULL, &type, (LPBYTE)&value, &size) != ERROR_SUCCESS){
		RegCloseKey(key);
		return 1;
	}

	RegCloseKey(key);

	if(value){
		return 1;
	}

	return 0;
}

static void select_idct_function(DECODE_PICTURE_PARAMETER *p)
{
	HKEY key;
	DWORD type;
	DWORD size;
	DWORD value;
	DWORD simd;

	if(RegOpenKeyEx(HKEY_CURRENT_USER, "Software\\marumo\\mpeg2vid_vfp", 0, KEY_READ, &key) != ERROR_SUCCESS){
		p->idct_func = idct_int32;
		return;
	}

	type = REG_DWORD;
	size = sizeof(DWORD);
	
	if(RegQueryValueEx(key, "idct_func", NULL, &type, (LPBYTE)&value, &size) != ERROR_SUCCESS){
		RegCloseKey(key);
		p->idct_func = idct_int32;
		return;
	}

	if(RegQueryValueEx(key, "simd", NULL, &type, (LPBYTE)&simd, &size) != ERROR_SUCCESS){
		simd = 0;
	}
	
	RegCloseKey(key);

	switch(value){
	case 0:
		if(simd & 2){
			p->idct_func = idct_sse32;
		}else{
			p->idct_func = idct_double;
		}
		return; 
	case 1:
		p->idct_func = idct_int32;
		return;
	default:
		p->idct_func = idct_int32;
	}
	
	return;
}

static int get_simd_mode()
{
	HKEY key;
	DWORD type;
	DWORD size;
	DWORD value;

	if(RegOpenKeyEx(HKEY_CURRENT_USER, "Software\\marumo\\mpeg2vid_vfp", 0, KEY_READ, &key) != ERROR_SUCCESS){
		return 0;
	}

	type = REG_DWORD;
	size = sizeof(DWORD);
	
	if(RegQueryValueEx(key, "simd", NULL, &type, (LPBYTE)&value, &size) != ERROR_SUCCESS){
		RegCloseKey(key);
		return 0;
	}

	RegCloseKey(key);

	if(value & 1){
		return 1;
	}

	return 0;
}

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 delete_data_bgr_buffer(BGR_BUFFER *p)
{
	BGR_BUFFER_ELEMENT *w;

	w = p->head->next;

	if(p->head->i_frame){
		p->i_frame_count -= 1;
	}
	free(p->head->data);
	free(p->head);

	if(w){
		p->head = w;
		p->head->prev = NULL;
	}else{
		p->head = NULL;
		p->tail = NULL;
	}
}

static void add_data_bgr_buffer(BGR_BUFFER *p, unsigned char *data, __int64 index, int i_frame)
{
	BGR_BUFFER_ELEMENT *w, *prev, *next;

	w = (BGR_BUFFER_ELEMENT *)malloc(sizeof(BGR_BUFFER_ELEMENT));
	w->index = index;
	w->i_frame = i_frame;
	w->data = data;
	w->prev = NULL;
	w->next = NULL;
	
	if(i_frame){
		p->i_frame_count += 1;
	}

	prev = p->tail;
	next = NULL;
	while(prev){
		if(prev->index < index){
			prev->next = w;
			w->prev = prev;
			w->next = next;
			if(next == NULL){
				p->tail = w;
			}else{
				next->prev = w;
			}
			if(p->i_frame_count > 1 || ( (!p->head->i_frame) && (p->i_frame_count > 0) ) ){
				delete_data_bgr_buffer(p);
			}
			return;
		}else{
			next = prev;
			prev = prev->prev;
		}
	}

	if(p->tail == NULL){
		p->head = w;
		p->tail = w;
	}else{
		w->next = p->head;
		p->head->prev = w;
		p->head = w;
	}

	if(p->i_frame_count > 1){
		delete_data_bgr_buffer(p);
	}
}

static BGR_BUFFER_ELEMENT *search_bgr_buffer(BGR_BUFFER *p, __int64 index)
{
	BGR_BUFFER_ELEMENT *r;

	r = p->tail;

	while(r){
		if(r->index == index){
			return r;
		}else{
			r = r->prev;
		}
	}

	return NULL;
}

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(next_start_code(&(p->bitstream))){
		code = get_bits(&(p->bitstream), 32);
		if(code == 0x100){
			offset = video_stream_tell(&(p->bitstream)) - 4;
			
			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;
		}
	}
}