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


CEncodeVorbis::CEncodeVorbis(IUnknown *pUnk, HRESULT *phr)
   : CSimpleWriter(L"OggVorbis Writer", pUnk, CLSID_EncodeVorbis, phr)
{  

   
   ::ZeroMemory(&m_inFormat, sizeof(m_inFormat));
   ::ZeroMemory(&m_outFormat, sizeof(m_outFormat));
   m_bInitialize = false;

   m_bCheckVolume = false;
   m_bUseTrackgain = false;
   m_bUseAlbumgain = false;
   m_dPeak = 0.0;
}

CEncodeVorbis::~CEncodeVorbis()  
{  

   
   if(m_bInitialize == true)
   {
      
      ogg_stream_clear(&m_os);
      vorbis_block_clear(&m_vb);
      vorbis_dsp_clear(&m_vd);
      vorbis_comment_clear(&m_vc);
      vorbis_info_clear(&m_vi);

      m_bInitialize = false;
   }
}  

HRESULT CEncodeVorbis::CheckVolumeMode(bool bCheckMode)
{  
   m_bCheckVolume = bCheckMode;
   return S_OK;
}

HRESULT CEncodeVorbis::GetReplayGain(double *pdPeak, double *pdGain)
{
   *pdPeak = (m_dPeak / 32768.0);
   *pdGain = (double)GetTitleGain();

   return S_OK;
}

HRESULT CEncodeVorbis::GetGainResult(unsigned int *pArrayResult)
{
   GetResult(pArrayResult);
   return S_OK;
}

HRESULT CEncodeVorbis::SetReplaygain(bool bUseTrackgain, bool bUseAlbumgain)
{  
   m_bUseTrackgain = bUseTrackgain;
   m_bUseAlbumgain = bUseAlbumgain;

   return S_OK;
}

HRESULT CEncodeVorbis::SetOutFormat(void *pFormat)
{
   
   VORBISFORMAT *pVorbisFormat = (VORBISFORMAT *)pFormat;

   int ret;
   vorbis_info vi;
   vorbis_info_init(&vi);

   if(-0.11f <= pVorbisFormat->fQuality && pVorbisFormat->fQuality <= 1.01f)
   {  
      ret = vorbis_encode_init_vbr(&vi, (long)pVorbisFormat->nChannels, pVorbisFormat->nSamplesPerSec, pVorbisFormat->fQuality);
   }
   else
   {  
      ret = vorbis_encode_init(&vi, (long)pVorbisFormat->nChannels, pVorbisFormat->nSamplesPerSec, pVorbisFormat->nMaxBitsPerSec, pVorbisFormat->nAvgBitsPerSec, pVorbisFormat->nMinBitsPerSec);
   }

   vorbis_info_clear(&vi);

   if(ret < 0)
   {  
      return E_FAIL;
   }

   
   ::CopyMemory(&m_outFormat, pFormat, sizeof(m_outFormat));

   return S_OK;
}

HRESULT CEncodeVorbis::OnConnectInPin(const CMediaType *pmtIn)
{  

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

   if(pmtIn->subtype != MEDIASUBTYPE_PCM && pmtIn->subtype != MEDIASUBTYPE_IEEE_FLOAT)
      return S_FALSE;

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

   WAVEFORMATEX *pwf = (WAVEFORMATEX *)pmtIn->Format();
   if(pwf->nChannels <= 2 && pwf->nSamplesPerSec <= 48000)
   {  

      ::CopyMemory(&m_inFormat, pwf, sizeof(m_inFormat));

      m_outFormat.nChannels = pwf->nChannels;
      m_outFormat.nSamplesPerSec = pwf->nSamplesPerSec;

      return S_OK;
   }

   return S_FALSE;
}

