/**************************************************************************
 * Copyright (C) 2008 Cocha                                               *
 * http://sourceforge.jp/projects/ecodecotool/                            *
 *                                                                        *
 *  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. *
 *                                                                        *
 **************************************************************************/

#include "DecodeAc3.h"

struct ac3FrameSize
{
   int nBitrate;
   int nFrameSize[3];
};

static const struct ac3FrameSize ac3FrameSizeTable[] =
{
   {  32 ,{  64 ,  69 ,  96}},
   {  32 ,{  64 ,  70 ,  96}},
   {  40 ,{  80 ,  87 , 120}},
   {  40 ,{  80 ,  88 , 120}},
   {  48 ,{  96 , 104 , 144}},
   {  48 ,{  96 , 105 , 144}},
   {  56 ,{ 112 , 121 , 168}},
   {  56 ,{ 112 , 122 , 168}},
   {  64 ,{ 128 , 139 , 192}},
   {  64 ,{ 128 , 140 , 192}},
   {  80 ,{ 160 , 174 , 240}},
   {  80 ,{ 160 , 175 , 240}},
   {  96 ,{ 192 , 208 , 288}},
   {  96 ,{ 192 , 209 , 288}},
   { 112 ,{ 224 , 243 , 336}},
   { 112 ,{ 224 , 244 , 336}},
   { 128 ,{ 256 , 278 , 384}},
   { 128 ,{ 256 , 279 , 384}},
   { 160 ,{ 320 , 348 , 480}},
   { 160 ,{ 320 , 349 , 480}},
   { 192 ,{ 384 , 417 , 576}},
   { 192 ,{ 384 , 418 , 576}},
   { 224 ,{ 448 , 487 , 672}},
   { 224 ,{ 448 , 488 , 672}},
   { 256 ,{ 512 , 557 , 768}},
   { 256 ,{ 512 , 558 , 768}},
   { 320 ,{ 640 , 696 , 960}},
   { 320 ,{ 640 , 697 , 960}},
   { 384 ,{ 768 , 835 ,1152}},
   { 384 ,{ 768 , 836 ,1152}},
   { 448 ,{ 896 , 975 ,1344}},
   { 448 ,{ 896 , 976 ,1344}},
   { 512 ,{1024 ,1114 ,1536}},
   { 512 ,{1024 ,1115 ,1536}},
   { 576 ,{1152 ,1253 ,1728}},
   { 576 ,{1152 ,1254 ,1728}},
   { 640 ,{1280 ,1393 ,1920}},
   { 640 ,{1280 ,1394 ,1920}}
};

