//-----------------------------------------------------------------------------
//
//	Musepack Decoder
//
//	Author : Igor Janos
//
//-----------------------------------------------------------------------------
#include "stdafx.h"


//-----------------------------------------------------------------------------
//
//	CMPCDecoder class
//
//-----------------------------------------------------------------------------

CUnknown *CMPCDecoder::CreateInstance(LPUNKNOWN pUnk, HRESULT *phr)
{
	if (phr) *phr = NOERROR;

	return new CMPCDecoder(pUnk, phr);
}

CMPCDecoder::CMPCDecoder(LPUNKNOWN pUnk, HRESULT *phr) :
	CTransformFilter(_T("MPC Decoder"), pUnk, CLSID_MusepackDecoder),
	demux(NULL),
	decoder_specific(NULL),
	decoder_specific_size(0)
{
	if (phr) *phr = NOERROR;

	m_pInput = new CTransformInputPin(NAME("Input"), this, phr, L"Input");
	m_pOutput = new CTransformOutputPin(NAME("Output"), this, phr, L"Output");

	bitrate_decay = 0.9;
	bitrate_cnt = 1;
	bitrate_blur = 0;
	frames_decoded = 0;
}

CMPCDecoder::~CMPCDecoder()
{
}

STDMETHODIMP CMPCDecoder::NonDelegatingQueryInterface(REFIID riid,void **ppv)
{
    CheckPointer(ppv,E_POINTER);

	if (riid == IID_ISpecifyPropertyPages) {
		return GetInterface((ISpecifyPropertyPages*)this, ppv);
	} else
	if (riid == IID_IMusepackDecoder) {
		return GetInterface((IMusepackDecoder*)this, ppv);
	} else
		return CBaseFilter::NonDelegatingQueryInterface(riid,ppv);
}

STDMETHODIMP CMPCDecoder::GetPages(CAUUID *pPages)
{
    CheckPointer(pPages,E_POINTER);

    pPages->cElems = 1;
    pPages->pElems = (GUID *) CoTaskMemAlloc(sizeof(GUID));
    if (pPages->pElems == NULL) {
        return E_OUTOFMEMORY;
    }

	*(pPages->pElems) = CLSID_MusepackDecoderPage;
    return NOERROR;

} // GetPages

HRESULT CMPCDecoder::CheckInputType(const CMediaType *mtIn)
{
	// we only accept these
	if (mtIn->majortype != MEDIATYPE_Audio) return E_FAIL;
	if (mtIn->subtype != MEDIASUBTYPE_MusepackPacket) return E_FAIL;
	if (mtIn->formattype != FORMAT_WaveFormatEx) return E_FAIL;

	return NOERROR;
}

HRESULT CMPCDecoder::CheckTransform(const CMediaType *mtIn, const CMediaType *mtOut)
{
	HRESULT hr = CheckInputType(mtIn);
	if (FAILED(hr)) return hr;

	// output is fixed
	if (mtOut->majortype != MEDIATYPE_Audio) return E_FAIL;
	if (mtOut->subtype != MEDIASUBTYPE_PCM) return E_FAIL;

	return NOERROR;
}

HRESULT CMPCDecoder::SetMediaType(PIN_DIRECTION direction, const CMediaType *pmt)
{
	if (direction == PINDIR_INPUT) {
		if (pmt->formattype != FORMAT_WaveFormatEx) return E_FAIL;

		// keep a local copy so our decoder specific info doesn't go away
		mtIn = *pmt;

		WAVEFORMATEX *wfx = (WAVEFORMATEX*)mtIn.pbFormat;

		// we assemble output format
		wfOut.cbSize = 0;
		wfOut.nChannels = wfx->nChannels;
		wfOut.nSamplesPerSec = wfx->nSamplesPerSec;
		wfOut.wFormatTag = WAVE_FORMAT_PCM;
		wfOut.wBitsPerSample = 16;				// currently only 16-bit

		// get hold of decoder specific info
		decoder_specific = ((BYTE*)wfx) + sizeof(WAVEFORMATEX);
		decoder_specific_size = wfx->cbSize;

		mpc_demux	*dmx;

		int ret = OpenDemux(&dmx, &stream_info);
		if (ret < 0) {
			mpc_demux_exit(dmx);
			return E_FAIL;
		}

		// if the stream_info says something different...
		if (stream_info.sample_freq != wfx->nSamplesPerSec ||
			stream_info.channels != wfx->nChannels
			) {

			// update values
			wfOut.nChannels = stream_info.channels;
			wfOut.nSamplesPerSec = stream_info.sample_freq;
		}

		wfOut.nBlockAlign = (wfOut.wBitsPerSample * wfOut.nChannels) >> 3;
		wfOut.nAvgBytesPerSec = wfOut.nBlockAlign * wfOut.nSamplesPerSec;

		// done with the demuxer
		mpc_demux_exit(dmx);
	}
	return NOERROR;
}

