#define WIN32_LEAN_AND_MEAN		// Exclude rarely-used stuff from Windows headers
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <fcntl.h>

#ifdef WIN32
#include <io.h>
#endif

#define VIDEO_ID		0xE0
#define AUDIO_ID		0xC0
#define NON_EMPTY_ID	0x03
#define CHECKSUM		0x3FC
#define MAGIC_SIGNATURE 0x91231ebc
#define CHUNK_SIZE		(0x20000)
#define min(a, b)		((a) <= (b) ? (a) : (b))

//#define SCAN_TO_NEXT_HEADER
//#define EXHAUSTIVE_SCAN_TO_HEADER
#define CHECK_SEQ


int verbose = 0;
double begin_audio_time = 0;
double begin_video_time = 0;

double last_audio_time = 0;
double last_video_time = 0;

unsigned long int	predict_audio_sequence = 0;
unsigned long int	predict_video_sequence = 0;

int previous_chunk_bad = 0;
int quit_after_time_calc = 0;

/* Report Time-Offset Only */
/* Added by Stealth Dave */
#define TIME_OFFSET_STR	"--time-offset"

/*

  Ver 0.1
    - first release
	- skip chunks with bad seq entirely.
	- better video header location (thanks to displaystream)
	- non-audio/video record with payload! (type = 0x03) => fixes "audio failure" and "failure" problems
	- audio seq matching
	- dump seq mismatched chunks
	- audio start from the beginning of first chunk
	- read from standard in

  Ver 0.2
	- small code cleanup
	- relax 0 byte audio record for seq. reset
	- do memory checks to avoid crashes in bad data
	- add large seq warnings
	- add win32 ifdef for clean compile on linux

  Ver 0.3
    - really skip sequences with bad memory pointers
	- provide estimates for initial time offset

  Ver 0.4
    - added '--time-offset' flag to return only the audio offset information (Stealth Dave)
	- ignore video seq mismatches of 0x54 and 0x58
	- timestamp messages
	- Turn off warnings for consequetive chunks
	- if a chunk is skipped, reset the sequence to the last good value

  Ver 0.5
	- fixed audio delay calculation when no audio was found in the first record.
    - fixed PES time decoder
	- turn off large seq mismatch message on consequetive chunks
	- added total time display at the end.
	- --time-offset now quits as soon as it prints the message

  TODO:
	- large seq matching

  WEIRDNESS:
    - Why are some chunks duplicated? (1 case in 1.2 GB file)
	- Why are some chunks near power of 2 boundaries (2048, 4096, 8192, etc) out of sequence?
	- Why are there lots of out of sequence chunks at the end of the stream?


  INFO:
	From Charlie, re: time in PES headers
	The whole timecode is 33 bits long (yes, wonderful isn't it) and it's
	stored in 5 bytes.  It looks like this (starting with the MSB):

	what's there:    How many bits it is:		Byte in PES Header data
	  '0010'              4								9
	  PTS [32..30]        3								9
	  marker bit          1								9
	  PTS [29..15]       15								10, 11
	  marker bit          1								11
	  PTS [14..0]        15								12, 13
	  marker bit          1								13

	"PTS" stands for Presentation Time Stamp it indicates the exact time
	that the audio or video should be presented to the viewer.  The PTS
	increments at a rate of 90KHz


*/

struct TyStreamHeader
{
	int				type;
	int				PESHeader;
	int				UnRecognized;
	int				size;
	unsigned char	*header;
	unsigned char	*data;
};

static void parse_chunk
			(
				FILE			*audio_out_fd,
				FILE			*video_out_fd,
				FILE			*bad_chunk_out_fd,
				unsigned char	*buf,
				unsigned long	chunk_num
			);

int			main(int argc, char *argv[]);

/*
 =======================================================================================================================
 =======================================================================================================================
 */

