//-----------------------------------------------------------------------------
//
//	Musepack Demuxer
//
//	Author : Igor Janos
//
//-----------------------------------------------------------------------------
#include "stdafx.h"
#include <math.h>

#define OLD_GAIN_REF		64.82


//-----------------------------------------------------------------------------
//
//	CMPCFile class
//
//-----------------------------------------------------------------------------

CMPCFile::CMPCFile() :
	duration_10mhz(0),
	reader(NULL),
	seek_table(NULL),
	seek_table_position(0)
{
	extradata_max_size = 16*1024;			// way too much..
	extradata_size = 0;
	extradata = (uint8*)malloc(extradata_max_size);
	bits_to_skip = 0;

	gain_album_db = 0;
	gain_album_peak_db = 0;
	gain_title_db = 0;
	gain_title_peak_db = 0;
}

CMPCFile::~CMPCFile()
{
	ClearChapters();
	if (seek_table) { free(seek_table); seek_table = NULL; }
	if (extradata) { free(extradata); extradata = NULL; }
}

void CMPCFile::ClearChapters()
{
	for (int i=0; i<chapters.size(); i++) {
		CMPCChapter	*chap = chapters[i];
		delete chap;
	}
	chapters.clear();
}

int CMPCFile::Open_SV8()
{
	bool		done;
	int			ret;

	// means no seek table
	seek_table_position = 0;
	seek_table_size = 0;

	done = false;

	// init extradata
	extradata[0] = 'M';
	extradata[1] = 'P';
	extradata[2] = 'C';
	extradata[3] = 'K';
	extradata_size = 4;

	// now loop through a few packets
	CMPCPacket		packet;
	do {
		ret = packet.Load(reader);
		if (ret < 0) return -1;			// end of file reached before initialization is done ?

		switch (packet.key) {
		case MPC_KEY('S','H'):	
			ret = ReadStreamHeader(&packet); 
			StoreExtraDataPacket(&packet);
			break;
		case MPC_KEY('R','G'):  
			ret = ReadReplaygain(&packet); 
			StoreExtraDataPacket(&packet);
			break;
		case MPC_KEY('S','O'):	
			ret = ReadSeekOffset(&packet); 
			break;
		case MPC_KEY('E','I'):
			StoreExtraDataPacket(&packet);
			break;

		// audio packet - initialization is over
		case MPC_KEY('A','P'):	done = true; break;
		}

		// if parsing failed we quit
		if (ret < 0) return ret;
	} while (!done);

	// insert a dummy AP packet
	uint8 *out = extradata + extradata_size;
	out[0] = 'A';
	out[1] = 'P';
	out[2] = 0x03;
	extradata_size += 3;

	// if there was a seek table, load it
	if (seek_table_position != 0) {
		int64 first_ap_pos, avail;
		reader->GetPosition(&first_ap_pos, &avail);

		reader->Seek(seek_table_position/8);
		ret = packet.Load(reader);
		ret = ReadSeekTable(&packet); 
		if (ret < 0) return ret;

		// if there is a sequence of chapters it MUST follow the seeking table
		while (true) {
			ret = packet.Load(reader);
			if (ret < 0) break;

			if (packet.key == MPC_KEY('C','T')) {
				CMPCChapter		*chap = new CMPCChapter();
				ret = chap->Load(&packet);
				if (ret < 0) {
					delete chap;
				} else {					
					if (chap->caption == _T("")) {
						chap->caption.Format(_T("Chapter %d"), (int)chapters.size() + 1);
					}
					chapters.push_back(chap);
				}
			} else {
				break;
			}
		}

		// seek back to first AP
		reader->Seek(first_ap_pos);
	}

	// we're at start
	current_sample = 0;
	return 0;
}