HRESULT CEncodeVorbis::OnStart(HANDLE hFile)
{
   int eos = 0, ret;

   if(m_bCheckVolume == true || m_bUseTrackgain == true || m_bUseAlbumgain == true)
   {  
      
      InitGainAnalysis(m_outFormat.nSamplesPerSec);
   }

   if(m_bCheckVolume == true)
   {  
      return S_OK;
   }

   
   vorbis_info_init(&m_vi);

   if(-0.11f <= m_outFormat.fQuality && m_outFormat.fQuality <= 1.01f)
   {  
      ret = vorbis_encode_init_vbr(&m_vi, (long)m_outFormat.nChannels, m_outFormat.nSamplesPerSec, m_outFormat.fQuality);
   }
   else
   {  
      ret = vorbis_encode_init(&m_vi, (long)m_outFormat.nChannels, m_outFormat.nSamplesPerSec, m_outFormat.nMaxBitsPerSec, m_outFormat.nAvgBitsPerSec, m_outFormat.nMinBitsPerSec);
   }

   if(ret < 0)
   {  

      vorbis_info_clear(&m_vi);
      return E_FAIL;
   }

   
   vorbis_comment_init(&m_vc);

   if(m_bUseTrackgain == true)
   {  
      char szText[16];
      double dGain;

      dGain = 1.0;
      sprintf_s(szText, 16, "%1.8f", dGain);
      vorbis_comment_add_tag(&m_vc, "REPLAYGAIN_TRACK_PEAK", szText);

      dGain = -1.0;
      sprintf_s(szText, 16, "%+1.8f dB", dGain);
      vorbis_comment_add_tag(&m_vc, "REPLAYGAIN_TRACK_GAIN", szText);
   }

   if(m_bUseAlbumgain == true)
   {  

      char szText[16];
      double dGain;

      dGain = 1.0;
      sprintf_s(szText, 16, "%1.8f", dGain);
      vorbis_comment_add_tag(&m_vc, "REPLAYGAIN_ALBUM_PEAK", szText);

      dGain = -1.0;
      sprintf_s(szText, 16, "%+1.8f dB", dGain);
      vorbis_comment_add_tag(&m_vc, "REPLAYGAIN_ALBUM_GAIN", szText);
   }

   
   vorbis_analysis_init(&m_vd, &m_vi);
   vorbis_block_init(&m_vd, &m_vb);

   
   srand((unsigned int)time(NULL));
   ogg_stream_init(&m_os, rand());

   ogg_packet header;
   ogg_packet header_comm;
   ogg_packet header_code;

   vorbis_analysis_headerout(&m_vd, &m_vc, &header, &header_comm, &header_code);
   ogg_stream_packetin(&m_os, &header);
   ogg_stream_packetin(&m_os, &header_comm);
   ogg_stream_packetin(&m_os, &header_code);

   while(!eos)
   {
      int result = ogg_stream_flush(&m_os, &m_og);

      if(result==0)
         break;

      if(m_bCheckVolume == false)
      {  

         BOOL flag;
         DWORD dwWriteSize;

         flag = ::WriteFile(hFile , (LPVOID)m_og.header, m_og.header_len, &dwWriteSize , NULL);
         if(flag == FALSE)
         {
            vorbis_info_clear(&m_vi);
            return E_FAIL;
         }

         flag = ::WriteFile(hFile , (LPVOID)m_og.body, m_og.body_len, &dwWriteSize , NULL);
         if(flag == FALSE)
         {
            vorbis_info_clear(&m_vi);
            return E_FAIL;
         }
      }
   }

   m_bInitialize = true;

   return S_OK;
}

HRESULT CEncodeVorbis::Encode(HANDLE hFile, int nBufferSize)
{  

   int eos = 0;

   vorbis_analysis_wrote(&m_vd, nBufferSize);

   while(vorbis_analysis_blockout(&m_vd, &m_vb)==1)
   {
      vorbis_analysis(&m_vb, NULL);
      vorbis_bitrate_addblock(&m_vb);

      while(vorbis_bitrate_flushpacket(&m_vd, &m_op))
      {
         ogg_stream_packetin(&m_os, &m_op);
    
         while(!eos)
         {
            int result = ogg_stream_pageout(&m_os, &m_og);
            if(result == 0)
               break;

            if(m_bCheckVolume == false)
            {  
               BOOL flag;
               DWORD dwWriteSize;

               flag = ::WriteFile(hFile , (LPVOID)m_og.header, m_og.header_len, &dwWriteSize , NULL);
               if(flag == FALSE) return E_FAIL;

               flag = ::WriteFile(hFile , (LPVOID)m_og.body, m_og.body_len, &dwWriteSize , NULL);
               if(flag == FALSE) return E_FAIL;

               if(ogg_page_eos(&m_og))
                  eos = 1;
            }
         }
      }
   }

   return S_OK;
}