static void print_time(double time)
{
	int hour;
	int minute;
	int second;
	int millisecond;

	hour = (int)(time/3600000);
	minute = (int) (time - hour * 3600000)/60000;
	second = (int) (time - (hour * 360000 + minute * 60000))/1000;
	millisecond = (int) (time - (hour * 360000 + minute * 60000 + second * 1000));
	fprintf(stderr, "%2.2d:%2.2d:%2.2d.%d", hour, minute, second, millisecond);
}

static void print_last_time()
{
	fprintf(stderr, " last audio time ");
	print_time(last_audio_time);
	fprintf(stderr, " last video time ");
	print_time(last_video_time);
	fprintf(stderr, "\n");
}

static unsigned long int getCurrentSequence(TyStreamHeader *tys)
{
	unsigned long int tempSequence = (((unsigned long int) tys->header[5]) << 16) +
		(((unsigned long int) tys->header[6]) << 8) + ((unsigned long int) tys->header[7]);
	return tempSequence;
}

static int check_audio_sequence(TyStreamHeader *audio_tys, unsigned long chunk_num,
								int record_num)
{

	if (audio_tys->PESHeader) return 0; // no matching on PES Headers
#ifdef CHECK_SEQ
	unsigned long int current_seq = getCurrentSequence(audio_tys);


	if (verbose) fprintf(stderr, "chunk %d record %d seq %x pred %x\n", chunk_num, record_num,
		current_seq, predict_audio_sequence);

	if(!predict_audio_sequence || (current_seq == 0x2D1C40 && predict_audio_sequence != 0x2D1C40))
	{
		/* start/restart */
		if (audio_tys->size != 0 && predict_audio_sequence != 0)
		{
			fprintf(stderr, "Warning: audio sequence 0x2D1C40 at chunk %d record %d, but non-zero size\n",
				chunk_num, record_num);
			print_last_time();
		}
		predict_audio_sequence = current_seq;

	}

	if (audio_tys->size == 0 && current_seq != 0x2D1C40)
	{
		fprintf(stderr, "hmpf: audio sequence not 0x2D1C40 at chunk %d record %d\n",
			chunk_num, record_num);
		print_last_time();
	}

	if(current_seq != predict_audio_sequence)
	{
		if (!previous_chunk_bad)
		{
			fprintf
			(
				stderr,
				"Audio sequence number bogus %x!=%x at chunk %d record %d",
				predict_audio_sequence,
				current_seq,
				chunk_num,
				record_num
			);
			print_last_time();
		}
		return 1; // not matched
	}
	predict_audio_sequence += audio_tys->size;

#endif
	return 0; // matched
}

static int check_video_sequence(TyStreamHeader *video_tys, unsigned long chunk_num, int record_num)
{

	if (video_tys->PESHeader) return 0; // no matching on PES Headers
#ifdef CHECK_SEQ
	unsigned long int current_seq = getCurrentSequence(video_tys);


	if (verbose) fprintf(stderr, "chunk %d record %d seq %x pred %x\n", chunk_num, record_num,
		current_seq, predict_video_sequence);
	if((!predict_video_sequence && current_seq >= 0x200000) || (current_seq & 0xF0F0F0) == 0x205040)	/* start/restart */
		predict_video_sequence = current_seq;

	if(video_tys->header[4] > 0x30)
	{
		/*~~~~~~~~~~~~~~*/
		int k, chksum = 0;
		/*~~~~~~~~~~~~~~*/

		for(k = 0; k < 8; k++) chksum += video_tys->header[k];
		if(chksum != CHECKSUM)
		{
			fprintf(stderr, "Checksum bogus %x!=%x at record %d", chksum, CHECKSUM, chunk_num);
			print_last_time();
			return 1;
		}
	}
	else
	{
		if(current_seq != predict_video_sequence && video_tys->header[5] >= 0x20)
		{
			if (current_seq - predict_video_sequence == 0x54 ||
				current_seq - predict_video_sequence == 0x58)
			{
				predict_video_sequence = current_seq;
				fprintf
				(
					stderr,
					"Warning Video Sequence number out of wack by a little at chunk %d record %d",
					chunk_num,
					record_num
				);
				print_last_time();

			}
			else
			{
				if (!previous_chunk_bad)
				{
					fprintf
					(
						stderr,
						"Video Sequence number bogus %x!=%x at chunk %d record %d",
						predict_video_sequence,
						current_seq,
						chunk_num,
						record_num
					);
					print_last_time();
				}
				return 1; // not matched
			}
		}
	}
	if
	(
		video_tys->header[5] >= 0x20	/* data */
		||	video_tys->header[2] == 0x42	/* marker */
		||	video_tys->header[2] == 0x8c) /* frame beginning */
	{
		predict_video_sequence += video_tys->size;
	}

#endif
	return 0; // matched
}