int CMPCFile::Open_SV7()
{
	int		ret;
	bool	done;

	// means no seek table
	seek_table_position = 0;
	seek_table_size = 0;

	done = false;

	// init extradata
	extradata[0] = 'M';
	extradata[1] = 'P';
	extradata[2] = '+';
	extradata[3] = stream_version;
	extradata_size = 4;

	uint8		hdr[8*4];
	
	__int64 c, a;
	reader->GetPosition(&c, &a);
	ret = reader->Read(extradata + 4, 6*4);
	if (ret < 0) return -1;			// not enough data
	extradata_size += 6*4;
	reader->Seek(c);
	ret = reader->ReadSwapped(hdr, 6*4);
	Bitstream	b(hdr);

	b.NeedBits32();
	int audio_frames = (b.UGetBits(16)<<16) | b.UGetBits(16);

	b.NeedBits32();
	b.UGetBits(1);				// intensity stereo, should be 0
	b.UGetBits(1);				// mid-side
	b.UGetBits(6);				// max-band
	b.UGetBits(4);				// profile
	b.UGetBits(2);				// link
	int samplerateidx = b.UGetBits(2);			// samplerate

	switch (samplerateidx) {
	case 0x00:	sample_rate = 44100; break;
	case 0x01:	sample_rate = 48000; break;
	case 0x02:	sample_rate = 37800; break;
	case 0x03:	sample_rate = 32000; break;
	}
	channels = 2;
	block_pwr = 0;
	b.UGetBits(16);				// max-level

	b.NeedBits32();
	gain_title_db = ((int16)b.SGetBits(16)) / 100.0;		// title gain
	int title_peak = b.UGetBits(16);
	if (title_peak != 0) {		
		gain_title_peak_db = (title_peak / 32767.0);
	} else {
		gain_title_peak_db = 0;
	}

	b.NeedBits32();
	gain_album_db = ((int16)b.SGetBits(16)) / 100.0;		// album gain
	int album_peak = b.UGetBits(16);
	if (album_peak != 0) {
		gain_album_peak_db = (album_peak / 32767.0);
	} else {
		gain_album_peak_db = 0;
	}


	b.NeedBits32();
	true_gapless = b.UGetBits(1);			// true gapless
	int last_frame = b.UGetBits(11);		// last-frame samples
	audio_block_frames = 1;
	total_samples = (audio_frames * 1152);
	if (true_gapless) {
		total_samples -= (1152 - last_frame);
	} else {
		total_samples -= 481;		// synth delay
	}
	duration_10mhz = (total_samples * 10000000) / sample_rate;

	fast_seeking = b.UGetBits(1);
	b.UGetBits(16);				// 19 zero bits
	b.UGetBits(3);

	b.NeedBits32();
	int ver = b.UGetBits(8);

	//-------------------------------------------------------------------------
	// Seek table
	//-------------------------------------------------------------------------
	seek_pwr = 6;
	if (block_pwr > seek_pwr) seek_pwr = block_pwr;	

	// calculate size as seen in mpc_demux.c
	int64 tmp = 2+total_samples / (1152 << seek_pwr);
	if (tmp < seek_table_size) tmp = seek_table_size;

	// alloc memory for seek table
	seek_table = (uint64*)malloc(tmp * sizeof(uint64));
	if (!seek_table) return -1;

	int64 pos, av;
	reader->GetPosition(&pos, &av);
	seek_table[0] = pos*8 - b.BitsLeft();		// current position in bits
	seek_table_size = 1;

	// loop through file to get complete seeking table
	CMPCPacket	packet;
	int64 seek_frames = 0;
	Seek(0);
	do {
		if (seek_frames*1152 >= total_samples) break;

		// new entry in seeking table
		if (seek_frames == (seek_table_size << seek_pwr)) {
			__int64 c, a;
			reader->GetPosition(&c, &a);
			seek_table[seek_table_size] = (c*8) + bits_to_skip;
			seek_table_size++;
		}

		ret = packet.Load_SV7(reader, bits_to_skip, true);
		if (ret<0) break;
		seek_frames += 1;
	} while (true);

	// seek to begin
	Seek(0);

	header_position = 0;
	bits_to_skip = 0;
	return 0;
}


// I/O for MPC file
int CMPCFile::Open(CMPCReader *reader)
{
	HRESULT		hr;
	int			ret;

	//-------------------------------------------------------------------------
	//
	//	Load APE Tag
	//
	//-------------------------------------------------------------------------

	// try to load the APE Tag
	__int64		size, avail;
	int			tagsize = 0;
	reader->GetSize(&avail, &size);

	tag.Clear();
	if (size >= 32) {
		BYTE		temp[32];

		reader->Seek(size - 32);
		reader->Read(temp, 32);

		ret = tag.ReadFooter(temp, 32, &tagsize);
		if (ret == 0 && tagsize > 0) {
			// seek this many bytes back
			reader->Seek(size - 32 - tagsize);
			
			// also load the footer
			tagsize += 32;
			BYTE		*apetag = (BYTE*)malloc(tagsize);
			if (apetag) {
				reader->Read(apetag, tagsize);

				ret = tag.ReadTag(apetag, tagsize);
				if (ret < 0) {
					tag.Clear();
				}

				free(apetag);
			}
		}
	}

	//-------------------------------------------------------------------------
	//
	//	Load MPC file
	//
	//-------------------------------------------------------------------------

	// keep a local copy of the reader
	this->reader = reader;

	// According to stream specification the first 4 bytes should be 'MPCK'
	uint32	magick;
	reader->Seek(0);
	ret = reader->GetMagick(magick);			if (ret < 0) return ret;
	if (magick == 0x4d50434b) {					// MPCK
		return Open_SV8();
	} else
	if ((magick&0xffffff00) == 0x4d502b00) {		// MP+
		// stream version... only 7 is supported
		stream_version = magick & 0x0f;
		if (stream_version != 7) return -1;
		return Open_SV7();
	}
	return -1;
}