void CEncodeVorbis::Convert8to32f(BYTE *pbInBuffer, float **ppfOutBuffer, int nEncodeSamples)
{
   int i, j;
   int nChannels = m_outFormat.nChannels;

   for(i=0;i<nEncodeSamples;i++)
   {
      for(j=0;j<nChannels;j++)
      {
         if(*pbInBuffer > 128)
            ppfOutBuffer[j][i] = (float)(*pbInBuffer - 128)/ 127.f;
         else
            ppfOutBuffer[j][i] = (float)(*pbInBuffer - 128)/ 128.f;

         *pbInBuffer++;
      }
   }
}

void CEncodeVorbis::Convert16to32f(short *psInBuffer, float **ppfOutBuffer, int nEncodeSamples)
{
   int i, j;
   int nChannels = m_outFormat.nChannels;

   for(i=0;i<nEncodeSamples;i++)
   {
      for(j=0;j<nChannels;j++)
      {
         if(*psInBuffer > 0)
            ppfOutBuffer[j][i] = (float)*psInBuffer / 32767.f;
         else
            ppfOutBuffer[j][i] = (float)*psInBuffer / 32768.f;

         *psInBuffer++;
      }
   }
}

void CEncodeVorbis::Convert24to32f(_int24 *pbInBuffer, float **ppfOutBuffer, int nEncodeSamples)
{
   int i, j, nData;
   int nChannels = m_outFormat.nChannels;

   for(i=0;i<nEncodeSamples;i++)
   {
      for(j=0;j<nChannels;j++)
      {
         if(pbInBuffer->b[2] & 0x80)
         {
            nData = int(0xff000000 | (pbInBuffer->b[2] << 16) | (pbInBuffer->b[1] << 8) | pbInBuffer->b[0]);
            ppfOutBuffer[j][i] = (float)nData / 8388608.f;
         }
         else
         {
            nData = int((pbInBuffer->b[2] << 16) | (pbInBuffer->b[1] << 8) | pbInBuffer->b[0]);
            ppfOutBuffer[j][i] = (float)nData / 8388607.f;
         }

         pbInBuffer++;
      }
   }
}

void CEncodeVorbis::Convert32to32f(int *pnInBuffer, float **ppfOutBuffer, int nEncodeSamples)
{
   int i, j;
   int nChannels = m_outFormat.nChannels;

   for(i=0;i<nEncodeSamples;i++)
   {
      for(j=0;j<nChannels;j++)
      {
         if(*pnInBuffer > 0)
            ppfOutBuffer[j][i] = *pnInBuffer / 2147483647.f;
         else
            ppfOutBuffer[j][i] = *pnInBuffer / 2147483648.f;

         *pnInBuffer++;
      }
   }
}

void CEncodeVorbis::Convert32fto32f(float *pfInBuffer, float **ppfOutBuffer, int nEncodeSamples)
{
   int i, j;
   int nChannels = m_outFormat.nChannels;

   for(i=0;i<nEncodeSamples;i++)
   {
      for(j=0;j<nChannels;j++)
      {
         ppfOutBuffer[j][i] = *pfInBuffer;
         *pfInBuffer++;
      }
   }
}

void CEncodeVorbis::Convert64to32f(double *pdInBuffer, float **ppfOutBuffer, int nEncodeSamples)
{
   int i, j;
   int nChannels = m_outFormat.nChannels;

   for(i=0;i<nEncodeSamples;i++)
   {
      for(j=0;j<nChannels;j++)
      {
         ppfOutBuffer[j][i] = (float)*pdInBuffer;
         *pdInBuffer++;
      }
   }
}

