/* 
 *	Copyright (C) 2003-2006 Gabest
 *	http://www.gabest.org
 *
 *  This Program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2, or (at your option)
 *  any later version.
 *   
 *  This Program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 *  GNU General Public License for more details.
 *   
 *  You should have received a copy of the GNU General Public License
 *  along with GNU Make; see the file COPYING.  If not, write to
 *  the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA. 
 *  http://www.gnu.org/copyleft/gpl.html
 *
 */

#include "StdAfx.h"
#include "BaseMuxerOutputPin.h"

#include <mmreg.h>
#include <aviriff.h>
#include <atlpath.h>

#include <initguid.h>
#include <moreuuids.h>

//
// CBaseMuxerOutputPin
//

CBaseMuxerOutputPin::CBaseMuxerOutputPin(LPCWSTR pName, CBaseFilter* pFilter, CCritSec* pLock, HRESULT* phr)
	: CBaseOutputPin(_T("CBaseMuxerOutputPin"), pFilter, pLock, phr, pName)
{
}

IBitStream* CBaseMuxerOutputPin::GetBitStream()
{
	if(!m_pBitStream)
	{
		if(CComQIPtr<IStream> pStream = GetConnected())
		{
			m_pBitStream = new CBitStream(pStream, true);
		}
	}

	return m_pBitStream;
}

HRESULT CBaseMuxerOutputPin::BreakConnect()
{
	m_pBitStream = NULL;

	return __super::BreakConnect();
}

HRESULT CBaseMuxerOutputPin::DecideBufferSize(IMemAllocator* pAlloc, ALLOCATOR_PROPERTIES* pProperties)
{
    ASSERT(pAlloc);
    ASSERT(pProperties);

    HRESULT hr = NOERROR;

	pProperties->cBuffers = 1;
	pProperties->cbBuffer = 1;

    ALLOCATOR_PROPERTIES Actual;
    if(FAILED(hr = pAlloc->SetProperties(pProperties, &Actual))) return hr;

    if(Actual.cbBuffer < pProperties->cbBuffer) return E_FAIL;
    ASSERT(Actual.cBuffers == pProperties->cBuffers);

    return NOERROR;
}

HRESULT CBaseMuxerOutputPin::CheckMediaType(const CMediaType* pmt)
{
	return pmt->majortype == MEDIATYPE_Stream && pmt->subtype == MEDIASUBTYPE_NULL
		? S_OK
		: E_INVALIDARG;
}

HRESULT CBaseMuxerOutputPin::GetMediaType(int iPosition, CMediaType* pmt)
{
    CAutoLock cAutoLock(m_pLock);

	if(iPosition < 0) return E_INVALIDARG;
	if(iPosition > 0) return VFW_S_NO_MORE_ITEMS;

	pmt->ResetFormatBuffer();
	pmt->InitMediaType();
	pmt->majortype = MEDIATYPE_Stream;
	pmt->subtype = MEDIASUBTYPE_NULL;
	pmt->formattype = FORMAT_None;

	return S_OK;
}

HRESULT CBaseMuxerOutputPin::DeliverEndOfStream()
{
	m_pBitStream = NULL;

	return __super::DeliverEndOfStream();
}

STDMETHODIMP CBaseMuxerOutputPin::Notify(IBaseFilter* pSender, Quality q)
{
	return E_NOTIMPL;
}

//
// CBaseMuxerRawOutputPin
//

CBaseMuxerRawOutputPin::CBaseMuxerRawOutputPin(LPCWSTR pName, CBaseFilter* pFilter, CCritSec* pLock, HRESULT* phr)
	: CBaseMuxerOutputPin(pName, pFilter, pLock, phr)
{
}

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

	return 
		QI(IBaseMuxerRelatedPin)
		__super::NonDelegatingQueryInterface(riid, ppv);
}

