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

#define WAVE_FORMAT_IEEE_FLOAT 3;

CDecodeWv::CDecodeWv(IUnknown *pUnk, HRESULT *phr)
   : CSimpleTransform(L"Wv Decoder", pUnk, CLSID_DecodeWv, phr)
{  

   ::ZeroMemory(&m_outFormat, sizeof(m_outFormat));
   m_pBufferDecoder = NULL;
   m_pFrameBuffer = NULL;
   m_pnOutBuffer = NULL;
   m_rtCurrent = 0;
   m_nDecodeVersion = 0x0407;
}
CDecodeWv::~CDecodeWv()  
{  

   if(m_pnOutBuffer != NULL)
   {
      delete []m_pnOutBuffer;
      m_pnOutBuffer = NULL;
   }

   if(m_pFrameBuffer != NULL)
   {
      frame_buffer_free(m_pFrameBuffer);
      m_pFrameBuffer = NULL;
   }

   if(m_pBufferDecoder != NULL)
   {
      wavpack_buffer_decoder_free(m_pBufferDecoder);
      m_pBufferDecoder = NULL;
   }
}  
HRESULT CDecodeWv::OnStart(void)
{  
   m_rtCurrent = 0;
   return S_OK;
}
HRESULT CDecodeWv::OnSeek(void)
{  
   m_rtCurrent = 0;
   return S_OK;
}
HRESULT CDecodeWv::OnTransform(IMediaSample *pInSample, IMediaSample *pOutSample)
{  

   BYTE *pbInBuffer = NULL;
   BYTE *pbOutBuffer = NULL;

   pInSample->GetPointer(&pbInBuffer);

   pOutSample->GetPointer(&pbOutBuffer);

   int i, j, nSampleSize, nUseArraySize, nResult;
   common_frame_data commonFrameData;

   nResult = reconstruct_wavpack_frame(m_pFrameBuffer, &commonFrameData, (char *)pbInBuffer, (uint32_t)pInSample->GetActualDataLength(), true, 0, m_nDecodeVersion);

   if(nResult <= 0)
   {
      pOutSample->SetActualDataLength(0);
      return S_OK;
   }

   if(m_pnOutBuffer == NULL)
   {  
      m_pnOutBuffer = new int[nResult * m_outFormat.nChannels];
   }

   nResult = wavpack_buffer_decoder_load_frame(m_pBufferDecoder, m_pFrameBuffer->data, m_pFrameBuffer->len, NULL, m_pFrameBuffer->len);

   nSampleSize = wavpack_buffer_decoder_unpack(m_pBufferDecoder, m_pnOutBuffer, commonFrameData.block_samples);

   if(nSampleSize <= 0)
   {
      pOutSample->SetActualDataLength(0);
      return S_OK;
   }

   nUseArraySize = nSampleSize * m_outFormat.nChannels;

   if(m_outFormat.wBitsPerSample == 16)
   {
      short *ps = (short *)pbOutBuffer;

      for(i=0;i<nUseArraySize;i++)
         ps[i] = m_pnOutBuffer[i];
   }
   else if(m_outFormat.wBitsPerSample == 8)
   {
      for(i=0;i<nUseArraySize;i++)
         pbOutBuffer[i] = (BYTE)(m_pnOutBuffer[i] - 128);
   }
   else if(m_outFormat.wBitsPerSample == 24)
   {
      for(i=0,j=0; i<nUseArraySize; i++,j+=3)
      {
         pbOutBuffer[j]   = (BYTE)(m_pnOutBuffer[i] & 0x000000ff);
         pbOutBuffer[j+1] = (BYTE)((m_pnOutBuffer[i] & 0x0000ff00) >> 8);
         pbOutBuffer[j+2] = (BYTE)((m_pnOutBuffer[i] & 0x00ff0000) >> 16);
      }
   }
   else 
   {
      if(WavpackGetMode(m_pBufferDecoder->wpc) & MODE_FLOAT)
      {  

         int nData;
         float fData;

         for(i=0;i<nUseArraySize;i++)
         {
            ::CopyMemory(&fData, &m_pnOutBuffer[i], 4);

            nData = (int)(fData * 0x7fffffff);

            ::CopyMemory(pbOutBuffer, &nData, 4);
            pbOutBuffer+=4;
         }
      }
      else
         ::CopyMemory(pbOutBuffer, m_pnOutBuffer, nUseArraySize * 4);
   }

   int nOutLength = nUseArraySize * (m_outFormat.wBitsPerSample / 8);

   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 CDecodeWv::OnConnectInPin(const CMediaType *pmtIn)
{  

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

   if(pmtIn->subtype != MEDIASUBTYPE_WAVPACK4)
      return S_FALSE;

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

   if(m_pnOutBuffer != NULL)
   {
      delete []m_pnOutBuffer;
      m_pnOutBuffer = NULL;
   }

   if (m_pFrameBuffer != NULL)
   {
      frame_buffer_free(m_pFrameBuffer);
      m_pFrameBuffer = NULL;
   }

   if(m_pBufferDecoder != NULL)
   {
      wavpack_buffer_decoder_free(m_pBufferDecoder);
      m_pBufferDecoder = NULL;
   }

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

   if(pwf->cbSize < 2)
      goto error;

   short s;
   ::CopyMemory(&s, pmtIn->Format() + sizeof(WAVEFORMATEX), 2);
   m_nDecodeVersion = s;

   m_pBufferDecoder = wavpack_buffer_decoder_new();
   if(m_pBufferDecoder == NULL)
      goto error;

   m_pFrameBuffer = frame_buffer_new();
   if(m_pFrameBuffer == NULL)
      goto error;

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

   m_outFormat.cbSize = 0;
   m_outFormat.wFormatTag = WAVE_FORMAT_PCM;

   return S_OK;

error:

   if(m_pnOutBuffer != NULL)
   {
      delete []m_pnOutBuffer;
      m_pnOutBuffer = NULL;
   }

   if (m_pFrameBuffer != NULL)
   {
      frame_buffer_free(m_pFrameBuffer);
      m_pFrameBuffer = NULL;
   }

   if(m_pBufferDecoder != NULL)
   {
      wavpack_buffer_decoder_free(m_pBufferDecoder);
      m_pBufferDecoder = NULL;
   }

   return S_FALSE;
}
HRESULT CDecodeWv::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 = nInBufferSize * 2;

   return S_OK;
}
CUnknown * WINAPI CDecodeWv::CreateInstance(LPUNKNOWN punk, HRESULT *phr)
{  

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

   return dynamic_cast<CUnknown *>(pNewObject);
}


