#include "funcprotos.h"
#include "quicktime.h"




int quicktime_trak_init(quicktime_trak_t *trak)
{
	quicktime_tkhd_init(&(trak->tkhd));
	quicktime_edts_init(&(trak->edts));
	quicktime_mdia_init(&(trak->mdia));
	return 0;
}

int quicktime_trak_init_video(quicktime_t *file, 
							quicktime_trak_t *trak, 
							int frame_w, 
							int frame_h, 
							float frame_rate,
							char *compressor)
{

	quicktime_tkhd_init_video(file, 
		&(trak->tkhd), 
		frame_w, 
		frame_h);
	quicktime_mdia_init_video(file, 
		&(trak->mdia), 
		frame_w, 
		frame_h, 
		frame_rate, 
		compressor);
	quicktime_edts_init_table(&(trak->edts));

	return 0;
}

int quicktime_trak_init_audio(quicktime_t *file, 
							quicktime_trak_t *trak, 
							int channels, 
							int sample_rate, 
							int bits, 
							char *compressor)
{
	quicktime_mdia_init_audio(file, &(trak->mdia), 
		channels, 
		sample_rate, 
		bits, 
		compressor);
	quicktime_edts_init_table(&(trak->edts));
	return 0;
}

int quicktime_trak_delete(quicktime_trak_t *trak)
{
	quicktime_mdia_delete(&(trak->mdia));
	quicktime_edts_delete(&(trak->edts));
	quicktime_tkhd_delete(&(trak->tkhd));
	return 0;
}


int quicktime_trak_dump(quicktime_trak_t *trak)
{
	printf(" track\n");
	quicktime_tkhd_dump(&(trak->tkhd));
	quicktime_edts_dump(&(trak->edts));
	quicktime_mdia_dump(&(trak->mdia));

	return 0;
}

// Used when reading a file
quicktime_trak_t* quicktime_add_trak(quicktime_t *file)
{
	quicktime_moov_t *moov = &(file->moov);
	if(moov->total_tracks < MAXTRACKS)
	{
		moov->trak[moov->total_tracks] = calloc(1, sizeof(quicktime_trak_t));
		quicktime_trak_init(moov->trak[moov->total_tracks]);
		moov->total_tracks++;
	}
	return moov->trak[moov->total_tracks - 1];
}

int quicktime_delete_trak(quicktime_moov_t *moov)
{
	if(moov->total_tracks)
	{
		moov->total_tracks--;
		quicktime_trak_delete(moov->trak[moov->total_tracks]);
		free(moov->trak[moov->total_tracks]);
	}
	return 0;
}


int quicktime_read_trak(quicktime_t *file, quicktime_trak_t *trak, quicktime_atom_t *trak_atom)
{
	quicktime_atom_t leaf_atom;

	do
	{
		quicktime_atom_read_header(file, &leaf_atom);

//printf("quicktime_read_trak %llx %llx\n", quicktime_position(file), quicktime_ftell(file));
/* mandatory */
		if(quicktime_atom_is(&leaf_atom, "tkhd"))
			{ quicktime_read_tkhd(file, &(trak->tkhd)); }
		else
		if(quicktime_atom_is(&leaf_atom, "mdia"))
			{ quicktime_read_mdia(file, &(trak->mdia), &leaf_atom); }
		else
/* optional */
		if(quicktime_atom_is(&leaf_atom, "clip"))
			{ quicktime_atom_skip(file, &leaf_atom); }
		else
		if(quicktime_atom_is(&leaf_atom, "matt"))
			{ quicktime_atom_skip(file, &leaf_atom); }
		else
		if(quicktime_atom_is(&leaf_atom, "edts"))
			{ quicktime_read_edts(file, &(trak->edts), &leaf_atom); }
		else
		if(quicktime_atom_is(&leaf_atom, "load"))
			{ quicktime_atom_skip(file, &leaf_atom); }
		else
		if(quicktime_atom_is(&leaf_atom, "tref"))
			{ quicktime_atom_skip(file, &leaf_atom); }
		else
		if(quicktime_atom_is(&leaf_atom, "imap"))
			{ quicktime_atom_skip(file, &leaf_atom); }
		else
		if(quicktime_atom_is(&leaf_atom, "udta"))
			{ quicktime_atom_skip(file, &leaf_atom); }
		else
			quicktime_atom_skip(file, &leaf_atom);
//printf("quicktime_read_trak %llx %llx\n", quicktime_position(file), leaf_atom.end);
	}while(quicktime_position(file) < trak_atom->end);

	return 0;
}