void CMPCFile::StoreExtraDataPacket(CMPCPacket *packet)
{
	// we add current packet to extradata
	uint8	*out = extradata + extradata_size;
	memcpy(out, packet->packet, packet->packet_size);
	extradata_size += packet->packet_size;
}

// parsing packets
int CMPCFile::ReadReplaygain(CMPCPacket *packet)
{
	Bitstream	b(packet->payload);
	b.NeedBits();

	int version = b.UGetBits(8);
	if (version != 1) return 0;		// unsupported RG version. not critical to us...

	int16	val;
	b.NeedBits();	val = b.SGetBits(16);	gain_title_db = val / 256.0;
	b.NeedBits();	val = b.UGetBits(16);	gain_title_peak_db = val;
	b.NeedBits();	val = b.SGetBits(16);	gain_album_db = val / 256.0;
	b.NeedBits();	val = b.UGetBits(16);	gain_album_peak_db = val;

	if (gain_title_db != 0) gain_title_db = OLD_GAIN_REF - gain_title_db;
	if (gain_album_db != 0) gain_album_db = OLD_GAIN_REF - gain_album_db;
	if (gain_title_peak_db != 0) {
		gain_title_peak_db = pow((double)10.0, (double)(gain_title_peak_db / (20*256))) / (double)(1 << 15);
	}
	if (gain_album_peak_db != 0) {
		gain_album_peak_db = pow((double)10.0, (double)(gain_album_peak_db / (20*256))) / (double)(1 << 15);
	}

	return 0;
}

int CMPCFile::ReadSeekOffset(CMPCPacket *packet)
{
	Bitstream	b(packet->payload);

	seek_table_position = b.GetMpcSize() * 8;
	seek_table_position += packet->file_position;

	// success 
	return 0;
}

int CMPCFile::ReadSeekTable(CMPCPacket *packet)
{
	Bitstream	b(packet->payload);

	// calculate size as seen in mpc_demux.c
	int64 tmp = b.GetMpcSize();	b.NeedBits();
	seek_table_size = tmp;
	seek_pwr = block_pwr + b.UGetBits(4);
	tmp = 2+total_samples / (1152 << seek_pwr);
	if (tmp < seek_table_size) tmp = seek_table_size;

	// alloc memory for seek table
	seek_table = (uint64*)malloc(tmp * sizeof(uint64));

	uint64 *table = seek_table;

	tmp = b.GetMpcSize();
	table[0] = (uint32)(tmp*8 + header_position);
	if (seek_table_size == 1) return 0;

	tmp = b.GetMpcSize();
	table[1] = (uint64)(tmp*8 + header_position);
	for (int i=2; i<seek_table_size; i++) {
		int code = b.Get_Golomb(12);
		if (code&1) code = -(code&(-1 << 1));
		code <<= 2;
		table[i] = code + 2*table[i-1] - table[i-2];
	}

	return 0;
}