void wavpack_buffer_format_samples(wavpack_buffer_decoder* wbd, uchar *dst, long *src, uint32_t samples)
{
    long temp;
    int bps = WavpackGetBytesPerSample(wbd->wpc);
    uint32_t samcnt = samples * WavpackGetNumChannels (wbd->wpc);

    switch (bps) {
        
    case 1:
        while (samcnt--)
            *dst++ = (uchar)(*src++ + 128);
        
        break;
        
    case 2:
        while (samcnt--) {
            *dst++ = (uchar)(temp = *src++);
            *dst++ = (uchar)(temp >> 8);
        }
        
        break;
        
    case 3:
        while (samcnt--) {
            *dst++ = (uchar)(temp = *src++);
            *dst++ = (uchar)(temp >> 8);
            *dst++ = (uchar)(temp >> 16);
        }
        
        break;
    }
}

int32_t frame_stream_reader_read_bytes(void *id, void *data, int32_t bcount)
{
    uint32_t bytes_left = 0;
    uint32_t bytes_to_read = 0;
    frame_stream_reader* fsr = (frame_stream_reader*)id;
    if(fsr->buffer == NULL)
    {
        return 0;
    }
    bytes_left = (fsr->length - fsr->position);
    bytes_to_read = min(bytes_left, (uint32_t)bcount);
    
    if(bytes_to_read > 0)
    {    
        wp_memcpy(data, fsr->buffer + fsr->position, bytes_to_read);
        fsr->position += bytes_to_read;
    }

    return bytes_to_read;
}


uint32_t frame_stream_reader_get_pos(void *id)
{
    frame_stream_reader* fsr = (frame_stream_reader*)id;
    return fsr->position;
}


int frame_stream_reader_set_pos_abs(void *id, uint32_t pos)
{
    frame_stream_reader* fsr = (frame_stream_reader*)id;
    fsr->position = min(pos, fsr->length);
    if(pos > fsr->length)
    {
        return -1;
    }
    else
    {
        return 0;
    }
}


