/*
 * (C) 2003-2006 Gabest
 * (C) 2006-2016 see Authors.txt
 *
 * This file is part of MPC-BE.
 *
 * MPC-BE 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 3 of the License, or
 * (at your option) any later version.
 *
 * MPC-BE 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#include "stdafx.h"
#include <algorithm>
#include "BaseMuxerOutputPin.h"

#include <MMReg.h>
#include <aviriff.h>
#include <atlpath.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 = DNew 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() - (ULONG)(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_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 = int(pPacket->pData.GetCount());

	if (mt.subtype == MEDIASUBTYPE_RAW_AAC1 && 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] = BYTE((profile << 6) | (srate_idx << 2) | ((channels & 4) >> 2));
		hdr[3] = BYTE(((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 > (DWORD)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%02u:%02u:%02u,%03d --> %02u:%02u:%02u,%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);

		size_t 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,%u:%02u:%02u.%02d,%u:%02u:%02u.%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_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 = (std::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();

		// This code is probably totally broken for anything but 16 bits
		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_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);

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

					_ftprintf_s(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_s(f, _T("\nlangidx: 0\n\nid: %s, index: 0\n"), iso6391.GetString());

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

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

					fclose(f);
				}
			}
		}
	}
}