HRESULT CMPCDecoder::BreakConnect(PIN_DIRECTION dir)
{
	if (dir == PINDIR_INPUT) {
		memset(&wfOut, 0, sizeof(WAVEFORMATEX));

		// disconnect output
		if (m_pOutput->IsConnected()) {
			m_pOutput->GetConnected()->Disconnect();
			m_pOutput->Disconnect();
		}

		// cancel out decoder specific info
		decoder_specific = NULL;
		decoder_specific_size = 0;
	}
	return __super::BreakConnect(dir);
}

HRESULT CMPCDecoder::DecideBufferSize(IMemAllocator *pAlloc, ALLOCATOR_PROPERTIES *pProp)
{
	// since the frames are of fixed size we can be sure
	// that the size won't be larger than....
	pProp->cbBuffer = (wfOut.nBlockAlign * 1152) * 1.5;		// + some extra space
	pProp->cBuffers = 3;

	ALLOCATOR_PROPERTIES	act;
	HRESULT hr = pAlloc->SetProperties(pProp, &act);
	if (FAILED(hr)) return hr;

	if (act.cbBuffer < pProp->cbBuffer) return E_FAIL;
	return NOERROR;
}

HRESULT CMPCDecoder::GetMediaType(int iPosition, CMediaType *pmt)
{
	if (iPosition < 0) return E_INVALIDARG;
	if (iPosition > 0) return VFW_S_NO_MORE_ITEMS;

	// don't accept connections until input pin is connected
	if (m_pInput->IsConnected() == FALSE) return E_FAIL;

	pmt->majortype  = MEDIATYPE_Audio;
	pmt->subtype    = MEDIASUBTYPE_PCM;
	pmt->formattype	= FORMAT_WaveFormatEx;

	// we already know...
	WAVEFORMATEX	*pwf = (WAVEFORMATEX*)pmt->AllocFormatBuffer(sizeof(WAVEFORMATEX));
	memcpy(pwf, &wfOut, sizeof(WAVEFORMATEX));

	return NOERROR;
}

HRESULT CMPCDecoder::GetDeliveryBuffer(IMediaSample **sample)
{
    IMediaSample *pOutSample;
    HRESULT hr = m_pOutput->GetDeliveryBuffer(&pOutSample, NULL, NULL, 0);
    *sample = pOutSample;
    if (FAILED(hr)) {
        return hr;
    }

	// check for type change
	AM_MEDIA_TYPE *mt;
	if (pOutSample->GetMediaType(&mt) == NOERROR) {
		CMediaType _mt(*mt);
		SetMediaType(PINDIR_OUTPUT, &_mt);
		DeleteMediaType(mt);
	}

	return NOERROR;
}

HRESULT CMPCDecoder::StartStreaming()
{
	HRESULT hr = __super::StartStreaming();
	if (FAILED(hr)) return hr;

	ASSERT(!demux);

	// now let's open the demuxer
	int ret = OpenDemux(&demux, &stream_info);
	if (ret < 0) return E_FAIL;

	bitrate_decay = 0.9;
	bitrate_cnt = 1;
	bitrate_blur = 0;
	frames_decoded = 0;

	return hr;
}

HRESULT CMPCDecoder::StopStreaming()
{
	HRESULT hr = __super::StopStreaming();
	if (FAILED(hr)) return hr;

	// done with demuxer
	mpc_demux_exit(demux);
	demux = NULL;

	bitrate_decay = 0.9;
	bitrate_cnt = 1;
	bitrate_blur = 0;
	frames_decoded = 0;

	return hr;
}

#define FLOAT_CLIP(x)		(x < -1.0 ? -1.0 : x > 1.0 ? 1.0 : x)

static void convert_float_to_int16_c(float *src, short *dst, int count)
{
	count -= 8;
	while (count > 0) {
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		count -= 8;
	}
	count += 8;

	// last few must be done one-by-one
	while (count > 0) {
		*dst++ = (short)(FLOAT_CLIP(*src) * 32767);	src++;
		count --;
	}
}