int quicktime_write_trak(quicktime_t *file, 
	quicktime_trak_t *trak, 
	long moov_time_scale)
{
	long duration;
	long timescale;
	quicktime_atom_t atom;
	quicktime_atom_write_header(file, &atom, "trak");
	quicktime_trak_duration(trak, &duration, &timescale);

/*printf("quicktime_write_trak duration %d\n", duration); */
/* get duration in movie's units */
	trak->tkhd.duration = (long)((float)duration / timescale * moov_time_scale);
	trak->mdia.mdhd.duration = duration;
	trak->mdia.mdhd.time_scale = timescale;

	quicktime_write_tkhd(file, &(trak->tkhd));
	quicktime_write_edts(file, &(trak->edts), trak->tkhd.duration);
	quicktime_write_mdia(file, &(trak->mdia));

	quicktime_atom_write_footer(file, &atom);

	return 0;
}

int64_t quicktime_track_end(quicktime_trak_t *trak)
{
/* get the byte endpoint of the track in the file */
	int64_t size = 0;
	int64_t chunk, chunk_offset, chunk_samples;
	quicktime_stsz_t *stsz = &(trak->mdia.minf.stbl.stsz);
	quicktime_stsc_t *stsc = &(trak->mdia.minf.stbl.stsc);
	quicktime_stco_t *stco;

/* get the last chunk offset */
/* the chunk offsets contain the HEADER_LENGTH themselves */
	stco = &(trak->mdia.minf.stbl.stco);
	chunk = stco->total_entries;
	size = chunk_offset = stco->table[chunk - 1].offset;

/* get the number of samples in the last chunk */
	chunk_samples = stsc->table[stsc->total_entries - 1].samples;

/* get the size of last samples */
	if(stsz->sample_size) {
		int sample_bits = trak->mdia.minf.stbl.stsd.table[0].sample_size;
		int channels = trak->mdia.minf.stbl.stsd.table[0].channels;
/* assume audio so calculate the sample size */
		size += chunk_samples * stsz->sample_size * channels * sample_bits / 8;
	}
	else {
/* assume video */
		int64_t i = stsz->total_entries;
		while( --i >= chunk_samples ) size += stsz->table[i].size;
	}

	return size;
}

long quicktime_track_samples(quicktime_t *file, quicktime_trak_t *trak)
{
/*printf("file->rd %d file->wr %d\n", file->rd, file->wr); */
	if(file->wr)
	{
/* get the sample count when creating a new file */
 		quicktime_stsc_table_t *table = trak->mdia.minf.stbl.stsc.table;
		long total_entries = trak->mdia.minf.stbl.stsc.total_entries;
		long chunk = trak->mdia.minf.stbl.stco.total_entries;
		long sample;

		if(chunk)
		{
			sample = quicktime_sample_of_chunk(trak, chunk);
			sample += table[total_entries - 1].samples;
		}
		else 
			sample = 0;

		return sample;
	}
	else
	{
/* get the sample count when reading only */
		quicktime_stts_t *stts = &(trak->mdia.minf.stbl.stts);
		int64_t total = 0;
		int i;

		if(trak->mdia.minf.is_audio)
		{

// Get total sample duration
			for(i = 0; i < stts->total_entries; i++)
			{
				total += stts->table[i].sample_count *
					stts->table[i].sample_duration;
			}

			return total;
		}
		else
		if(trak->mdia.minf.is_video)
		{
/* Get total number of samples */
			for(i = 0; i < stts->total_entries; i++)
			{
				total += stts->table[i].sample_count;
			}
			return total;
		}
		return total;
	}
}