int CMPCFile::ReadStreamHeader(CMPCPacket *packet)
{
	Bitstream	b(packet->payload);
	b.NeedBits();

	// let's do some reading
	uint32		crc;
	crc = b.UGetBits(16);	b.NeedBits();
	crc <<= 16;
	crc |= b.UGetBits(16);

	// TODO: CRC checking.
	b.NeedBits();
	stream_version = b.UGetBits(8);

	switch (stream_version) {
	case 8:
		{
			//-----------------------------------------------------------------
			//	Stream Version 8
			//-----------------------------------------------------------------
			total_samples = b.GetMpcSize();
			int64	silence = b.GetMpcSize();

			b.NeedBits();
			uint8	freq = b.UGetBits(3);
			switch (freq) {
			case 0:	sample_rate = 44100; break;
			case 1:	sample_rate = 48000; break;
			case 2:	sample_rate = 37800; break;
			case 3: sample_rate = 32000; break;
			default:
				// invalid value
				return -1;				
			}

			// let's calculate duration
			duration_10mhz = (total_samples * 10000000) / sample_rate;

			b.DumpBits(5);					// bands
			channels = b.UGetBits(4)+1;		// channels

			b.NeedBits();
			b.DumpBits(1);					// mid side
			block_pwr = b.UGetBits(3) * 2;
			audio_block_frames = 1 << (block_pwr);

			// store absolute position of stream header
			header_position = packet->file_position - 4*8;
		}
		break;
	default:
		// stream version not supported
		return -2;
	}

	// everything is okay
	return 0;
}

// parsing out packets
int CMPCFile::ReadAudioPacket(CMPCPacket *packet, int64 *cur_sample)
{
	// we just load packets until we face the SE packet. Then we return -1
	int ret;

	// end of file ?
	if (current_sample >= total_samples) return -2;

	do {
		switch (stream_version) {
		case 7:	ret = packet->Load_SV7(reader, bits_to_skip); break;
		case 8: ret = packet->Load(reader); break;
		default: ret = -2; break;
		}
		if (ret < 0) {
			return ret;
		}

		// keep track of samples...
		if (cur_sample) *cur_sample = current_sample;
		if (stream_version == 8) {
			switch (packet->key) {
			case MPC_KEY('A','P'):	
				{
					current_sample += (1152*audio_block_frames);
					return 0;			// we got one
				}
			case MPC_KEY('S','E'):
				{
					return -2;
				}
				break;
			}
		} else {
			current_sample += (1152*audio_block_frames);
			return 0;
		}

		// skip other packets
	} while (1);

	// unexpected...
	return 0;
}

int CMPCFile::Seek(int64 seek_sample)
{
	/*
		This will be a little tricky. I would like to support also 
		less precise seeking with files that don't have seeking table. 
		But if there is one, we	would use that one.

		Original implementation also deals with introductory silence.
		But since I don't know of any dshow filters capable of
		gapless playback we would ignore this. Of course it can
		be implemented later.

		It is absolutely okay to do AP-level seeking. The caller
		knows the exact time it wants to seek to and we will provide
		the nearest AP. If the requested time is in the middle of AP,
		the (current_time - requested_time) expression would be negative
		and decoder will skip samples until positive time values are encountered.
		This way we can do sample-precise seeking.
	*/

	// cannot seek
	if (seek_table == NULL || seek_table_size == 0) return -1;

	int packet_num = seek_sample / (1152*audio_block_frames);

	// we need to seek back a little to avoid artifacts
	if (stream_version == 7) {
		if (packet_num > 32) {
			packet_num -= 32;
		} else {
			packet_num = 0;
		}
	}

	int i = (packet_num >> (seek_pwr - block_pwr));
	if (i >= seek_table_size) {
		i = seek_table_size-1;
	}


	// seek to position
	uint64 pos = seek_table[i] >> 3;
	bits_to_skip = seek_table[i] & 0x07;
	reader->Seek(pos);

	// shift back
	i = i << (seek_pwr - block_pwr);
	current_sample = i*1152*audio_block_frames;

	return 0;
}


//-----------------------------------------------------------------------------
//
//	CMPCPacket
//
//-----------------------------------------------------------------------------

CMPCPacket::CMPCPacket() :
	file_position(0),
	packet(NULL),
	payload(NULL),
	packet_size(0),
	payload_size(0),
	key(0)
{
}

CMPCPacket::~CMPCPacket()
{
	// just let it go
	Release();
}

void CMPCPacket::Release()
{
	if (packet) { free(packet); packet = NULL; }
	payload = NULL;
	packet_size = 0;
	payload_size = 0;
	key = 0;
	file_position = 0;
}

