/**************************************************************************
 * 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 "EcoDecoWriter.h"

CEcoDecoWriter::CEcoDecoWriter(IUnknown *pUnk, HRESULT *phr)
   : CSimpleWriter(L"EcoDecoWriter", pUnk, CLSID_EcoDecoWriter, phr)
{  

   m_nWait = 0;
   m_llInputSize = 0;
   m_llWriteDataSize = 0;
   m_nOutputType = 0;
   m_nPacketNumber = 0;

   m_bCheckMode = false;
   m_bUseTrackgain = false;
   m_bUseAlbumgain = false;

   ::ZeroMemory(&m_inFormat, sizeof(m_inFormat));
   ::ZeroMemory(&m_volbisFormat, sizeof(m_volbisFormat));
}
CEcoDecoWriter::~CEcoDecoWriter()  
{  

}  
HRESULT CEcoDecoWriter::SetWait(int nWait)
{
   CAutoLock lock(&m_Lock);

   m_nWait = nWait;

   return S_OK;
}
HRESULT CEcoDecoWriter::GetTransformedBytes(LONGLONG *pllBytes)
{
   CAutoLock lock(&m_Lock);

   *pllBytes = m_llInputSize;

   return S_OK;
}
HRESULT CEcoDecoWriter::GetInFormat(WAVEFORMATEX *pFormat)
{  
   CheckPointer(pFormat, E_POINTER);
   ::CopyMemory(pFormat, &m_inFormat, sizeof(m_inFormat));

   return S_OK;
}
int CEcoDecoWriter::GetOutputMode(void)
{
   return m_nOutputType;
}
void CEcoDecoWriter::SetOutputMode(int nOutputType)
{
   m_nOutputType = nOutputType;
}
HRESULT CEcoDecoWriter::CheckMode(bool bCheckMode)
{  
   m_bCheckMode = bCheckMode;
   return S_OK;
}
HRESULT CEcoDecoWriter::SetReplaygain(bool bUseTrackgain, bool bUseAlbumgain)
{  
   m_bUseTrackgain = bUseTrackgain;
   m_bUseAlbumgain = bUseAlbumgain;

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

   if(m_nOutputType == 0)
      return S_FALSE;

   if( (m_nOutputType == CONTAINER_RAW_WAV) && (pmtIn->subtype == MEDIASUBTYPE_PCM || pmtIn->subtype == MEDIASUBTYPE_IEEE_FLOAT) )
   {  
      ::CopyMemory(&m_inFormat, (WAVEFORMATEX *)pmtIn->Format(), sizeof(WAVEFORMATEX));
      return S_OK;
   }
   else if(m_nOutputType == CONTAINER_RAW_AAC && pmtIn->subtype == MEDIASUBTYPE_AAC)
   {  
      ::CopyMemory(&m_inFormat, (WAVEFORMATEX *)pmtIn->Format(), sizeof(WAVEFORMATEX));
      return S_OK;
   }
   else if(m_nOutputType == CONTAINER_RAW_OGG && pmtIn->subtype == MEDIASUBTYPE_Vorbis)
   {  
      ::CopyMemory(&m_volbisFormat, (VORBISFORMAT *)pmtIn->Format(), sizeof(m_volbisFormat));
      return S_OK;
   }
   else if(m_nOutputType == CONTAINER_RAW)
   {
      if(pmtIn->subtype == MEDIASUBTYPE_MP3 || pmtIn->subtype == MEDIASUBTYPE_MPEG1AudioPayload ||
         pmtIn->subtype == MEDIASUBTYPE_AVI_AC3 || pmtIn->subtype == MEDIASUBTYPE_DOLBY_AC3 ||
         pmtIn->subtype == MEDIASUBTYPE_MPEG2_AUDIO)
      {
         return S_OK;
      }
   }

   return S_FALSE;
}
void CEcoDecoWriter::MakeAdtsHeader(void)
{
   m_adtsHeader[0] = 0xff;
   m_adtsHeader[1] = 0xf9;
   m_adtsHeader[2] = 0x40;

   switch(m_inFormat.nSamplesPerSec)
   {
      case 96000: break;
      case 88200: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x01 << 2); break;
      case 64000: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x02 << 2); break;
      case 48000: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x03 << 2); break;
      case 44100: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x04 << 2); break;
      case 32000: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x05 << 2); break;
      case 24000: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x06 << 2); break;
      case 22050: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x07 << 2); break;
      case 16000: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x08 << 2); break;
      case 12000: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x09 << 2); break;
      case 11025: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x0a << 2); break;
      case  8000: m_adtsHeader[2] = m_adtsHeader[2] | (BYTE)(0x0b << 2); break;
   }

   m_adtsHeader[3] = 0x00;

   if(m_inFormat.nChannels >= 4)
   {
      m_adtsHeader[2] = m_adtsHeader[2] | 0x01;
      m_adtsHeader[3] = m_adtsHeader[3] | ( (m_inFormat.nChannels & 0x03) << 6 );
   }
   else
      m_adtsHeader[3] = ( m_inFormat.nChannels << 6 );

   m_adtsHeader[4] = 0;
   m_adtsHeader[5] = 0x1f;
   m_adtsHeader[6] = 0xfc;
}
HRESULT CEcoDecoWriter::WriteWavHeader(HANDLE hFile)
{
   BOOL flag;
   DWORD dwWriteSize;
   DWORD dwBuffer;

   flag = ::WriteFile(hFile , "RIFF" , 4, &dwWriteSize , NULL);
   if(flag == FALSE) return E_FAIL;

   dwBuffer = 0;
   flag = ::WriteFile(hFile , &dwBuffer  , 4, &dwWriteSize , NULL);
   if(flag == FALSE) return E_FAIL;

   flag = ::WriteFile(hFile , "WAVE" , 4, &dwWriteSize , NULL);
   if(flag == FALSE) return E_FAIL;

   flag = ::WriteFile(hFile , "fmt " , 4 , &dwWriteSize , NULL);
   if(flag == FALSE) return E_FAIL;

   dwBuffer = 16;
   flag = ::WriteFile(hFile, &dwBuffer, 4, &dwWriteSize , NULL);
   if(flag == FALSE) return E_FAIL;

   flag = ::WriteFile(hFile, &m_inFormat.wFormatTag, 2, &dwWriteSize , NULL);
   if(flag == FALSE) return S_FALSE;

   flag = ::WriteFile(hFile, &m_inFormat.nChannels, 2, &dwWriteSize , NULL);
   if(flag == FALSE) return S_FALSE;

   flag = ::WriteFile(hFile, &m_inFormat.nSamplesPerSec, 4, &dwWriteSize , NULL);
   if(flag == FALSE) return S_FALSE;

   flag = ::WriteFile(hFile, &m_inFormat.nAvgBytesPerSec, 4, &dwWriteSize , NULL);
   if(flag == FALSE) return S_FALSE;

   flag = ::WriteFile(hFile, &m_inFormat.nBlockAlign, 2, &dwWriteSize , NULL);
   if(flag == FALSE) return S_FALSE;

   flag = ::WriteFile(hFile, &m_inFormat.wBitsPerSample, 2, &dwWriteSize , NULL);
   if(flag == FALSE) return S_FALSE;

   if(m_bUseTrackgain == true || m_bUseAlbumgain == true)
   {  
      flag = ::WriteFile(hFile, "rgad", 4, &dwWriteSize , NULL);
      if(flag == FALSE) return S_FALSE;

      dwBuffer = 8;
      flag = ::WriteFile(hFile , &dwBuffer, 4, &dwWriteSize , NULL);
      if(flag == FALSE) return E_FAIL;

      dwBuffer = 0;
      flag = ::WriteFile(hFile , &dwBuffer, 4, &dwWriteSize , NULL);
      if(flag == FALSE) return E_FAIL;

      dwBuffer = 0;
      flag = ::WriteFile(hFile , &dwBuffer, 4, &dwWriteSize , NULL);
      if(flag == FALSE) return E_FAIL;
   }

   flag = ::WriteFile(hFile, "data", 4, &dwWriteSize , NULL);
   if(flag == FALSE) return S_FALSE;

   dwBuffer = 0;
   flag = ::WriteFile(hFile , &dwBuffer, 4, &dwWriteSize , NULL);
   if(flag == FALSE) return E_FAIL;

   return S_OK;
}
HRESULT CEcoDecoWriter::WriteOggPage(HANDLE hFile, IMediaSample *pInSample, bool bEos)
{
   BOOL flag;
   int nResult;
   DWORD dwWriteSize;

   ogg_page   og;
   ogg_packet op; 

   pInSample->GetPointer(&op.packet);
   op.bytes = pInSample->GetActualDataLength();

   if(op.bytes == 0)
      return S_OK;

   if(bEos == false)
      op.e_o_s = 0;
   else
      op.e_o_s = 1;

   op.b_o_s = 0;
   op.packetno = m_nPacketNumber;
   m_nPacketNumber++;

   if(op.packet[0] & PACKET_HEADER_BIT)
   {  

      if(op.packet[0] == PACKET_HEADER_IDENTIFICATION)
         op.b_o_s = 1; 

      op.granulepos = 0;
   }
   else
   {  

      REFERENCE_TIME rtStart, rtStop;
      pInSample->GetTime(&rtStart, &rtStop);

      op.granulepos = (ogg_int64_t)( ((double)rtStart * m_volbisFormat.nSamplesPerSec / 10000000.0) + 0.5); 
   }

   nResult = ogg_stream_packetin(&m_os, &op);
    
   while(true)
   {
      nResult = ogg_stream_pageout(&m_os, &og);
      if(nResult == 0)
         break;

      flag = ::WriteFile(hFile, og.header, og.header_len, &dwWriteSize , NULL);
      if(flag == FALSE)
         return S_FALSE;

      flag = ::WriteFile(hFile, og.body, og.body_len, &dwWriteSize , NULL);
      if(flag == FALSE)
         return S_FALSE;
      
      if(ogg_page_eos(&og) != 0)
         break;
   }

   return S_OK;
}
HRESULT CEcoDecoWriter::OnStart(HANDLE hFile)
{
   if(m_bCheckMode == true)
      return S_OK;

   m_llInputSize = 0;
   m_llWriteDataSize = 0;
   m_nPacketNumber = 0;

   if(m_nOutputType == CONTAINER_RAW_WAV)
   {  

      HRESULT hr;
      hr = WriteWavHeader(hFile);
      if(FAILED(hr))
         return E_FAIL;
   }
   else if(m_nOutputType == CONTAINER_RAW_AAC)
   {
      MakeAdtsHeader();
   }
   else if(m_nOutputType == CONTAINER_RAW_OGG)
   {
      ::srand((unsigned int)time(NULL));
      ogg_stream_init(&m_os,rand());
   }

   return S_OK;
}
HRESULT CEcoDecoWriter::OnReceive(HANDLE hFile, IMediaSample *pInSample)
{
   if(m_nWait > 0)
      ::Sleep(m_nWait);

   m_llInputSize += pInSample->GetActualDataLength();

   if(m_bCheckMode == true)
      return S_OK;

   if(m_nOutputType == CONTAINER_RAW_OGG)
   {
      WriteOggPage(hFile, pInSample, false);
   }
   else
   {
      BYTE *pbData;
      DWORD dwWritten;

      pInSample->GetPointer(&pbData);

      if(m_nOutputType == CONTAINER_RAW_AAC)
      {
         int nLength;
         nLength = pInSample->GetActualDataLength() + 7;

         m_adtsHeader[4] = (BYTE)(nLength >> 3);
         m_adtsHeader[5] = ((BYTE)(nLength & 0x07) << 5) | 0x1f; 

         ::WriteFile(hFile, (PVOID)m_adtsHeader, 7, &dwWritten, NULL);
      }

      ::WriteFile(hFile, (PVOID)pbData, (DWORD)pInSample->GetActualDataLength(), &dwWritten, NULL);

      if(m_nOutputType == CONTAINER_RAW_WAV)
      {  
         m_llWriteDataSize += pInSample->GetActualDataLength();
      }
   }

   return S_OK;
}
HRESULT CEcoDecoWriter::OnStop(HANDLE hFile, bool bEos)
{
   if(m_bCheckMode == true)
   {  
      return S_OK;
   }

   if(hFile == NULL)
      return S_OK;

   if(bEos == true)
   {
      if(m_nOutputType == CONTAINER_RAW_WAV)
      {  

         BOOL flag;
         DWORD dwWriteSize;

         DWORD dwTotalSize = 36 + (DWORD)m_llWriteDataSize;
         DWORD dwDataPosition = 40;

         if(m_bUseTrackgain == true || m_bUseAlbumgain == true)
         {
            dwTotalSize += 16;
            dwDataPosition += 16;
         }

         ::SetFilePointer(hFile, 4, NULL, FILE_BEGIN);
         flag = ::WriteFile(hFile , &dwTotalSize, 4, &dwWriteSize , NULL);
         if(flag == FALSE) return E_FAIL;

         ::SetFilePointer(hFile, dwDataPosition, NULL, FILE_BEGIN);
         flag = ::WriteFile(hFile , &m_llWriteDataSize, 4, &dwWriteSize , NULL);
         if(flag == FALSE) return E_FAIL;
      }

      if(m_nOutputType == CONTAINER_RAW_OGG)
      {
         ogg_stream_clear(&m_os);
      }
   }

   return S_OK;
}
STDMETHODIMP CEcoDecoWriter::OnQueryInterface(REFIID riid, void ** ppv)
{
   if(riid == IID_IEcoDecoWriterInterface)
      return GetInterface((IEcoDecoWriterInterface *)this, ppv);
   else if(riid == IID_IEcoDecoInterface)
      return GetInterface((IEcoDecoInterface *)this, ppv);

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

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

   return dynamic_cast<CUnknown *>(pNewObject);
}