void CBaseMuxerRawOutputPin::MuxHeader(const CMediaType& mt)
{
	CComQIPtr<IBitStream> pBitStream = GetBitStream();
	if(!pBitStream) return;

	const BYTE utf8bom[3] = {0xef, 0xbb, 0xbf};

	if((mt.subtype == FOURCCMap('1CVA') || mt.subtype == FOURCCMap('1cva')) && mt.formattype == FORMAT_MPEG2_VIDEO)
	{
		MPEG2VIDEOINFO* vih = (MPEG2VIDEOINFO*)mt.Format();

		for(DWORD i = 0; i < vih->cbSequenceHeader-2; i += 2)
		{
			pBitStream->BitWrite(0x00000001, 32);
			WORD size = (((BYTE*)vih->dwSequenceHeader)[i+0]<<8) | ((BYTE*)vih->dwSequenceHeader)[i+1];
			pBitStream->ByteWrite(&((BYTE*)vih->dwSequenceHeader)[i+2], size);
			i += size;
		}
	}
	else if(mt.subtype == MEDIASUBTYPE_UTF8)
	{
		pBitStream->ByteWrite(utf8bom, sizeof(utf8bom));
	}
	else if(mt.subtype == MEDIASUBTYPE_SSA || mt.subtype == MEDIASUBTYPE_ASS || mt.subtype == MEDIASUBTYPE_ASS2)
	{
		SUBTITLEINFO* si = (SUBTITLEINFO*)mt.Format();
		BYTE* p = (BYTE*)si + si->dwOffset;

		if(memcmp(utf8bom, p, 3) != 0) 
			pBitStream->ByteWrite(utf8bom, sizeof(utf8bom));;

		CStringA str((char*)p, mt.FormatLength() - (p - mt.Format()));
		pBitStream->StrWrite(str + '\n', true);

		if(str.Find("[Events]") < 0) 
			pBitStream->StrWrite("\n\n[Events]\n", true);
	}
	else if(mt.subtype == MEDIASUBTYPE_SSF)
	{
		DWORD dwOffset = ((SUBTITLEINFO*)mt.pbFormat)->dwOffset;
		try {m_ssf.Parse(ssf::MemoryInputStream(mt.pbFormat + dwOffset, mt.cbFormat - dwOffset, false, false));}
		catch(ssf::Exception&) {}
	}
	else if(mt.subtype == MEDIASUBTYPE_VOBSUB)
	{
		m_idx.RemoveAll();
	}
	else if(mt.majortype == MEDIATYPE_Audio 
		&& (mt.subtype == MEDIASUBTYPE_PCM 
		|| mt.subtype == MEDIASUBTYPE_DVD_LPCM_AUDIO
		|| mt.subtype == FOURCCMap(WAVE_FORMAT_EXTENSIBLE) 
		|| mt.subtype == FOURCCMap(WAVE_FORMAT_IEEE_FLOAT))
		&& mt.formattype == FORMAT_WaveFormatEx)
	{
		pBitStream->BitWrite('RIFF', 32);
		pBitStream->BitWrite(0, 32); // file length - 8, set later
		pBitStream->BitWrite('WAVE', 32);

		pBitStream->BitWrite('fmt ', 32);
		pBitStream->ByteWrite(&mt.cbFormat, 4);
		pBitStream->ByteWrite(mt.pbFormat, mt.cbFormat);

		pBitStream->BitWrite('data', 32);
		pBitStream->BitWrite(0, 32); // data length, set later
	}
}