static long unsigned int get32bit(unsigned char *buf)
{
	unsigned long int temp = (((unsigned long int) buf[0]) << 24) +
		(((unsigned long int) buf[1]) << 16) +
		(((unsigned long int) buf[2]) << 8) +
		((unsigned long int) buf[3]);
	return temp;
}

static double getTime(unsigned char *buf) /* buf is the pes header data, returns value in ms */
{
	double result = 0;
	unsigned char temp;

	temp = (buf[9] & 0xE) >> 1; /* strip out the initial and marker */
	result = ((double) temp) * ((double) (1L << 30))/ ((double)90); /* clock ticks at 90 kHz */
	temp = buf[10];
	result += ((double) temp)* ((double) (1L << 22))/((double)90);
	temp = (buf[11] & 0xFE) >> 1;
	result += ((double) temp)* ((double) (1L << 15))/((double)90);
	temp = buf[12];
	result += ((double) temp)* ((double) (1L << 7))/((double)90);
	temp = (buf[13] & 0xFE) >> 1;
	result += ((double) temp)/((double)90);

	return result;

}

static int print_audio_delay(struct TyStreamHeader TyStream[], int num_recs, int first_video)
{
	int i;
	static double first_video_time;
	double min_video_time = 1e30;
	static int found_audio = 0;
	static int found_video = 0;
	static int found_min_video = 0;

	for(i=0;i<num_recs;i++)
	{
		if (!found_audio && TyStream[i].PESHeader && TyStream[i].type == AUDIO_ID)
		{
			found_audio = 1;
			begin_audio_time = getTime(TyStream[i].data);
		}
		if (i >= first_video && TyStream[i].PESHeader && TyStream[i].type == VIDEO_ID)
		{
			double video_time;
			video_time = getTime(TyStream[i].data);
			if (min_video_time > video_time)
			{
				min_video_time = video_time;
			}
			if (!found_video)
			{
				found_video = 1;
				first_video_time = video_time;
			}


		}
	}
	if (found_video)
	{
		if (!found_min_video)
		{
			found_min_video = 1;
			begin_video_time = min_video_time;
		}
	}

	if (found_audio && found_video)
	{
		fprintf(stderr, "estimated audio delay from first video %f ms\n", begin_audio_time-first_video_time);
		fprintf(stderr, "estimated audio delay from earliest video %f ms\n", begin_audio_time-begin_video_time);
		if (quit_after_time_calc) exit(0);
	}


	return(found_audio && found_video);

}