int CMPCPacket::Load_SV7(CMPCReader *reader, int &bits_to_skip, bool only_parse)
{
	uint8	temp[14*1024];
	uint8	outtemp[14*1024];
	__int64	cur, av;
	int		t;
	int		total_bits;

	Release();

	reader->GetPosition(&cur, &av);
	int ret = reader->ReadSwapped(temp, min(sizeof(temp), (av-cur)));
	if (ret < 0) return 0;

	Bitstream	b(temp);
	Bitstream	o(outtemp);
	b.NeedBits32();

	// We load bits from input file and store it into our buffer.
	// There they will always start at byte alignment so they can 
	// be transferred to decoder in a pleasant way.

	// skip bits from previous frame
	if (bits_to_skip > 0) {	b.DumpBits(bits_to_skip); }
	file_position = (cur*8) + bits_to_skip;

	// 20-bit frame length
	b.NeedBits24();
	int	frame_len_bits = b.UGetBits(20);
	o.PutBits(frame_len_bits, 20);

	if (only_parse == false) {
		// load data
		int left_bits = frame_len_bits;
		while (left_bits > 16) {
			b.NeedBits();
			t = b.UGetBits(16);
			o.PutBits(t, 16);
			left_bits -= 16;
		}
		b.NeedBits();
		t = b.UGetBits(left_bits);
		o.PutBits(t, left_bits);

		// zero-align 32-bit word
		total_bits = (20 + frame_len_bits + 31) &~ 31;
		int zero_bits = total_bits - frame_len_bits;
		o.PutBits(0, zero_bits);
		o.WriteBits();
	}

	// bits to skip in the next frame
	__int64 finalpos = file_position + 20 + frame_len_bits;
	bits_to_skip = (finalpos & 7);
	reader->Seek(finalpos >> 3);

	// copy data
	if (only_parse == false) {
		packet_size  = total_bits >> 3;
		payload_size = packet_size;
		packet = (uint8*)malloc(packet_size);
		payload = packet;				// pointer to packet payload

		// copy data
		uint8 *src = outtemp;
		uint8 *dst = packet;
		int left = packet_size;
		while (left > 3) {
			dst[0] = src[3];
			dst[1] = src[2];
			dst[2] = src[1];
			dst[3] = src[0];
			src += 4;
			dst += 4;
			left-= 4;
		}
	}

	return 0;
}

int CMPCPacket::Load(CMPCReader *reader)
{
	uint16	key_val;
	int64	size_val, avail;
	int32	size_len;
	int		ret;
	Release();
	reader->GetPosition(&file_position, &avail);

	// end of stream
	if (file_position >= avail) return -2;
	file_position *= 8;

	ret = reader->GetKey(key_val);					
	if (ret < 0) return ret;

	ret = reader->GetSizeElm(size_val, size_len);	
	if (ret < 0) return ret;

	// if the key is not valid, quit
	if (!reader->KeyValid(key_val)) return -1;
	key = key_val;

	// now load the packet
	packet_size  = size_val;
	payload_size = size_val - 2 - size_len;
	packet = (uint8*)malloc(packet_size);
	payload = packet + 2 + size_len;				// pointer to packet payload

	// roll back the bytes
	reader->Seek(file_position/8);
	ret = reader->Read(packet, packet_size);		
	if (ret < 0) return ret;
	return 0;
}

//-----------------------------------------------------------------------------
//
//	CMPCChapter
//
//-----------------------------------------------------------------------------

CMPCChapter::CMPCChapter()
{
	sample_offset		= 0;
	chapter_gain_db		= 0;
	chapter_peak_db		= 0;
}

CMPCChapter::~CMPCChapter()
{
}

int CMPCChapter::Load(CMPCPacket *packet)
{
	Bitstream	b(packet->payload);
	b.NeedBits();

	sample_offset = b.GetMpcSize();

	int16	val;
	b.NeedBits();	val = b.SGetBits(16);	chapter_gain_db = val / 256.0;
	b.NeedBits();	val = b.UGetBits(16);	chapter_peak_db = val;

	if (chapter_gain_db != 0) chapter_gain_db = OLD_GAIN_REF - chapter_gain_db;
	if (chapter_peak_db != 0) {
		chapter_peak_db = pow((double)10.0, (double)(chapter_peak_db / (20*256))) / (double)(1 << 15);
	}

	BYTE		*buf = b.Position();
	BYTE		*end = packet->payload + packet->payload_size;

	// let's try to load the APE Tag
	APE_Tag		tag;
	int ret = tag.ReadTag(buf, end-buf, true);
	if (ret == 0) {

		CString		title	= tag.FindName(_T("title"));
		CString		track	= tag.FindName(_T("track"));
		
		if (track != _T("")) {
			caption.Format(_T("%s [%s]"), title.GetBuffer(), track.GetBuffer());
		} else {
			caption = title;
		}
		caption.Trim();
	
	} else {
		caption = _T("");
	}

	return 0;
}