long quicktime_sample_of_chunk(quicktime_trak_t *trak, long chunk)
{
	quicktime_stsc_table_t *table = trak->mdia.minf.stbl.stsc.table;
	long total_entries = trak->mdia.minf.stbl.stsc.total_entries;
	long chunk1entry, chunk2entry;
	long chunk1, chunk2, chunks, total = 0;

	for(chunk1entry = total_entries - 1, chunk2entry = total_entries; 
		chunk1entry >= 0; 
		chunk1entry--, chunk2entry--)
	{
		chunk1 = table[chunk1entry].chunk;

		if(chunk > chunk1)
		{
			if(chunk2entry < total_entries)
			{
				chunk2 = table[chunk2entry].chunk;

				if(chunk < chunk2) chunk2 = chunk;
			}
			else
				chunk2 = chunk;

			chunks = chunk2 - chunk1;

			total += chunks * table[chunk1entry].samples;
		}
	}

	return total;
}

// For AVI
int quicktime_avg_chunk_samples(quicktime_t *file, quicktime_trak_t *trak)
{
	int chunk = trak->mdia.minf.stbl.stco.total_entries - 1;
	long total_samples;

	if(chunk >= 0)
	{
		total_samples = quicktime_sample_of_chunk(trak, chunk);
		return total_samples / (chunk + 1);
	}
	else
	{
		total_samples = quicktime_track_samples(file, trak);
		return total_samples;
	}
}

int quicktime_chunk_of_sample(int64_t *chunk_sample, int64_t *chunk, 
				quicktime_trak_t *trak, int64_t sample)
{
	quicktime_stsc_table_t *table = trak->mdia.minf.stbl.stsc.table;
	int64_t total_entries = trak->mdia.minf.stbl.stsc.total_entries;
	if(!total_entries) {
		*chunk_sample = 0;
		*chunk = 0;
		return 0;
	}

	int64_t last_chunk = 1, last_samples = 0;
	int64_t chunk_samples = 0, i = 0;
	while( i < total_entries ) {
		int64_t next_chunk = table[i].chunk;
		int64_t chunks = next_chunk - last_chunk;
		int64_t samples = chunks * last_samples;
		int64_t end_sample = chunk_samples + samples;
		if( sample < end_sample ) break;
		chunk_samples = end_sample;
		last_chunk = next_chunk;
		last_samples = table[i++].samples;
	}

	*chunk = !last_samples ? 1 : (sample - chunk_samples) / last_samples + last_chunk;
	*chunk_sample = chunk_samples + (*chunk - last_chunk) * last_samples;
	return 0;
}

int64_t quicktime_chunk_to_offset(quicktime_t *file, quicktime_trak_t *trak, long chunk)
{
	quicktime_stco_table_t *table = trak->mdia.minf.stbl.stco.table;
	int64_t table_entries = trak->mdia.minf.stbl.stco.total_entries;
	if( chunk >= table_entries ) chunk = table_entries-1;
	int64_t result = chunk < 0 ? (int64_t)HEADER_LENGTH * 2 : table[chunk-1].offset;
// Skip chunk header for AVI.  Skip it here instead of in read_chunk because some
// codecs can't use read_chunk
	if(file->use_avi) {
//printf("quicktime_chunk_to_offset 1 %llx %llx\n", result, file->mdat.atom.start);
		result += 8 + file->mdat.atom.start;
	}
	return result;
}

long quicktime_offset_to_chunk(int64_t *chunk_offset, 
	quicktime_trak_t *trak, int64_t offset)
{
	quicktime_stco_table_t *table = trak->mdia.minf.stbl.stco.table;
	int64_t i = trak->mdia.minf.stbl.stco.total_entries;
	while( --i >= 0 && table[i].offset > offset );
	*chunk_offset = i>=0 ? table[i].offset : HEADER_LENGTH * 2;
	return i>=0 ? i+1 : 1;
}

int quicktime_chunk_bytes(quicktime_t *file, int64_t *chunk_offset,
				int chunk, quicktime_trak_t *trak)
{
	*chunk_offset = quicktime_chunk_to_offset(file, trak, chunk);
	quicktime_set_position(file, *chunk_offset - 4);
	int result = quicktime_read_int32_le(file);
	return result;
}