void CEncodeVorbis::Analyze(BYTE *pbInBuffer, int nInBufferLength, WAVEFORMATEX *pwf)
{  

   int i, j, n;

   WORD wBytesPerSample = pwf->wBitsPerSample / 8;

   float *pfData;
   float *pfLeft = new float[nInBufferLength / pwf->nBlockAlign];
   float *pfRight = NULL;

   if(pwf->nChannels == 2)
      pfRight = new float[nInBufferLength / pwf->nBlockAlign];

   if(pwf->wFormatTag == WAVE_FORMAT_PCM)
   {
      if(pwf->wBitsPerSample == 16)
      {
         short sData;

         for(n=0;n<pwf->nChannels;n++)
         {
            if(n == 0) pfData = pfLeft;
            else       pfData = pfRight;

            for(i=(wBytesPerSample*n),j=0; i<nInBufferLength; i+=pwf->nBlockAlign,j++)
            {
               ::CopyMemory(&sData, &pbInBuffer[i], 2);
               pfData[j] = (float)sData;
            }
         }
      }
      else if(pwf->wBitsPerSample == 8)
      {
         for(n=0;n<pwf->nChannels;n++)
         {
            if(n == 0) pfData = pfLeft;
            else       pfData = pfRight;

            for(i=(wBytesPerSample*n),j=0; i<nInBufferLength; i+=pwf->nBlockAlign,j++)
            {
               if(pbInBuffer[i] < 128)
                  pfData[j] = ((float)pbInBuffer[i] - 128.0f) * 256.0f;
               else
                  pfData[j] = ((float)pbInBuffer[i] - 128.0f) * 32767.0f / 127.0f;
            }
         }
      }
      else if(pwf->wBitsPerSample == 24)
      {
         for(n=0;n<pwf->nChannels;n++)
         {
            if(n == 0) pfData = pfLeft;
            else       pfData = pfRight;

            for(i=(wBytesPerSample*n),j=0; i<nInBufferLength; i+=pwf->nBlockAlign,j++)
            {
               if(pbInBuffer[i+2] & 0x80)
                  pfData[j] = (float)( ((pbInBuffer[i+2] << 16) | (pbInBuffer[i+1] << 8) | pbInBuffer[i]) - 0x01000000 ) / 256.0f;
               else
                  pfData[j] = (float)( (pbInBuffer[i+2] << 16) | (pbInBuffer[i+1] << 8) | pbInBuffer[i] ) / 8388607.0f * 32767.0f;
            }
         }
      }
      else 
      {
         int nData;

         for(n=0;n<pwf->nChannels;n++)
         {
            if(n == 0) pfData = pfLeft;
            else       pfData = pfRight;

            for(i=(wBytesPerSample*n),j=0; i<nInBufferLength; i+=pwf->nBlockAlign,j++)
            {
               ::CopyMemory(&nData, &pbInBuffer[i], 4);

               if(nData < 0)
                  pfData[j] = (float)nData / 65536.0f;
               else
                  pfData[j] = (float)nData / 2147483647.0f * 32767.0f;
            }
         }
      }
   }
   else 
   {
      if(pwf->wBitsPerSample == 32)
      {
         for(n=0;n<pwf->nChannels;n++)
         {
            if(n == 0) pfData = pfLeft;
            else       pfData = pfRight;

            for(i=(wBytesPerSample*n),j=0; i<nInBufferLength; i+=pwf->nBlockAlign,j++)
            {
               ::CopyMemory(&pfData[j], &pbInBuffer[i], 4);

               if(pfData[j] < 0)
                  pfData[j] *= 32768.0f;
               else
                  pfData[j] *= 32767.0f;
            }
         }
      }
      else 
      {
         double dData;

         for(n=0;n<pwf->nChannels;n++)
         {
            if(n == 0) pfData = pfLeft;
            else       pfData = pfRight;

            for(i=(wBytesPerSample*n),j=0; i<nInBufferLength; i+=pwf->nBlockAlign,j++)
            {
               ::CopyMemory(&dData, &pbInBuffer[i], 8);

               if(dData < 0)
                  pfData[j] = (float)(dData * 32768.0);
               else
                  pfData[j] = (float)(dData * 32767.0);
            }
         }
      }
   }


   
   for(i=0;i<nInBufferLength / pwf->nBlockAlign;i++)
   {
      if(m_dPeak < ::abs(pfLeft[i]))
         m_dPeak = ::abs(pfLeft[i]);
   }

   if(pwf->nChannels == 2)
   {
      for(i=0;i<nInBufferLength / pwf->nBlockAlign;i++)
      {
         if(m_dPeak < ::abs(pfRight[i]))
            m_dPeak = ::abs(pfRight[i]);
      }
   }

   AnalyzeSamples(pfLeft, pfRight, nInBufferLength / pwf->nBlockAlign, pwf->nChannels);

   
   delete []pfLeft;

   if(pwf->nChannels == 2)
      delete []pfRight;
}