CDecodeAc3::CDecodeAc3(IUnknown *pUnk, HRESULT *phr)
   : CSimpleTransform(L"Aac Decoder", pUnk, CLSID_DecodeAc3, phr)
{  

   ::ZeroMemory(&m_outFormat, sizeof(m_outFormat));
   m_outFormat.wFormatTag = WAVE_FORMAT_IEEE_FLOAT;
   m_outFormat.wBitsPerSample = 32;

   m_pbInBuffer = NULL;
   m_nInArraySize = 0;
   m_nInBufferLength = 0;

   m_bSeek = false;
   m_rtCurrent = 0;

   if(phr != NULL)
      *phr = S_OK;
}
CDecodeAc3::~CDecodeAc3()  
{  

   SAFE_DELETE(m_pbInBuffer);
}  
HRESULT CDecodeAc3::GetOutFormat(void *pFormat)
{  
   CheckPointer(pFormat, E_POINTER);
   ::CopyMemory(pFormat, &m_outFormat, sizeof(m_outFormat));
   return S_OK;
}
HRESULT CDecodeAc3::SetOutFormat(void *pFormat)
{  

   CheckPointer(pFormat, E_POINTER);
   WAVEFORMATEX *pwf = (WAVEFORMATEX *)pFormat;

   if(pwf->wFormatTag == WAVE_FORMAT_PCM && pwf->wBitsPerSample == 16)
   {
      ::CopyMemory(&m_outFormat, pFormat, sizeof(m_outFormat));
      return S_OK;
   }
   else if(pwf->wFormatTag == WAVE_FORMAT_IEEE_FLOAT && pwf->wBitsPerSample == 32)
   {
      ::CopyMemory(&m_outFormat, pFormat, sizeof(m_outFormat));
      return S_OK;
   }

   return E_FAIL;
}
HRESULT CDecodeAc3::OnStart(void)
{  
   m_rtCurrent = 0;

   bool bResult;

   bResult = mt2spk(m_pInput->CurrentMediaType(), m_inSpeakers);
   if(bResult == FALSE)
      return E_FAIL;

   bResult = mt2spk(m_pOutput->CurrentMediaType(), m_outSpeakers);
   if(bResult == FALSE)
      return E_FAIL;

   bResult = m_decoderGraph.set_input(m_inSpeakers);
   if(bResult == FALSE)
      return E_FAIL;

   bResult = m_decoderGraph.set_user(m_outSpeakers);
   if(bResult == FALSE)
      return E_FAIL;

   m_decoderGraph.proc.set_input_order(std_order);
   m_decoderGraph.proc.set_output_order(win_order);
   m_decoderGraph.reset();
   m_chunk.set_empty(m_outSpeakers);

   m_nInBufferLength = 0;

   return S_OK;
}
HRESULT CDecodeAc3::OnSeek(void)
{  
   m_rtCurrent = 0;
   m_bSeek = true;
   return S_OK;
}
HRESULT CDecodeAc3::OnTransform(IMediaSample *pInSample, IMediaSample *pOutSample)
{  

   HRESULT hr;
   BYTE *pbInBuffer = NULL;
   BYTE *pbOutBuffer = NULL;


   int nSrcLength = pInSample->GetActualDataLength();

   pInSample->GetPointer(&pbInBuffer);

   hr = pOutSample->GetPointer(&pbOutBuffer);

   if(m_pbInBuffer == NULL || m_bSeek == true)
   {

      int i;

      for(i=1;i<(nSrcLength-1);i++)
      {
         if(pbInBuffer[i] == 0x0B && pbInBuffer[i+1]==0x77)
            break;
      }

      if(i == (nSrcLength-1))
      {
         pOutSample->SetActualDataLength(0);
         return S_OK;
      }

      pbInBuffer += i;
      nSrcLength -= i;
      m_nInBufferLength = 0;

      if(m_pbInBuffer == NULL)
      {  

         if(nSrcLength < 4)
         {
            pOutSample->SetActualDataLength(0);
            return S_OK;
         }

         int nBitrateCode = pbInBuffer[4] & 0x3F;
         int nFrameSizeCode = (pbInBuffer[4] & 0xC0) >> 6;

         if(nBitrateCode < 38 && nFrameSizeCode < 3)
            m_nInArraySize = ac3FrameSizeTable[nBitrateCode].nFrameSize[nFrameSizeCode] * 2;
         else
            m_nInArraySize = 2048;

         m_pbInBuffer = new BYTE[m_nInArraySize];
      }
      else
      {  
         m_bSeek = false;
      }
   }

   bool bResult;
   int nOutLength = 0;
   int nCopyLength;

   while(nSrcLength > 0)
   {
      if( (m_nInBufferLength + nSrcLength) < m_nInArraySize)
      {
         nCopyLength = nSrcLength;
         ::CopyMemory(&m_pbInBuffer[m_nInBufferLength], pbInBuffer, nCopyLength);
         m_nInBufferLength += nCopyLength;
         nSrcLength = 0;
      }
      else
      {
         nCopyLength = m_nInArraySize - m_nInBufferLength;
         ::CopyMemory(&m_pbInBuffer[m_nInBufferLength], pbInBuffer, nCopyLength);
         m_nInBufferLength = 0;
         nSrcLength -= nCopyLength;
         pbInBuffer += nCopyLength;

         m_chunk.set_rawdata(m_inSpeakers, m_pbInBuffer, m_nInArraySize);
         bResult = m_decoderGraph.process(&m_chunk);
         if(bResult == false)
            return E_FAIL;

         m_chunk.set_empty(m_outSpeakers);

         bResult = m_decoderGraph.get_chunk(&m_chunk);
         if(bResult == false)
         {
            m_chunk.set_empty(m_outSpeakers);
            return E_FAIL;
         }

         if(m_chunk.size > 0)
         {
            ::CopyMemory(pbOutBuffer, m_chunk.rawdata, m_chunk.size);
            nOutLength += m_chunk.size;
            pbOutBuffer += m_chunk.size;
            m_chunk.drop(m_chunk.size);
         }
      }
   }

   pOutSample->SetSyncPoint(TRUE);
   pOutSample->SetDiscontinuity(FALSE);
   pOutSample->SetPreroll(FALSE);
   pOutSample->SetActualDataLength(nOutLength);

   REFERENCE_TIME rtStop;

   rtStop = m_rtCurrent + (REFERENCE_TIME)((LONGLONG)nOutLength * ONE_SECOND / m_outFormat.nAvgBytesPerSec);
   pOutSample->SetTime(&m_rtCurrent, &rtStop);
   m_rtCurrent = rtStop;

   return S_OK;
}
HRESULT CDecodeAc3::OnConnectInPin(const CMediaType *pmtIn)
{  

   if(pmtIn->majortype != MEDIATYPE_Audio)
      return S_FALSE;

   if(pmtIn->subtype != MEDIASUBTYPE_DOLBY_AC3 && pmtIn->subtype != MEDIASUBTYPE_AVI_AC3)
      return S_FALSE;

   if(pmtIn->formattype != FORMAT_WaveFormatEx)
      return S_FALSE;

   WAVEFORMATEX *pwf = (WAVEFORMATEX *)pmtIn->Format();

   m_outFormat.cbSize = 0;
   m_outFormat.nSamplesPerSec = pwf->nSamplesPerSec;

   if(pwf->nChannels <= 2)
      m_outFormat.nChannels = pwf->nChannels;
   else
      m_outFormat.nChannels = 2;

   m_outFormat.nBlockAlign = m_outFormat.nChannels * (m_outFormat.wBitsPerSample / 8);
   m_outFormat.nAvgBytesPerSec = m_outFormat.nSamplesPerSec * m_outFormat.nBlockAlign;

   return S_OK;
}
HRESULT CDecodeAc3::OnConnectOutPin(const CMediaType *pmtIn, int nInBufferSize, CMediaType *pmtOut, int *pnOutBufferSize)
{  

   pmtOut->SetType(&MEDIATYPE_Audio);
   pmtOut->SetSubtype(&MEDIASUBTYPE_PCM);
   pmtOut->SetTemporalCompression(FALSE);
   pmtOut->SetSampleSize(0);
   pmtOut->SetFormatType(&FORMAT_WaveFormatEx);
   pmtOut->SetFormat((BYTE*)&m_outFormat, sizeof(m_outFormat));

   *pnOutBufferSize = m_outFormat.nAvgBytesPerSec * 2;

   return S_OK;
}
STDMETHODIMP CDecodeAc3::OnQueryInterface(REFIID riid, void ** ppv)
{
   if(riid == IID_IDecodeInterface)
      return GetInterface((IDecodeInterface *)this, ppv);

   return CTransformFilter::NonDelegatingQueryInterface(riid, ppv);
}
CUnknown * WINAPI CDecodeAc3::CreateInstance(LPUNKNOWN punk, HRESULT *phr)
{  

   ASSERT(phr);
    
   CDecodeAc3 *pNewObject = new CDecodeAc3(punk, phr);
   if (pNewObject == NULL)
   {
      if(phr != NULL)
         *phr = E_OUTOFMEMORY;
   }

   return dynamic_cast<CUnknown *>(pNewObject);
}