int64_t quicktime_sample_range_size(quicktime_trak_t *trak, int64_t chunk_sample, int64_t sample)
{
	quicktime_stsz_table_t *table = trak->mdia.minf.stbl.stsz.table;
	int64_t total_entries = trak->mdia.minf.stbl.stsz.total_entries;
	if( sample >= total_entries ) return 0;
	int64_t total = 0, i = chunk_sample;
	while( i < sample ) total += table[i++].size;
	return total;
}

int64_t quicktime_sample_to_offset(quicktime_t *file, 
	quicktime_trak_t *trak, int64_t sample)
{
	int64_t chunk, chunk_sample;
	quicktime_chunk_of_sample(&chunk_sample, &chunk, trak, sample);
	int64_t offset = quicktime_chunk_to_offset(file, trak, chunk);
	int sample_sz = trak->mdia.minf.stbl.stsz.sample_size;
/* if all frames have the same size */
	if( sample_sz ) {
		int sample_bits = trak->mdia.minf.stbl.stsd.table[0].sample_size;
		int frame_sz = sample_bits * trak->mdia.minf.stbl.stsd.table[0].channels;
// frame_sz is in bits, so we divide by 8 at the end
		offset += (sample - chunk_sample) * frame_sz / 8;
	}
	else
		offset += quicktime_sample_range_size(trak, chunk_sample, sample);
	return offset;
}

long quicktime_offset_to_sample(quicktime_trak_t *trak, int64_t offset)
{
	int64_t sample_offset;
	int64_t chunk = quicktime_offset_to_chunk(&sample_offset, trak, offset);
	int64_t sample = quicktime_sample_of_chunk(trak, chunk);
	quicktime_stsz_table_t *table = trak->mdia.minf.stbl.stsz.table;

	int sample_sz = trak->mdia.minf.stbl.stsz.sample_size;
	if( !sample_sz ) {
		int64_t total_entries1 = trak->mdia.minf.stbl.stsz.total_entries - 1;
		while( sample_offset < offset && sample < total_entries1 )
			sample_offset += table[sample++].size;
	}
	else
		sample += (offset - sample_offset) / sample_sz;
	
	return sample;
}

void quicktime_write_chunk_header(quicktime_t *file, 
	quicktime_trak_t *trak, 
	quicktime_atom_t *chunk)
{
	if(file->use_avi)
	{
/* Get tag from first riff strl */
		quicktime_riff_t *first_riff = file->riff[0];
		quicktime_hdrl_t *hdrl = &first_riff->hdrl;
		quicktime_strl_t *strl = hdrl->strl[trak->tkhd.track_id - 1];
		char *tag = strl->tag;

/* Create new RIFF object at 1 Gig mark */
		quicktime_riff_t *riff = file->riff[file->total_riffs - 1];
		if(quicktime_position(file) - riff->atom.start > 0x40000000)
		{
			quicktime_finalize_riff(file, riff);
			quicktime_init_riff(file);
		}

		

/* Write AVI header */
		quicktime_atom_write_header(file, chunk, tag);
	}
	else
	{
		chunk->start = quicktime_position(file);
	}
}

void quicktime_write_chunk_footer(quicktime_t *file, 
	quicktime_trak_t *trak,
	int current_chunk,
	quicktime_atom_t *chunk, 
	int samples)
{
	int64_t offset = chunk->start;
	int sample_size = quicktime_position(file) - offset;


// Write AVI footer
	if(file->use_avi)
	{
		quicktime_atom_write_footer(file, chunk);

// Save version 1 index entry for first RIFF only
		if(file->total_riffs < 2)
		{
			quicktime_update_idx1table(file, 
				trak, 
				offset, 
				sample_size);
		}

// Save partial index entry
		quicktime_update_ixtable(file, 
			trak, 
			offset, 
			sample_size);
	}

	if(offset + sample_size > file->mdat.atom.size)
		file->mdat.atom.size = offset + sample_size;

	quicktime_update_stco(&(trak->mdia.minf.stbl.stco), 
		current_chunk, 
		offset);

	if(trak->mdia.minf.is_video)
		quicktime_update_stsz(&(trak->mdia.minf.stbl.stsz), 
		current_chunk - 1, 
		sample_size);

	quicktime_update_stsc(&(trak->mdia.minf.stbl.stsc), 
		current_chunk, 
		samples);
}