static void parse_chunk(FILE *audio_out_fd, FILE *video_out_fd, FILE *bad_chunk_out_fd,
						unsigned char *buf, unsigned long chunk_num)
{
	/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/
	struct TyStreamHeader		TyStream[0xff];
	int							num_recs = (int) buf[0] /* +(((int)buf[1]) << 8) */ ;
	int							i;
	static int					foundfirstframe = 0;
	int							first_video = -1;
	static int					expect_more = 0;
	int							header_pos;
	int							data_pos;
	static unsigned long int	large_seq_high = 0;
	static unsigned long int	large_seq_low = 0;
	int							last_av_record = 0;
	int							check_large_seq = 1;
	static int					calc_audio_delay = 1;
	unsigned long int			previous_audio_seq = predict_audio_sequence;
	unsigned long int			previous_video_seq = predict_video_sequence;
	/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/
	static unsigned char audio_header[3] = {0xFF, 0xFD, 0xA8};

	if(buf[0] == 0xF5 && buf[1] == 0x46 && buf[2] == 0x7a && buf[3] == 0xbd)
	{
		fprintf(stderr, "Bad chunk (header?) %d, skipping\n", chunk_num);
		return;
	}

	if (buf[1] != 0x00)
	{
		fprintf(stderr, "Warning: second byte in chunk %d is non-zero %x:%x\n",chunk_num,
			(unsigned int) buf[0],(unsigned int)buf[1]);
	}

	if(verbose)
	{
		fprintf(stderr, "Records = %d\n", num_recs);
	}

	// first pass: fill in TyStream array and find the first record to be dumped.
	data_pos = num_recs*16+4;
	for(i=0;i<num_recs;i++)
	{
		header_pos = i*16+4;
		TyStream[i].header = buf + header_pos;
		TyStream[i].data = buf + data_pos;
		TyStream[i].type = buf[header_pos + 3];
		TyStream[i].PESHeader = 0;
/*
		if((buf[header_pos + 1] == 0x01) && (buf[header_pos + 2] == 0x03 || buf[header_pos + 2] == 0x06))
*/
		if ((TyStream[i].type == AUDIO_ID && TyStream[i].header[2] == 0x03) ||
			(TyStream[i].type == VIDEO_ID && TyStream[i].header[2] == 0x06))
		{
			TyStream[i].PESHeader = 1;
		}
		if (TyStream[i].type == AUDIO_ID || TyStream[i].type == VIDEO_ID)
		{
			if (check_large_seq)
			{
				unsigned long int high = get32bit(TyStream[i].header+8);
				unsigned long int low = get32bit(TyStream[i].header+12);

				if ((high != large_seq_high || low != large_seq_low) && !previous_chunk_bad)
				{
					fprintf(stderr, "Warning: large sequence mismatch. Expecting %x:%x, got %x:%x in chunk %d record %d",
						large_seq_high, large_seq_low,high,low, chunk_num, i);
					print_last_time();
				}
				check_large_seq = 0;
			}

			last_av_record = i;
			TyStream[i].size = (((unsigned int) buf[header_pos]) << 12) | (((unsigned int) buf[header_pos + 1]) << 4) | (((unsigned int) buf[header_pos + 2]) >> 4);
			data_pos += TyStream[i].size;
			if (data_pos > CHUNK_SIZE)
			{
				if (bad_chunk_out_fd != NULL)
					fwrite(buf, sizeof(char), CHUNK_SIZE, bad_chunk_out_fd);
				fprintf(stderr, "Data points past chunk size, skipping chunk %d", chunk_num);
				print_last_time();
				return;
			}
			if(!foundfirstframe)
			{
				if(TyStream[i].type == VIDEO_ID)
				{
					if((buf[header_pos + 2] & 0xF) == 0x7)
					{	/* video header record */
						foundfirstframe = 1;
						first_video = i;
						if (verbose) fprintf(stderr, "first video %d\n",first_video);
					}
				}
			}
			if (foundfirstframe) //note: this is not the same as "else"
			{
				int bad = 0;
				if (TyStream[i].type == VIDEO_ID)
				{
					bad = check_video_sequence(&(TyStream[i]), chunk_num, i);
				}
				else if (TyStream[i].type == AUDIO_ID)
				{
					bad = check_audio_sequence(&(TyStream[i]), chunk_num, i);
				}

				if (bad)
				{
					if (bad_chunk_out_fd != NULL)
						fwrite(buf, sizeof(char), CHUNK_SIZE, bad_chunk_out_fd);
					if (!previous_chunk_bad)
					{
						fprintf(stderr, "Chunk %d has bad sequence, ignoring chunk ", chunk_num);
						print_last_time();
					}
					previous_chunk_bad = bad;
					predict_audio_sequence = previous_audio_seq;
					predict_video_sequence = previous_video_seq;
					return;
				}
			}
		}
		else
		{
			if (TyStream[i].type == NON_EMPTY_ID)
			{
				TyStream[i].size = (((unsigned int) buf[header_pos]) << 12) | (((unsigned int) buf[header_pos + 1]) << 4) | (((unsigned int) buf[header_pos + 2]) >> 4);
				data_pos += TyStream[i].size;
			}
		}
	}
	if(!foundfirstframe)
	{
		fprintf(stderr, "No first frame, skipping chunk %d\n", chunk_num);
		return;
	}

	if (previous_chunk_bad)
	{
		fprintf(stderr, "chunk %d is good.\n", chunk_num);
	}
	previous_chunk_bad = 0;

	large_seq_high = get32bit(TyStream[last_av_record].header+8);
	large_seq_low = get32bit(TyStream[last_av_record].header+12);
	if (calc_audio_delay)
	{
		calc_audio_delay = !print_audio_delay(TyStream, num_recs, first_video);
	}
	for(i=0;i<num_recs;i++)
	{
		if (TyStream[i].type == AUDIO_ID)
		{
			if (!TyStream[i].PESHeader)
			{
				if(!expect_more && TyStream[i].data[0] != 0xFF && TyStream[i].data[1] != 0xFD
					&& TyStream[i].data[2] != 0xA8)
				{
					if(TyStream[i].size != 0)
					{
						fprintf
						(
							stderr,
							"audio failure %d:0x%x 0x%x:0x%x:0x%x chunk %d",
							i,
							TyStream[i].size,
							TyStream[i].data[0],
							TyStream[i].data[1],
							TyStream[i].data[2],
							chunk_num
						);
						print_last_time();
#ifdef SCAN_TO_NEXT_HEADER
						for(j=0;j<TyStream[i].size;j++)
						{
							if (TyStream[i].data[j] == 0xFF && TyStream[i].data[j+1] == 0xFD &&
								TyStream[i].data[j+2] == 0xA8)
							{
								fprintf(stderr, "found new header at offset %d\n",j);
								if (audio_out_fd != NULL)
									fwrite(&(TyStream[i].data[j]), sizeof(char), TyStream[i].size-j, audio_out_fd);
							}
						}
#endif
#ifdef EXHAUSTIVE_SCAN_TO_HEADER
						for(j=num_recs*16+4;j<CHUNK_SIZE;j++)
						{
							if (buf[j] == 0xFF && buf[j+1] == 0xFD &&
								buf[j+2] == 0xA8)
							{
								fprintf(stderr, "found audio header at offset %d relative %d\n",j, (buf+j)-TyStream[i].data);
							}
						}
#endif
					}
				}
				else
				{
					if(TyStream[i].data[0] == 0xFF && TyStream[i].data[1] == 0xFD &&
							TyStream[i].data[2] == 0xA8)
					{
						if(TyStream[i].size < 0x360)
						{
							expect_more = 1;
						}
						else expect_more = 0;
					}
					if (audio_out_fd != NULL)
						fwrite(TyStream[i].data, sizeof(char), TyStream[i].size, audio_out_fd);
				}
			}
			else /* PES Header */
			{
				last_audio_time = getTime(TyStream[i].data) - begin_audio_time;
				if (last_audio_time < 0)
				{
					fprintf(stderr, "Warning: found even earlier Audio frame, add %f to min audio delay chunk %d\n",
						last_audio_time, chunk_num);

				}
			}
		}
		else if (TyStream[i].type == VIDEO_ID)
		{
			if (i >= first_video && !TyStream[i].PESHeader)
			{
				if (video_out_fd != NULL)
					fwrite(TyStream[i].data, sizeof(char), TyStream[i].size, video_out_fd);
			}
			else if (TyStream[i].PESHeader)
			{
				last_video_time =  getTime(TyStream[i].data) - begin_video_time;
				if (last_video_time < 0)
				{
					fprintf(stderr, "Warning: found even earlier Video frame, add %f to min audio delay chunk %d\n",
						-last_video_time, chunk_num);
				}
			}

		}
		else
		{
			// unknown record type; ignore it
			if (TyStream[i].type > 0x03)
			{
				fprintf(stderr, "Warning: unknown record type %d in chunk %d record %d\n",
					TyStream[i].type, chunk_num, i);
			}
		}
	}
}

