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

CEncodeLame::CEncodeLame(IUnknown *pUnk, HRESULT *phr)
   : CSimpleWriter(L"MP3 Encoder", pUnk, CLSID_EncodeLame, phr)
{  

   ::ZeroMemory(&m_inFormat, sizeof(m_inFormat));
   ::ZeroMemory(&m_outFormat, sizeof(m_outFormat));
   ::ZeroMemory(&m_beConfig, sizeof(m_beConfig));

   m_nWait = 0;
   m_llInputSize = 0;

   m_pMP3Buffer = NULL;
   m_hbeStream = NULL;
   m_nMaxBuffer = 0;
   m_nCurrentPos = 0;

   m_outFormat.nMode = ECODECO_ENCODE_MODE_CBR;
   m_outFormat.nQuality = 0;
   m_outFormat.cbr.dwBitrate = 128;
   m_outFormat.vbr.nVBRQuality = 0;
   m_outFormat.vbr.dwMinBitrate =  32;
   m_outFormat.vbr.dwMaxBitrate = 320;
}
CEncodeLame::~CEncodeLame()
{  
   LameRelease();
}
HRESULT CEncodeLame::SetWait(int nWait)
{
   CAutoLock lock(&m_Lock);

   m_nWait = nWait;

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

   *pllBytes = m_llInputSize;

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

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

   return S_OK;
}
HRESULT CEncodeLame::SetOutFormat(void *pFormat)
{
   LAME_CONFIG *pLameConfig = (LAME_CONFIG *)pFormat;

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

   return S_OK;
}
HRESULT CEncodeLame::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->wFormatTag != WAVE_FORMAT_PCM || pwf->wBitsPerSample != 16 || pwf->nChannels > 2 || pwf->nSamplesPerSec > 48000)
   {  
      return S_FALSE;
   }

   ::CopyMemory(&m_inFormat, pwf, sizeof(m_inFormat));
   return S_OK;
}
HRESULT CEncodeLame::CheckOutFormat(WAVEFORMATEX *pInFormat, LAME_CONFIG *pLameConfig)
{  

   if(pInFormat->wFormatTag != WAVE_FORMAT_PCM || pInFormat->wBitsPerSample != 16 ||
      pInFormat->nChannels > 2 || pInFormat->nSamplesPerSec > 48000)
   {  
      return E_FAIL;
   }

   BE_ERR err;
   HINSTANCE hDLL;
   BE_CONFIG beConfig;
   HBE_STREAM hbeStream = NULL;

   DWORD	dwSamples = 0;
   DWORD dwMP3Buffer = 0;

   BEINITSTREAM   beInitStream;
   BECLOSESTREAM  beCloseStream;

   hDLL = ::LoadLibrary(TEXT("lame_enc.dll"));

   if(hDLL == NULL)
      return E_FAIL;

   beInitStream  = (BEINITSTREAM)::GetProcAddress(hDLL, "beInitStream");
   beCloseStream = (BECLOSESTREAM)::GetProcAddress(hDLL, "beCloseStream");

   if(!beInitStream || !beCloseStream)
   {
      ::FreeLibrary(hDLL);
      hDLL = NULL;
      return E_FAIL;
   }

   ::ZeroMemory(&beConfig, sizeof(BE_CONFIG));
   SetConfig(&beConfig, (LAME_CONFIG *)pLameConfig, pInFormat);

   err = beInitStream(&beConfig, &dwSamples, &dwMP3Buffer, &hbeStream);

   if(hbeStream != NULL)
      beCloseStream(hbeStream);

   ::FreeLibrary(hDLL);

	if(err != BE_ERR_SUCCESSFUL)
      return E_FAIL;

   return S_OK;
}
void CEncodeLame::SetConfig(BE_CONFIG *pbeConfig, LAME_CONFIG *pLameConfig, WAVEFORMATEX *pInFormat)
{  

   ::ZeroMemory(pbeConfig, sizeof(BE_CONFIG));

   pbeConfig->dwConfig = BE_CONFIG_LAME;
   pbeConfig->format.LHV1.dwStructVersion = 1;
   pbeConfig->format.LHV1.dwStructSize    = sizeof(BE_CONFIG);
   pbeConfig->format.LHV1.dwReSampleRate  = 0;
   pbeConfig->format.LHV1.dwPsyModel      = 0;
   pbeConfig->format.LHV1.dwEmphasis      = 0;
   pbeConfig->format.LHV1.bPrivate        = 0;
   pbeConfig->format.LHV1.bNoRes          = 0;
   pbeConfig->format.LHV1.bOriginal       = FALSE;
   pbeConfig->format.LHV1.bWriteVBRHeader = FALSE;

   if(pInFormat->nChannels == 2)
      pbeConfig->format.LHV1.nMode = BE_MP3_MODE_JSTEREO;
   else
      pbeConfig->format.LHV1.nMode = BE_MP3_MODE_MONO;

   pbeConfig->format.LHV1.bCRC       = FALSE;
   pbeConfig->format.LHV1.bCopyright = FALSE;
   pbeConfig->format.LHV1.nQuality   = (WORD)pLameConfig->nQuality;

   if(pLameConfig->nMode == ECODECO_ENCODE_MODE_CBR)
   {
      pbeConfig->format.LHV1.nPreset = LQP_CBR;
      pbeConfig->format.LHV1.bEnableVBR = FALSE;
      pbeConfig->format.LHV1.bWriteVBRHeader = FALSE;
      pbeConfig->format.LHV1.nVbrMethod = VBR_METHOD_DEFAULT;
      pbeConfig->format.LHV1.dwBitrate = pLameConfig->cbr.dwBitrate;
   }
   else
   {  

      pbeConfig->format.LHV1.nPreset = LQP_NOPRESET;
      pbeConfig->format.LHV1.bEnableVBR = TRUE;
      pbeConfig->format.LHV1.nVbrMethod = VBR_METHOD_DEFAULT;
      pbeConfig->format.LHV1.bWriteVBRHeader = TRUE;
      pbeConfig->format.LHV1.dwBitrate = pLameConfig->vbr.dwMinBitrate;
      pbeConfig->format.LHV1.dwMaxBitrate = pLameConfig->vbr.dwMaxBitrate;
      pbeConfig->format.LHV1.nVBRQuality = pLameConfig->vbr.nVBRQuality;
   }

   pbeConfig->format.LHV1.dwSampleRate = pInFormat->nSamplesPerSec;
   pbeConfig->format.LHV1.dwReSampleRate = pInFormat->nSamplesPerSec;

   if(pbeConfig->format.LHV1.dwSampleRate > 24000)
      pbeConfig->format.LHV1.dwMpegVersion = MPEG1;
   else
      pbeConfig->format.LHV1.dwMpegVersion = MPEG2;
}
void CEncodeLame::LameRelease(void)
{
   if(m_hbeStream != NULL)
   {
      beCloseStream(m_hbeStream);
      m_hbeStream = NULL;
   }

   if(m_pMP3Buffer != NULL)
   {
      delete []m_pMP3Buffer;
      m_pMP3Buffer = NULL;
   }
}
HRESULT CEncodeLame::OnStart(HANDLE hFile)
{
   LameRelease();

   m_nCurrentPos = 0;

   m_hDLL = ::LoadLibrary(TEXT("lame_enc.dll"));

   if(m_hDLL == NULL)
      goto error;

   beInitStream     = (BEINITSTREAM)    ::GetProcAddress(m_hDLL, "beInitStream");
   beEncodeChunk    = (BEENCODECHUNK)   ::GetProcAddress(m_hDLL, "beEncodeChunk");
   beDeinitStream   = (BEDEINITSTREAM)  ::GetProcAddress(m_hDLL, "beDeinitStream");
   beCloseStream    = (BECLOSESTREAM)   ::GetProcAddress(m_hDLL, "beCloseStream");
   beWriteVBRHeader = (BEWRITEVBRHEADER)::GetProcAddress(m_hDLL, "beWriteVBRHeader");
   beWriteInfoTag   = (BEWRITEINFOTAG)  ::GetProcAddress(m_hDLL, "beWriteInfoTag");

   if(!beInitStream || !beEncodeChunk || !beDeinitStream || !beCloseStream || !beWriteVBRHeader)
      goto error;

   HRESULT hr;
   ALLOCATOR_PROPERTIES InProps;
   IMemAllocator * pInAlloc = NULL;

   hr = m_pInput->GetAllocator(&pInAlloc);
   if(FAILED(hr)) return hr;

   hr = pInAlloc->GetProperties(&InProps);
   SAFE_RELEASE(pInAlloc);

   int nInBufferSize = (int)InProps.cbBuffer;

   BE_ERR err;
   DWORD	dwSamples = 0;
   DWORD dwMP3BufferSize = 0;

   SetConfig(&m_beConfig, &m_outFormat, &m_inFormat);

   err = beInitStream(&m_beConfig, &dwSamples, &dwMP3BufferSize, &m_hbeStream);

   if(err != BE_ERR_SUCCESSFUL)
      return E_FAIL;

   m_nMaxBuffer = dwSamples * 2;

   m_pMP3Buffer = new BYTE[dwMP3BufferSize];

   return S_OK;

error:

   LameRelease();
   return E_FAIL;
}
HRESULT CEncodeLame::OnReceive(HANDLE hFile, IMediaSample *pInSample)
{
   BE_ERR err;
   int nRest, nUseData;

   if(m_nWait > 0)
   {
      int nWait = int((m_llInputSize + pInSample->GetActualDataLength()) * 10 / m_inFormat.nAvgBytesPerSec) - int(m_llInputSize * 10 / m_inFormat.nAvgBytesPerSec);

      if(nWait > 0)
         ::Sleep(nWait * m_nWait * 2);
   }

   m_llInputSize += pInSample->GetActualDataLength();

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

   nRest = pInSample->GetActualDataLength();

   while(0 < nRest)
   {  

      if(nRest < m_nMaxBuffer)
         nUseData = nRest;
      else
         nUseData = m_nMaxBuffer;

      DWORD dwWrite = 0;

      err = beEncodeChunk(m_hbeStream, nUseData / 2, (short *)pbInBuffer, m_pMP3Buffer, &dwWrite); 
      if(err != BE_ERR_SUCCESSFUL)
         return E_FAIL;

      if(dwWrite > 0)
      {
         BOOL flag;
         DWORD dwWritten;

         flag = ::WriteFile(hFile, (PVOID)m_pMP3Buffer, dwWrite, &dwWritten, NULL);
         if(flag == FALSE)
            return E_FAIL;
      }

      nRest -= nUseData;

      if(0 < nRest)
         pbInBuffer += nUseData;
   }

   return S_OK;
}
HRESULT CEncodeLame::OnStop(HANDLE hFile, bool bEos)
{
   if(bEos == true)
   {
      BE_ERR err;
      DWORD dwWrite;

      err = beDeinitStream(m_hbeStream, m_pMP3Buffer, &dwWrite);

      if(err != BE_ERR_SUCCESSFUL)
         return S_FALSE;

      if(dwWrite > 0)
      {
         BOOL flag;
         DWORD dwWritten;

         flag = ::WriteFile(hFile, (PVOID)m_pMP3Buffer, dwWrite, &dwWritten, NULL);
         if(flag == FALSE) return S_FALSE;
      }

      LameRelease();

      if(m_beConfig.format.LHV1.bEnableVBR == TRUE)
      {  

         ::CloseHandle(m_hFile);
         m_hFile = NULL;

         char szFileName[MAX_PATH];
         ::WideCharToMultiByte(CP_ACP, 0, m_awInputFileName, -1, szFileName, MAX_PATH, NULL, NULL);
         beWriteVBRHeader(szFileName);
      }
   }

   return S_OK;
}
void CEncodeLame::WriteHeader(WCHAR *pwFileName)
{  

   BE_ERR err;

   if(m_beConfig.format.LHV1.bWriteVBRHeader == TRUE)
   {
      char szFileName[MAX_PATH];
      ::WideCharToMultiByte(CP_ACP, 0, pwFileName, -1, szFileName, MAX_PATH, NULL, NULL);

      HINSTANCE hDLL = NULL;
      BEWRITEVBRHEADER beWriteVBRHeader;

      hDLL = ::LoadLibrary(TEXT("lame_enc.dll"));
      if(hDLL == NULL)
         return;

      beWriteVBRHeader = (BEWRITEVBRHEADER)::GetProcAddress(hDLL, "beWriteVBRHeader");

      if(!beWriteVBRHeader)
      {
         ::FreeLibrary(hDLL);
         hDLL = NULL;
         return;
      }

      err = beWriteVBRHeader(szFileName);

      ::FreeLibrary(hDLL);
      hDLL = NULL;
	}
}
STDMETHODIMP CEncodeLame::OnQueryInterface(REFIID riid, void ** ppv)
{
   if(riid == IID_IEcoDecoInterface)
      return GetInterface((IEcoDecoInterface *)this, ppv);
   else if(riid == IID_IEncodeLameInterface)
      return GetInterface((IEncodeLameInterface*)this, ppv);

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

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

   return dynamic_cast<CUnknown *>(pNewObject);
}