int quicktime_write_vbr_frame(quicktime_t *file, int track,
	char *data, int data_size, int samples)
{
	quicktime_audio_map_t *track_map = &(file->atracks[track]);
	quicktime_trak_t *trak = track_map->track;
	quicktime_atom_t chunk_atom;
	int result = 0;
	if( !samples ) return result;

	quicktime_write_chunk_header(file, trak, &chunk_atom);
	result = !quicktime_write_data(file, data, data_size);
	int64_t offset = chunk_atom.start;

// AVI case
	if(file->use_avi)
	{
		quicktime_atom_write_footer(file, &chunk_atom);
// Save version 1 index entry for first RIFF only
		if(file->total_riffs < 2)
		{
			quicktime_update_idx1table(file, 
				trak, 
				offset, 
				data_size);
		}

// Save version 2 index entry
		quicktime_update_ixtable(file, 
			trak, 
			offset, 
			data_size);
	}

// Update MDAT size
	if(offset + data_size > file->mdat.atom.size)
		file->mdat.atom.size = offset + data_size;

// Update time to sample table
	quicktime_stts_append_audio(file, 
		&(trak->mdia.minf.stbl.stts), 
		samples);

	int64_t total_chunks = quicktime_stts_total_samples(file,
		&(trak->mdia.minf.stbl.stts));
	quicktime_update_stco(&(trak->mdia.minf.stbl.stco), 
		total_chunks, 
		offset);
	quicktime_update_stsc(&(trak->mdia.minf.stbl.stsc), 
		total_chunks, 
		1);
	quicktime_update_stsz(&(trak->mdia.minf.stbl.stsz), 
		total_chunks - 1, 
		data_size);


	return result;
}



/**
 * This is used for writing the header
 **/
int quicktime_trak_duration(quicktime_trak_t *trak, 
	long *duration, 
	long *timescale)
{
	quicktime_stts_t *stts = &(trak->mdia.minf.stbl.stts);
	int i;
	*duration = 0;

	for(i = 0; i < stts->total_entries; i++)
	{
		*duration += stts->table[i].sample_duration * 
			stts->table[i].sample_count;
	}

	*timescale = trak->mdia.mdhd.time_scale;
	return 0;
}

int quicktime_trak_fix_counts(quicktime_t *file, quicktime_trak_t *trak)
{
	if(!trak->mdia.minf.stbl.stts.is_vbr) {
		long samples = quicktime_track_samples(file, trak);
		trak->mdia.minf.stbl.stts.table[0].sample_count = samples;

		if(!trak->mdia.minf.stbl.stsz.total_entries) {
			trak->mdia.minf.stbl.stsz.sample_size = 1;
			trak->mdia.minf.stbl.stsz.total_entries = samples;
		}
	}
	return 0;
}

long quicktime_chunk_samples(quicktime_trak_t *trak, long chunk)
{
	quicktime_stsc_t *stsc = &(trak->mdia.minf.stbl.stsc);
	long i = stsc->total_entries;
	while( --i>=0 && stsc->table[i].chunk > chunk );
	return i < 0 ? 0 : stsc->table[i].samples;
}

int quicktime_trak_shift_offsets(quicktime_trak_t *trak, int64_t offset)
{
	quicktime_stco_t *stco = &(trak->mdia.minf.stbl.stco);
	int i;

	for(i = 0; i < stco->total_entries; i++)
	{
		stco->table[i].offset += offset;
	}
	return 0;
}

char* quicktime_compressor(quicktime_trak_t *track)
{
	return track->mdia.minf.stbl.stsd.table[0].format;
}


int quicktime_sample_duration(quicktime_trak_t *trak)
{
	quicktime_stts_t *stts = &trak->mdia.minf.stbl.stts;
	int i;
	int max_count = 0;
	int result = 1;
	for(i = 0; i < stts->total_entries; i++)
	{
		quicktime_stts_table_t *table = &stts->table[i];
		if(table->sample_count > max_count)
		{
			max_count = table->sample_count;
			result = table->sample_duration;
		}
	}
	return result;
}