void CBaseMuxerRawOutputPin::MuxPacket(const CMediaType& mt, const MuxerPacket* pPacket)
{
	CComQIPtr<IBitStream> pBitStream = GetBitStream();
	if(!pBitStream) return;

	const BYTE* pData = pPacket->pData.GetData();
	const int DataSize = pPacket->pData.GetCount();

	if(mt.subtype == MEDIASUBTYPE_AAC && mt.formattype == FORMAT_WaveFormatEx)
	{
		WAVEFORMATEX* wfe = (WAVEFORMATEX*)mt.Format();

		int profile = 0;

		int srate_idx = 11;
		if(92017 <= wfe->nSamplesPerSec) srate_idx = 0;
		else if(75132 <= wfe->nSamplesPerSec) srate_idx = 1;
		else if(55426 <= wfe->nSamplesPerSec) srate_idx = 2;
		else if(46009 <= wfe->nSamplesPerSec) srate_idx = 3;
		else if(37566 <= wfe->nSamplesPerSec) srate_idx = 4;
		else if(27713 <= wfe->nSamplesPerSec) srate_idx = 5;
		else if(23004 <= wfe->nSamplesPerSec) srate_idx = 6;
		else if(18783 <= wfe->nSamplesPerSec) srate_idx = 7;
		else if(13856 <= wfe->nSamplesPerSec) srate_idx = 8;
		else if(11502 <= wfe->nSamplesPerSec) srate_idx = 9;
		else if(9391 <= wfe->nSamplesPerSec) srate_idx = 10;

		int channels = wfe->nChannels;

		if(wfe->cbSize >= 2)
		{
			BYTE* p = (BYTE*)(wfe+1);
			profile = (p[0]>>3)-1;
			srate_idx = ((p[0]&7)<<1)|((p[1]&0x80)>>7);
			channels = (p[1]>>3)&15;
		}

		int len = (DataSize + 7) & 0x1fff;

		BYTE hdr[7] = {0xff, 0xf9};
		hdr[2] = (profile<<6) | (srate_idx<<2) | ((channels&4)>>2);
		hdr[3] = ((channels&3)<<6) | (len>>11);
		hdr[4] = (len>>3)&0xff;
		hdr[5] = ((len&7)<<5) | 0x1f;
		hdr[6] = 0xfc;

		pBitStream->ByteWrite(hdr, sizeof(hdr));
	}
	else if((mt.subtype == FOURCCMap('1CVA') || mt.subtype == FOURCCMap('1cva')) && mt.formattype == FORMAT_MPEG2_VIDEO)
	{
		const BYTE* p = pData;
		int i = DataSize;

		while(i >= 4)
		{
			DWORD len = (p[0]<<24)|(p[1]<<16)|(p[2]<<8)|p[3];

			i -= len + 4;
			p += len + 4;
		}

		if(i == 0)
		{
			p = pData;
			i = DataSize;

			while(i >= 4)
			{
				DWORD len = (p[0]<<24)|(p[1]<<16)|(p[2]<<8)|p[3];

				pBitStream->BitWrite(0x00000001, 32);

				p += 4; 
				i -= 4;

				if(len > i || len == 1) {len = i; ASSERT(0);}

				pBitStream->ByteWrite(p, len);

				p += len;
				i -= len;
			}

			return;
		}
	}
	else if(mt.subtype == MEDIASUBTYPE_UTF8 || mt.majortype == MEDIATYPE_Text)
	{
		CStringA str((char*)pData, DataSize);
		str.Trim();
		if(str.IsEmpty()) return;

		DVD_HMSF_TIMECODE start = RT2HMSF(pPacket->rtStart, 25);
		DVD_HMSF_TIMECODE stop = RT2HMSF(pPacket->rtStop, 25);

		str.Format("%d\n%02d:%02d:%02d,%03d --> %02d:%02d:%02d,%03d\n%s\n\n", 
			pPacket->index+1,
			start.bHours, start.bMinutes, start.bSeconds, (int)((pPacket->rtStart/10000)%1000), 
			stop.bHours, stop.bMinutes, stop.bSeconds, (int)((pPacket->rtStop/10000)%1000),
			CStringA(str));

		pBitStream->StrWrite(str, true);

		return;
	}
	else if(mt.subtype == MEDIASUBTYPE_SSA || mt.subtype == MEDIASUBTYPE_ASS || mt.subtype == MEDIASUBTYPE_ASS2)
	{
		CStringA str((char*)pData, DataSize);
		str.Trim();
		if(str.IsEmpty()) return;

		DVD_HMSF_TIMECODE start = RT2HMSF(pPacket->rtStart, 25);
		DVD_HMSF_TIMECODE stop = RT2HMSF(pPacket->rtStop, 25);

		int fields = mt.subtype == MEDIASUBTYPE_ASS2 ? 10 : 9;

		CAtlList<CStringA> sl;
		Explode(str, sl, ',', fields);
		if(sl.GetCount() < fields) return;

		CStringA readorder = sl.RemoveHead(); // TODO
		CStringA layer = sl.RemoveHead();
		CStringA style = sl.RemoveHead();
		CStringA actor = sl.RemoveHead();
		CStringA left = sl.RemoveHead();
		CStringA right = sl.RemoveHead();
		CStringA top = sl.RemoveHead();
		if(fields == 10) top += ',' + sl.RemoveHead(); // bottom
		CStringA effect = sl.RemoveHead();
		str = sl.RemoveHead();

		if(mt.subtype == MEDIASUBTYPE_SSA) layer = "Marked=0";

		str.Format("Dialogue: %s,%d:%02d:%02d.%02d,%d:%02d:%02d.%02d,%s,%s,%s,%s,%s,%s,%s\n",
			layer,
			start.bHours, start.bMinutes, start.bSeconds, (int)((pPacket->rtStart/100000)%100), 
			stop.bHours, stop.bMinutes, stop.bSeconds, (int)((pPacket->rtStop/100000)%100),
			style, actor, left, right, top, effect, 
			CStringA(str));

		pBitStream->StrWrite(str, true);

		return;
	}
	else if(mt.subtype == MEDIASUBTYPE_SSF)
	{
		float start = (float)pPacket->rtStart / 10000000;
		float stop = (float)pPacket->rtStop / 10000000;
		m_ssf.Append(ssf::WCharInputStream(UTF8To16(CStringA((char*)pData, DataSize))), start, stop, true);
		return;
	}
	else if(mt.subtype == MEDIASUBTYPE_VOBSUB)
	{
		bool fTimeValid = pPacket->IsTimeValid();

		if(fTimeValid)
		{
			idx_t i;
			i.rt = pPacket->rtStart;
			i.fp = pBitStream->GetPos();
			m_idx.AddTail(i);
		}

		int DataSizeLeft = DataSize;

		while(DataSizeLeft > 0)
		{
			int BytesAvail = 0x7ec - (fTimeValid ? 9 : 4);
			int Size = min(BytesAvail, DataSizeLeft);
			int Padding = 0x800 - Size - 20 - (fTimeValid ? 9 : 4);

			pBitStream->BitWrite(0x000001ba, 32);
			pBitStream->BitWrite(0x440004000401ui64, 48);
			pBitStream->BitWrite(0x000003f8, 32);
			pBitStream->BitWrite(0x000001bd, 32);

			if(fTimeValid)
			{
				pBitStream->BitWrite(Size+9, 16);
				pBitStream->BitWrite(0x8180052100010001ui64, 64);
			}
			else
			{
				pBitStream->BitWrite(Size+4, 16);
				pBitStream->BitWrite(0x810000, 24);
			}

			pBitStream->BitWrite(0x20, 8);

			pBitStream->ByteWrite(pData, Size);

			pData += Size;
			DataSizeLeft -= Size;

			if(Padding > 0)
			{
				Padding -= 6;
				ASSERT(Padding >= 0);
				pBitStream->BitWrite(0x000001be, 32);
				pBitStream->BitWrite(Padding, 16);
				while(Padding-- > 0) pBitStream->BitWrite(0xff, 8);
			}

			fTimeValid = false;
		}

		return;
	}
	else if(mt.subtype == MEDIASUBTYPE_DVD_LPCM_AUDIO)
	{
		WAVEFORMATEX* wfe = (WAVEFORMATEX*)mt.Format();

		for(int i = 0, bps = wfe->wBitsPerSample/8; i < DataSize; i += bps)
			for(int j = bps-1; j >= 0; j--)
				pBitStream->BitWrite(pData[i+j], 8);

		return;
	}
	// else // TODO: restore more streams (vorbis to ogg)

	pBitStream->ByteWrite(pData, DataSize);
}