HRESULT CMPCDecoder::Receive(IMediaSample *pSample)
{
	if (!m_pOutput->IsConnected()) {
		return E_FAIL;
	}

	// we don't accept data while flushing
	if (m_pInput->IsFlushing()) {
		return E_UNEXPECTED;
	}


	// we receive whole AP packets so just decode the frames :D
	BYTE *pdata = NULL;
	long lsize  = 0;
	pSample->GetPointer(&pdata);
	lsize = pSample->GetActualDataLength();

	// set the reader
	reader.SetData(pdata, lsize);
	if (stream_info.stream_version == 8) reader.AppendSE();

	// TODO: use some more clever way of doing this.
	MPC_SAMPLE_FORMAT	samples[1152 * 8];
	int					dec_samples;

	// handling timestamps
	REFERENCE_TIME		rtStart, rtStop;
	__int64				current_sample = 0;
	__int64				samples_to_skip = 0;
	bool				has_timestamp = false;

	if (pSample->GetTime(&rtStart, &rtStop) == NOERROR) {
		has_timestamp = true;
		// calculate relative sample_time
		current_sample = (rtStart * stream_info.sample_freq) / 10000000;

		if (current_sample < 0) {
			samples_to_skip = -current_sample;
		}
	}

	// manually clear the buffer for demuxer

	demux->bytes_total = 0; 
	demux->bits_reader.buff = demux->buffer;
	demux->bits_reader.count = 8;
	demux->block_bits = 0;
	demux->block_frames = 0;
	demux->d->decoded_samples = 0;
	demux->d->samples_to_skip = 0;

	int	total_samples_decoded = 0;

	// discontinuity after seek - reset scale factors
	if (pSample->IsDiscontinuity() == NOERROR) {
		memset(demux->d->SCF_Index_L, 0, sizeof(demux->d->SCF_Index_L));
		memset(demux->d->SCF_Index_R, 0, sizeof(demux->d->SCF_Index_R));
	}


	// now scan through the frames
	int frame_cnt = 1 << stream_info.block_pwr;
	for (int i=0; i<frame_cnt; i++) {

		//---------------------------------------------------------------------
		//	Decoding Part
		//---------------------------------------------------------------------

		mpc_frame_info	frame;
		mpc_status		err;

		memset(&frame, 0, sizeof(frame));
		frame.buffer = (MPC_SAMPLE_FORMAT*)samples;

		// decode one frame
		err = mpc_demux_decode(demux, &frame);
		if (err != MPC_STATUS_OK) {
			break;
		}
		if (frame.bits == -1) {
			break;				// done.
		}

		// frame counter
		frames_decoded++;


		//---------------------------------------------------------------------
		//	Skipping samples...
		//---------------------------------------------------------------------

		MPC_SAMPLE_FORMAT	*src_samples = frame.buffer;
		int					decoded_time_samples = frame.samples;
		total_samples_decoded += frame.samples;

		if (has_timestamp) {
			if (samples_to_skip > 0) {

				// skip sample counters
				__int64 to_skip = min(frame.samples, samples_to_skip);
				decoded_time_samples -= to_skip;
				samples_to_skip		 -= to_skip;
				current_sample	 	 += to_skip;
				src_samples			 += (to_skip * stream_info.channels);
			}
		}

		// if there's nothing left, continue
		if (decoded_time_samples <= 0) continue;

		// total number of decoded samples
		dec_samples = decoded_time_samples * stream_info.channels;

		//---------------------------------------------------------------------
		//	Convert and deliver
		//---------------------------------------------------------------------

		// deliver decoded data
		IMediaSample	*sample;
		HRESULT			hr;
		hr = GetDeliveryBuffer(&sample);
		if (FAILED(hr)) break;				// perhaps allocator is decommited

		BYTE *outp;
		sample->GetPointer(&outp);

		// this should definitely be true
		ASSERT(sample->GetSize() >= (dec_samples * sizeof(short)));

		// do format conversion
#ifdef MPC_FIXED_POINT
		// TODO: find ranges in headers somewhere
#else
		convert_float_to_int16_c((float*)src_samples, (short*)outp, dec_samples);
		sample->SetActualDataLength(dec_samples * sizeof(short));
#endif

		// IMediaSample properties
		if (has_timestamp) {
			REFERENCE_TIME	tstart, tstop;
			tstart = (current_sample * 10000000) / stream_info.sample_freq;
			tstop  = ((current_sample+decoded_time_samples) * 10000000) / stream_info.sample_freq;
			sample->SetTime(&tstart, &tstop);
		}
		sample->SetSyncPoint(TRUE);
		sample->SetDiscontinuity(FALSE);

		// advance time
		current_sample += decoded_time_samples;

		hr = m_pOutput->Deliver(sample);
		sample->Release();

		// something went wrong. stop decoding.
		if (FAILED(hr)) return hr;
	}

	// update bitrate stuff
	bitrate_blur *= bitrate_decay;
	bitrate_cnt  *= bitrate_decay;
	bitrate_blur += (lsize*8);
	bitrate_cnt  += total_samples_decoded;

	return NOERROR;
}

int CMPCDecoder::OpenDemux(mpc_demux **dmx, mpc_streaminfo *si)
{
	// some data needed
	if (!decoder_specific || decoder_specific_size <= 0) return -1;

	reader.SetData(decoder_specific, decoder_specific_size);
	(*dmx) = mpc_demux_init(&reader.rd);
	if (!(*dmx)) return -1;

	mpc_demux_get_info((*dmx), si);
	return 0;
}