HRESULT CEncodeVorbis::OnReceive(HANDLE hFile, IMediaSample *pInSample)
{
   HRESULT hr;
   float **ppfBuffer = NULL;

   int nRestSamples = (int)pInSample->GetActualDataLength() / m_inFormat.nBlockAlign;

   
   BYTE *pbInBuffer = NULL;
   pInSample->GetPointer(&pbInBuffer);

   if(m_bCheckVolume == true || m_bUseTrackgain == true || m_bUseAlbumgain == true)
   {  
      Analyze(pbInBuffer, pInSample->GetActualDataLength(), &m_inFormat);
   }

   if(m_bCheckVolume == true)
   {  
      return S_OK;
   }

   
   ppfBuffer = vorbis_analysis_buffer(&m_vd, nRestSamples);

   
   if(m_inFormat.wBitsPerSample == 16)
      Convert16to32f((short *)pbInBuffer, ppfBuffer, nRestSamples);
   else if(m_inFormat.wBitsPerSample == 8)
      Convert8to32f(pbInBuffer, ppfBuffer, nRestSamples);
   else if(m_inFormat.wBitsPerSample == 24)
      Convert24to32f((_int24 *)pbInBuffer, ppfBuffer, nRestSamples);
   else if(m_inFormat.wBitsPerSample == 32)
   {
      if(m_inFormat.wFormatTag == WAVE_FORMAT_PCM)
         Convert32to32f((int *)pbInBuffer, ppfBuffer, nRestSamples);
      else
         Convert32fto32f((float *)pbInBuffer, ppfBuffer, nRestSamples);
   }
   else
      Convert64to32f((double *)pbInBuffer, ppfBuffer, nRestSamples);

   
   hr = Encode(hFile, nRestSamples);
   if(FAILED(hr))
      return hr;

   return S_OK;
}

HRESULT CEncodeVorbis::OnStop(HANDLE hFile, bool bEos)
{
   if(bEos == false)
      return S_OK;

   if(m_bCheckVolume == true)
   {  
      return S_OK;
   }

   if(m_bInitialize == true)
   {  

      
      Encode(hFile, 0);

      
      ogg_stream_clear(&m_os);
      vorbis_block_clear(&m_vb);
      vorbis_dsp_clear(&m_vd);
      vorbis_comment_clear(&m_vc);
      vorbis_info_clear(&m_vi);

      m_bInitialize = false;
   }

   return S_OK;
}

STDMETHODIMP CEncodeVorbis::OnQueryInterface(REFIID riid, void ** ppv)
{
   if(riid == IID_IEncodeVorbisInterface)
      return GetInterface((IEncodeVorbisInterface *)this, ppv);

   return CBaseFilter::NonDelegatingQueryInterface(riid, ppv);
}

CUnknown * WINAPI CEncodeVorbis::CreateInstance(LPUNKNOWN punk, HRESULT *phr)
{  

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

   return dynamic_cast<CUnknown *>(pNewObject);
}