void CBaseMuxerRawOutputPin::MuxFooter(const CMediaType& mt)
{
	CComQIPtr<IBitStream> pBitStream = GetBitStream();
	if(!pBitStream) return;

	if(mt.majortype == MEDIATYPE_Audio 
	&& (mt.subtype == MEDIASUBTYPE_PCM 
	|| mt.subtype == MEDIASUBTYPE_DVD_LPCM_AUDIO
	|| mt.subtype == FOURCCMap(WAVE_FORMAT_EXTENSIBLE) 
	|| mt.subtype == FOURCCMap(WAVE_FORMAT_IEEE_FLOAT))
	&& mt.formattype == FORMAT_WaveFormatEx)
	{
		pBitStream->BitFlush();

		ASSERT(pBitStream->GetPos() <= 0xffffffff);
		UINT32 size = (UINT32)pBitStream->GetPos();

		size -= 8;
		pBitStream->Seek(4);
		pBitStream->ByteWrite(&size, 4);

		size -= sizeof(RIFFLIST) + sizeof(RIFFCHUNK) + mt.FormatLength();
		pBitStream->Seek(sizeof(RIFFLIST) + sizeof(RIFFCHUNK) + mt.FormatLength() + 4);
		pBitStream->ByteWrite(&size, 4);
	}
	else if(mt.subtype == MEDIASUBTYPE_SSF)
	{
		ssf::WCharOutputStream s;
		try {m_ssf.Dump(s);} catch(ssf::Exception&) {}
		CStringA str = UTF16To8(s.GetString());
		pBitStream->StrWrite(str, true);
	}
	else if(mt.subtype == MEDIASUBTYPE_VOBSUB)
	{
		if(CComQIPtr<IFileSinkFilter> pFSF = GetFilterFromPin(GetConnected()))
		{
			WCHAR* fn = NULL;
			if(SUCCEEDED(pFSF->GetCurFile(&fn, NULL)))
			{
				CPathW p(fn);
				p.RenameExtension(L".idx");
				CoTaskMemFree(fn);

				if(FILE* f = _tfopen(CString((LPCWSTR)p), _T("w")))
				{
					SUBTITLEINFO* si = (SUBTITLEINFO*)mt.Format();

					_ftprintf(f, _T("%s\n"), _T("# VobSub index file, v7 (do not modify this line!)"));

					fwrite(mt.Format() + si->dwOffset, mt.FormatLength() - si->dwOffset, 1, f);

					CString iso6391 = ISO6392To6391(si->IsoLang);
					if(iso6391.IsEmpty()) iso6391 = _T("--");
					_ftprintf(f, _T("\nlangidx: 0\n\nid: %s, index: 0\n"), iso6391);

					CString alt = CString(CStringW(si->TrackName));
					if(!alt.IsEmpty()) _ftprintf(f, _T("alt: %s\n"), alt);

					POSITION pos = m_idx.GetHeadPosition();
					while(pos)
					{
						const idx_t& i = m_idx.GetNext(pos);
						DVD_HMSF_TIMECODE start = RT2HMSF(i.rt, 25);
						_ftprintf(f, _T("timestamp: %02d:%02d:%02d:%03d, filepos: %09I64x\n"), 
							start.bHours, start.bMinutes, start.bSeconds, (int)((i.rt/10000)%1000),
							i.fp);
					}

					fclose(f);
				}
			}
		}
	}
}