STDMETHODIMP CMPCDecoder::GetStreamVersion(int *version)
{
	if (m_pInput->IsConnected()) {
		if (version) *version = stream_info.stream_version;
		return NOERROR;
	} else {
		if (version) *version = 0;
		return E_FAIL;
	}
}

STDMETHODIMP CMPCDecoder::GetSamplerate(int *samplerate)
{
	if (m_pInput->IsConnected()) {
		if (samplerate) *samplerate = stream_info.sample_freq;
		return NOERROR;
	} else {
		if (samplerate) *samplerate = 0;
		return E_FAIL;
	}
}

STDMETHODIMP CMPCDecoder::GetChannels(int *channels)
{
	if (m_pInput->IsConnected()) {
		if (channels) *channels = stream_info.channels;
		return NOERROR;
	} else {
		if (channels) *channels = 0;
		return E_FAIL;
	}
}

STDMETHODIMP CMPCDecoder::GetReplaygain(float *gain)
{
	if (m_pInput->IsConnected()) {
		float g = 0;
		if (stream_info.gain_title != 0) {
			g = (64.82 - stream_info.gain_title / 256.0 );
		}
		if (stream_info.gain_album != 0) {
			g = (64.82 - stream_info.gain_album / 256.0 );
		}

		if (gain) *gain = g;
		return NOERROR;
	} else {
		if (gain) *gain = 0;
		return E_FAIL;
	}
}

STDMETHODIMP CMPCDecoder::GetBitrate(float *bitrate)
{
	if (bitrate_cnt > 0 && m_pInput->IsConnected()) {

		if (bitrate) *bitrate = (stream_info.sample_freq * bitrate_blur / bitrate_cnt);
	} else {
		if (bitrate) *bitrate = 0;
	}
	return NOERROR;
}

STDMETHODIMP CMPCDecoder::GetFramesDecoded(__int64 *frames)
{
	if (frames) *frames = frames_decoded;
	return NOERROR;
}



//-----------------------------------------------------------------------------
//
//	CMPCReader class
//
//-----------------------------------------------------------------------------

static mpc_int32_t mpc_read_proc(mpc_reader *p_reader, void *ptr, mpc_int32_t size)
{
	CMPCReader *instance = (CMPCReader*)p_reader->data;
	ASSERT(instance);
	return instance->Read(ptr, size);
}

static mpc_bool_t mpc_seek_proc(mpc_reader *p_reader, mpc_int32_t offset)
{
	CMPCReader *instance = (CMPCReader*)p_reader->data;
	ASSERT(instance);
	return instance->Seek(offset);
}

static mpc_int32_t mpc_tell_proc(mpc_reader *p_reader)
{
	CMPCReader *instance = (CMPCReader*)p_reader->data;
	ASSERT(instance);
	return instance->Tell();
}

static mpc_int32_t mpc_get_size_proc(mpc_reader *p_reader)
{
	CMPCReader *instance = (CMPCReader*)p_reader->data;
	ASSERT(instance);
	return instance->GetSize();
}

static mpc_bool_t mpc_canseek_proc(mpc_reader *p_reader)
{
	CMPCReader *instance = (CMPCReader*)p_reader->data;
	ASSERT(instance);
	return instance->CanSeek();
}

CMPCReader::CMPCReader()
{
	buf = (BYTE*)malloc(256*1024);		// 256 KB should be enough
	size = 0;

	rd.data = (void*)this;
	rd.read = mpc_read_proc;
	rd.seek = mpc_seek_proc;
	rd.tell = mpc_tell_proc;
	rd.get_size = mpc_get_size_proc;
	rd.canseek = mpc_canseek_proc;
}

CMPCReader::~CMPCReader()
{
	free(buf);
}

int CMPCReader::Read(void *buf, int size)
{
	// return the number of bytes read
	int toread_max = min(this->size - this->pos, size);
	memcpy(buf, this->buf + this->pos, toread_max);
	this->pos += toread_max;
	return toread_max;
}

bool CMPCReader::Seek(int offset)
{
	// sorry no seeking here
	return false;
}

int CMPCReader::Tell()
{
	// that's where we are
	return pos;
}

int CMPCReader::GetSize()
{
	// size of our supplied buffer
	return size;
}

bool CMPCReader::CanSeek()
{
	// no. we are not seekable. upstream demuxer is.
	return false;
}

void CMPCReader::SetData(BYTE *data, int size)
{
	ASSERT(size < 256*1024);
	memcpy(buf, data, size);

	this->size = size;
	this->pos = 0;
}

void CMPCReader::AppendSE()
{
	BYTE *out = buf + size;
	out[0] = 'S';
	out[1] = 'E';
	out[2] = 0x03;
	size += 3;
}