/*
 =======================================================================================================================
    }
 =======================================================================================================================
 */
void usage(void)
{
	fprintf(stdout, "Usage: \n");
	fprintf(stdout, "splitstream stream.ty program.m2a program.m2v badchunk.ty\n\n");
	fprintf(stdout, "stream.ty can be '-' for stdin\n");
	fprintf(stdout, "other arguments can be '+' to avoid creating those\n\n");
	fprintf(stdout, "stream.ty is the unprocessed stream\n");
	fprintf(stdout, "program.m2a is the MPEG-1 Layer 2 audio file to be written\n");
	fprintf(stdout, "program.m2v is the MPEG-2 video file to be written\n");
	fprintf(stdout, "badchunk.ty is all the chunks that could not be processed\n");
	/* Report Time-Offset Only */
	/* Added by Stealth Dave */
	fprintf(stdout, "\nsplitstream %s + +\n  Returns only the calculated time offset\n", TIME_OFFSET_STR);
	exit(0);
}

FILE *open_file(char *filename)
{
	FILE * result;

	/* Report Time-Offset Only */
	/* Added by Stealth Dave */
	/*   don't open file if it is the TIME_OFFSET_STR constant */
	if (!strcmp(filename, "+") || !strcmp(filename,TIME_OFFSET_STR))
	{
		result = NULL;
	}
	else
	{
		result = fopen(filename, "wb");
		if(result == NULL)
		{
			fprintf(stderr, "could not open %s\n", filename);
			exit(1);
		}
	}

	return result;

}