int frame_stream_reader_set_pos_rel(void *id, int32_t delta, int mode)
{
    frame_stream_reader* fsr = (frame_stream_reader*)id;
    uint32_t new_pos = 0;
    switch(mode)
    {
        case SEEK_SET:
            new_pos = delta;            
            break;      
        case SEEK_CUR:
            new_pos = fsr->position + delta;
            break;
        case SEEK_END:
            new_pos = fsr->length + delta;
            break;
    }
    fsr->position = constrain(0, new_pos, fsr->length);

    if((new_pos < 0) || (new_pos > fsr->length))
    {
        return -1;
    }
    else
    {
        return 0;
    }
}


int frame_stream_reader_push_back_byte(void *id, int c)
{
    frame_stream_reader* fsr = (frame_stream_reader*)id;
    fsr->position = constrain(0, fsr->position - 1, fsr->length);
    return fsr->position;
}


uint32_t frame_stream_reader_get_length(void *id)
{
    frame_stream_reader* fsr = (frame_stream_reader*)id;
    return fsr->length;
}


int frame_stream_reader_can_seek(void *id)
{
    return 0;
}


frame_stream_reader* frame_stream_reader_new()
{
   frame_stream_reader *fsr = (frame_stream_reader *)wp_alloc(sizeof(frame_stream_reader));
    if(!fsr)
    {
        return NULL;
    }
    
    fsr->sr.read_bytes = frame_stream_reader_read_bytes;
    fsr->sr.get_pos = frame_stream_reader_get_pos;
    fsr->sr.set_pos_abs = frame_stream_reader_set_pos_abs;
    fsr->sr.set_pos_rel = frame_stream_reader_set_pos_rel;
    fsr->sr.push_back_byte = frame_stream_reader_push_back_byte;
    fsr->sr.get_length = frame_stream_reader_get_length;
    fsr->sr.can_seek = frame_stream_reader_can_seek;

    fsr->buffer = 0;
    fsr->position = 0;

    return fsr;
}


void frame_stream_reader_free(frame_stream_reader* fsr)
{
    if(fsr != NULL)
    {
        wp_free(fsr);
    }
}


void frame_stream_reader_set_buffer(void *id, char* buffer, int length)
{
    frame_stream_reader* fsr = (frame_stream_reader*)id;
    fsr->buffer = buffer;
    fsr->length = length;
    fsr->position = 0;
}


wavpack_buffer_decoder* wavpack_buffer_decoder_new()
{
   wavpack_buffer_decoder *wbd = (wavpack_buffer_decoder *)wp_alloc(sizeof(wavpack_buffer_decoder));
    if(!wbd)
    {
        return NULL;
    }

    wbd->fsr = frame_stream_reader_new();
    if(!wbd->fsr)
    {
        return NULL;
    }

    wbd->fsrc = frame_stream_reader_new();
    if(!wbd->fsrc)
    {
        return NULL;
    }

    wbd->wpc = NULL;

    return wbd;
}


int wavpack_buffer_decoder_load_frame(wavpack_buffer_decoder* wbd,
                                      char* data, int length,
                                      char* correction_data, int cd_length)
{
    frame_stream_reader_set_buffer(wbd->fsr, data, length);
    
    if(correction_data != 0)
    {
        frame_stream_reader_set_buffer(wbd->fsrc, correction_data, cd_length);
    }

    if(wbd->wpc == NULL)
    {
        wbd->wpc = WavpackOpenFileInputEx(
            (WavpackStreamReader*)wbd->fsr,
            wbd->fsr,
            (correction_data != NULL) ? wbd->fsrc : NULL,
            wbd->wavpack_error_msg, OPEN_STREAMING|OPEN_NORMALIZE, 0);
    }

    return (wbd->wpc != NULL);
}


uint32_t wavpack_buffer_decoder_unpack(wavpack_buffer_decoder* wbd,
                                       int32_t* buffer, uint32_t samples)
{
    return WavpackUnpackSamples (wbd->wpc, buffer, samples);
}


void wavpack_buffer_decoder_free(wavpack_buffer_decoder* wbd)
{
    if(wbd != NULL)
    {   
        frame_stream_reader_free(wbd->fsr);
        wbd->fsr = NULL;
        frame_stream_reader_free(wbd->fsrc);
        wbd->fsrc = NULL;
        if(wbd->wpc != NULL)
        {
            WavpackCloseFile(wbd->wpc);
            wbd->wpc = NULL;
        }
        wp_free(wbd);
    }
}


frame_buffer* frame_buffer_new()
{
    frame_buffer* fb = (frame_buffer *)wp_alloc(sizeof(frame_buffer));
    if(!fb)
    {
        return NULL;
    }

    frame_reserve_space(fb, 1024);

    return fb;
}