/*
 =======================================================================================================================
 =======================================================================================================================
 */
int main(int argc, char *argv[])
{
	/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/
	FILE			*in_fp, *video_out_fp, *audio_out_fp, *bad_chunk_out_fp;
	unsigned char	*buf;
	int				read = 0;
	int				chunk_count = 0;
	/*~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~*/

	if(argc != 5) usage();

	if (!strcmp(argv[1], "-"))
	{
#ifdef WIN32
		setmode(fileno(stdin), _O_BINARY);
#endif
		in_fp = stdin;


	}
	else
	{
		in_fp = fopen(argv[1], "rb");
	}
	if(in_fp == NULL)
	{
		fprintf(stderr, "could not open %s\n", argv[1]);
		exit(1);
	}

	/* Report Time-Offset Only */
	/* Added by Stealth Dave */
	if(!strcmp(argv[2],TIME_OFFSET_STR))
	{
		audio_out_fp = NULL;
		video_out_fp = NULL;
		bad_chunk_out_fp = NULL;
		quit_after_time_calc = 1;
	}
	else
	{
		audio_out_fp = open_file(argv[2]);
		video_out_fp = open_file(argv[3]);
		bad_chunk_out_fp = open_file(argv[4]);
	}

	buf = (unsigned char *) malloc(sizeof(char) * CHUNK_SIZE);

	while(!feof(in_fp))
	{
		read += fread(buf + read, sizeof(char), CHUNK_SIZE - read, in_fp);
		if(read == CHUNK_SIZE)
		{
			parse_chunk(audio_out_fp, video_out_fp, bad_chunk_out_fp, buf, chunk_count++);
			read = 0;
		}
	}
	if (previous_chunk_bad)
	{
		fprintf(stderr, "last bad chunk: %d\n", chunk_count - 1);
	}

	print_last_time();
	fprintf(stderr, "\n");

	if(read > 0) fprintf(stderr, "File not a multiple of 128 KB chunks!\n");

		fclose(in_fp);
	if (audio_out_fp != NULL) fclose(audio_out_fp);
	if (bad_chunk_out_fp != NULL) fclose(bad_chunk_out_fp);
	if (video_out_fp != NULL) fclose(video_out_fp);
	return 0;
}