void frame_buffer_free(frame_buffer* fb)
{
    if(fb->data != NULL)
    {
        wp_free(fb->data);
        fb->data = NULL;
    }
    wp_free(fb);
}


int frame_reserve_space(frame_buffer* dst, int len)
{
    if(dst->data == 0)
    {
        dst->data = (char *)wp_alloc(len);
        dst->total_len = len;
    }
    else if(dst->len + len > dst->total_len)
    {
        dst->data = (char *)wp_realloc(dst->data, dst->len + len);        
        dst->total_len = dst->len + len;
    }

    return (dst->data != NULL);
}


int frame_append_data(frame_buffer* dst, char* src, int len)
{
    if(!frame_reserve_space(dst, len))
    {
        return -1;
    }

    wp_memcpy(dst->data + dst->pos, src, len);
    dst->pos += len;
    dst->len += len;

    return 0;
}


int frame_append_data2(frame_buffer* dst, WavpackStreamReader *io, int len)
{
    frame_reserve_space(dst, len);
    
    if(io->read_bytes(io, (dst->data + dst->pos), len) != len)
    {
        return -1;
    }    

    dst->pos += len;
    dst->len += len;

    return 0;
}


void frame_reset(frame_buffer* dst)
{
  dst->pos = 0;
  dst->len = 0;
  dst->nb_block = 0;
}


int reconstruct_wavpack_frame(
    frame_buffer *frame,
    common_frame_data *common_data,
    char *pSrc,
    uint32_t SrcLength,
    int is_main_frame,
    int several_blocks,
    int version)
{
    WavpackHeader header;
    int i = 0;
    uint32_t stripped_header_len = 0;
    uint32_t current_flag = 0;
    uint32_t current_crc = 0;
    uint32_t block_size = 0;
    char *pBlock = pSrc;

    frame_reset(frame);

    if(is_main_frame == TRUE)
    {
      wp_memclear(common_data, sizeof(common_frame_data));
    }

    while(pBlock < pSrc + SrcLength)
    {
        i = 0;
        if(is_main_frame == TRUE)
        {
            if(pBlock == pSrc)
            {
                common_data->block_samples = ((uint32_t*)pBlock)[i++];
            }
            current_flag = ((uint32_t*)pBlock)[i++];
            common_data->array_flags[frame->nb_block] = current_flag;
        }
        else
        {
          current_flag = common_data->array_flags[frame->nb_block];
        }

        current_crc = ((uint32_t*)pBlock)[i++];

        if(several_blocks == TRUE)
        {
            block_size = ((uint32_t*)pBlock)[i++];
        }
        else
        {
            block_size = SrcLength - (i * sizeof(uint32_t));
        }
        stripped_header_len = (i * sizeof(uint32_t));

        wp_memclear(&header, sizeof(WavpackHeader));
        wp_memcpy(header.ckID, "wvpk", 4);
        header.ckSize = block_size + sizeof(WavpackHeader) - 8;
        header.version = version;
        header.total_samples = -1;
        header.block_samples = common_data->block_samples;
        header.flags = current_flag;
        header.crc = current_crc;

        frame_append_data(frame, (char*)&header, sizeof(WavpackHeader));
        frame_append_data(frame, pBlock + stripped_header_len, block_size);

        frame->nb_block++;

        pBlock += (stripped_header_len + block_size);
    }

    return frame->len;
}


int strip_wavpack_block(frame_buffer *frame,
                        WavpackHeader *wphdr,
                        WavpackStreamReader *io,
                        uint32_t block_data_size,
                        int is_main_frame,
                        int several_blocks)
{
    char* pDst = NULL;
    uint32_t stripped_header_len = 0;
    uint32_t total_len = 0;

    if(is_main_frame == TRUE)
    {
        if(frame->nb_block == 0)
        {
            frame_append_data(frame, (char*)&wphdr->block_samples, sizeof(wphdr->block_samples));
        }        
        frame_append_data(frame, (char*)&wphdr->flags, sizeof(wphdr->flags));
    }
    frame_append_data(frame, (char*)&wphdr->crc, sizeof(wphdr->crc));

    if(several_blocks == TRUE)
    {
        frame_append_data(frame, (char*)&block_data_size, sizeof(block_data_size));
    }
    
    stripped_header_len = frame->len;

    if(frame_append_data2(frame, io, block_data_size) == -1)
    {
        return -1;
    }

    frame->nb_block++;
        
    return stripped_header_len;
}
