/*
 *  wtv2mp3 converter
 *
 *  Copyright (c) 2008 Miles Benson
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library 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
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 59 Temple Place - Suite 330,
 * Boston, MA 02111-1307, USA.
 */

//Don't forget to change project settings:
//1. C++: add include path to DirectShow include folder (such as c:\dxsdk\include)
//2. Link: add link path to DirectShow lib folder (such as c:\dxsdk\lib).
//3. Link: add strmiids.lib and quartz.lib

#include "stdafx.h"
#include <DShow.h>
#include <atlbase.h>
#include <initguid.h>
#include <dvdmedia.h>

#include "iaudioprops.h"

BOOL hrcheck(HRESULT hr, const char* errtext)
{
    if (hr >= S_OK)
        return FALSE;
    TCHAR szErr[MAX_ERROR_TEXT_LEN];
    DWORD res = AMGetErrorText(hr, szErr, MAX_ERROR_TEXT_LEN);
    if (res)
        {
        printf("Error %x: %s\n",hr, errtext);
        wprintf(L"%s\n", szErr);
        }
    else
        printf("Error %x: %s\n", hr, errtext);
    return TRUE;
}

//change this macro to fit your style of error handling
#define CHECK_HR(hr, msg) if (hrcheck(hr, msg)) return hr;

CComPtr<IBaseFilter> CreateFilter(WCHAR* displayName)
{
    CComPtr<IBindCtx> pBindCtx;
    HRESULT hr = CreateBindCtx(0, &pBindCtx);
    if (hrcheck(hr, "Can't create bind context"))
        return NULL;

    ULONG chEaten = 0;
    CComPtr<IMoniker> pMoniker;
    hr = MkParseDisplayName(pBindCtx, displayName, &chEaten, &pMoniker);
    if (hrcheck(hr, "Can't create parse display name of the filter"))
        return NULL;

    CComPtr<IBaseFilter> pFilter;
    if (SUCCEEDED(hr))
    {
        hr = pMoniker->BindToObject(pBindCtx, NULL, IID_IBaseFilter, (void**)&pFilter);
        if (hrcheck(hr, "Can't bind moniker to filter object"))
            return NULL;
    }
    return pFilter;
}

CComPtr<IPin> GetPin(IBaseFilter *pFilter, LPCOLESTR pinname)
{
    CComPtr<IEnumPins>  pEnum;
    CComPtr<IPin>       pPin;
    // wprintf(L"Looking for pinname [%s]\n", pinname);
    HRESULT hr = pFilter->EnumPins(&pEnum);
    if (hrcheck(hr, "Can't enumerate pins."))
        return NULL;

    while(pEnum->Next(1, &pPin, 0) == S_OK)
    {
        PIN_INFO pinfo;
        pPin->QueryPinInfo(&pinfo);
        BOOL found = !_wcsicmp(pinname, pinfo.achName);
        if (pinfo.pFilter) pinfo.pFilter->Release();
        if (found)
            return pPin;
        pPin.Release();
    }
    printf("Pin not found!\n");
    return NULL;  
}

// {C9F5FE02-F851-4EB5-99EE-AD602AF1E619}
DEFINE_GUID(CLSID_StreamBufferSource,
0xC9F5FE02, 0xF851, 0x4EB5, 0x99, 0xEE, 0xAD, 0x60, 0x2A, 0xF1, 0xE6, 0x19); //sbe.dll


// {E1F1A0B8-BEEE-490D-BA7C-066C40B5E2B9}
DEFINE_GUID(CLSID_MicrosoftDTVDVDAudioDecoder,
0xE1F1A0B8, 0xBEEE, 0x490D, 0xBA, 0x7C, 0x06, 0x6C, 0x40, 0xB5, 0xE2, 0xB9); //msmpeg2adec.dll


// {B8D27088-FF5F-4B7C-98DC-0E91A1696286}
DEFINE_GUID(CLSID_LAMEAudioEncoder,
0xB8D27088, 0xFF5F, 0x4B7C, 0x98, 0xDC, 0x0E, 0x91, 0xA1, 0x69, 0x62, 0x86); //lame.ax



HRESULT BuildGraph(IGraphBuilder *pGraph, LPCOLESTR srcFile1, LPCOLESTR dstFile1, DWORD dwBitRate, DWORD dwQuality)
{
    HRESULT hr = S_OK;

    //graph builder
    CComPtr<ICaptureGraphBuilder2> pBuilder;
    hr = pBuilder.CoCreateInstance(CLSID_CaptureGraphBuilder2);
    CHECK_HR(hr, "Can't create Capture Graph Builder");
    hr = pBuilder->SetFiltergraph(pGraph);
    CHECK_HR(hr, "Can't SetFiltergraph");

    //add StreamBufferSource
    CComPtr<IBaseFilter> pStreamBufferSource;
    hr = pStreamBufferSource.CoCreateInstance(CLSID_StreamBufferSource);
    CHECK_HR(hr, "Can't create StreamBufferSource");
    hr = pGraph->AddFilter(pStreamBufferSource, L"StreamBufferSource");
    CHECK_HR(hr, "Can't add StreamBufferSource to graph");
    //set source filename
    CComQIPtr<IFileSourceFilter, &IID_IFileSourceFilter> pStreamBufferSource_src(pStreamBufferSource);
    if (!pStreamBufferSource_src)
        CHECK_HR(E_NOINTERFACE, "Can't get IFileSourceFilter");
    hr = pStreamBufferSource_src->Load(srcFile1, NULL);
    CHECK_HR(hr, "Can't load file");


    //add PBDA DTFilter
    CComPtr<IBaseFilter> pPBDADTFilter = CreateFilter(L"@device:sw:{4A56AF32-C21F-11DB-96FA-005056C00008}\\PBDA DTFilter");
    hr = pGraph->AddFilter(pPBDADTFilter, L"PBDA DTFilter");
    CHECK_HR(hr, "Can't add PBDA DTFilter to graph");


    //connect StreamBufferSource and PBDA DTFilter
    hr = pGraph->ConnectDirect(GetPin(pStreamBufferSource, L"DVR Out - 1"), GetPin(pPBDADTFilter, L"In(Enc/Tag)"), NULL);
    CHECK_HR(hr, "Can't connect StreamBufferSource and PBDA DTFilter");


    //add Microsoft DTV-DVD Audio Decoder
    CComPtr<IBaseFilter> pMicrosoftDTVDVDAudioDecoder;
    hr = pMicrosoftDTVDVDAudioDecoder.CoCreateInstance(CLSID_MicrosoftDTVDVDAudioDecoder);
    CHECK_HR(hr, "Can't create Microsoft DTV-DVD Audio Decoder");
    hr = pGraph->AddFilter(pMicrosoftDTVDVDAudioDecoder, L"Microsoft DTV-DVD Audio Decoder");
    CHECK_HR(hr, "Can't add Microsoft DTV-DVD Audio Decoder to graph");


    //connect PBDA DTFilter and Microsoft DTV-DVD Audio Decoder
    hr = pGraph->ConnectDirect(GetPin(pPBDADTFilter, L"Out"), GetPin(pMicrosoftDTVDVDAudioDecoder, L"XForm In"), NULL);
    CHECK_HR(hr, "Can't connect PBDA DTFilter and Microsoft DTV-DVD Audio Decoder");


    //add LAME Audio Encoder
    CComPtr<IBaseFilter> pLAMEAudioEncoder;
    hr = pLAMEAudioEncoder.CoCreateInstance(CLSID_LAMEAudioEncoder);
    CHECK_HR(hr, "Can't create LAME Audio Encoder");
    
    IAudioEncoderProperties *pIP;
    
    hr = pLAMEAudioEncoder->QueryInterface(IID_IAudioEncoderProperties, (void **)&pIP);
    CHECK_HR(hr, "Can't access LAME Audio Encoder Properties");
    hr = pIP->set_Bitrate(dwBitRate);
    CHECK_HR(hr, "Can't set the bitrate");
    hr = pIP->set_Quality(dwQuality);
    CHECK_HR(hr, "Can't set the quality");
    pIP->Release();
    
    
    hr = pGraph->AddFilter(pLAMEAudioEncoder, L"LAME Audio Encoder");
    CHECK_HR(hr, "Can't add LAME Audio Encoder to graph");


    //connect Microsoft DTV-DVD Audio Decoder and LAME Audio Encoder
    hr = pGraph->ConnectDirect(GetPin(pMicrosoftDTVDVDAudioDecoder, L"XFrom Out"), GetPin(pLAMEAudioEncoder, L"XForm In"), NULL);
    CHECK_HR(hr, "Can't connect Microsoft DTV-DVD Audio Decoder and LAME Audio Encoder");


    //add File writer
    CComPtr<IBaseFilter> pFilewriter;
    hr = pFilewriter.CoCreateInstance(CLSID_FileWriter);
    CHECK_HR(hr, "Can't create File writer");
    hr = pGraph->AddFilter(pFilewriter, L"File writer");
    CHECK_HR(hr, "Can't add File writer to graph");
    //set destination filename
    CComQIPtr<IFileSinkFilter, &IID_IFileSinkFilter> pFilewriter_sink(pFilewriter);
    if (!pFilewriter_sink)
        CHECK_HR(E_NOINTERFACE, "Can't get IFileSinkFilter");
    hr = pFilewriter_sink->SetFileName(dstFile1, NULL);
    CHECK_HR(hr, "Can't set filename");


    //connect LAME Audio Encoder and File writer
    hr = pGraph->ConnectDirect(GetPin(pLAMEAudioEncoder, L"XForm Out"), GetPin(pFilewriter, L"in"), NULL);
    CHECK_HR(hr, "Can't connect LAME Audio Encoder and File writer");


    return S_OK;
}

int _tmain(int argc, _TCHAR* argv[]) //use this line in VS2008
{
    USES_CONVERSION;

    int iSuccess = 0;
    
    // let's do some very quick and dirty argument parsing
    // more LAME options could be added, but this is what
    // I think is necessary for now.
    //
    // -b 192 source dest
    // source dest
    //
    
    _TCHAR *pSource = NULL;
    _TCHAR *pDest   = NULL;
    DWORD dwBitRate = 128;
    DWORD dwQuality = 5;
    for (int i = 1; i < argc; i++)
        {
        if (wcscmp(argv[i], L"-b") == 0)
            {
            if (i == argc - 1)
                break;
            i++;
            _TCHAR *pEnd;
            dwBitRate = wcstoul(argv[i], &pEnd, 10);
            if (*pEnd)
                dwBitRate = 0;
            }
        else if (wcscmp(argv[i], L"-q") == 0)
            {
            if (i == argc - 1)
                break;
            i++;
            _TCHAR *pEnd;
            dwQuality = wcstoul(argv[i], &pEnd, 10);
            if (*pEnd)
                dwQuality = 9999;
            }
        else if (pSource == NULL)
            pSource = argv[i];
        else if (pDest == NULL)
            pDest = argv[i];
        }
    
    if (pSource == NULL ||
        pDest   == NULL ||
        dwBitRate == 0  ||
        dwQuality < 0 || dwQuality > 9)
        {
        printf("Usage is <sourcefile.wtv> <destfile.mp3>\n");
        printf("  or     -b <bitrate> -q <quality> <sourcefile.wtv> <destfile.mp3>\n");
        printf("  where bitrate is eg 128 or 192\n");
        printf("  and quality is 0...9\n");
        exit(iSuccess ? 0 : 1);
        }

    CoInitialize(NULL);
    CComPtr<IGraphBuilder> graph;
    graph.CoCreateInstance(CLSID_FilterGraph);

    HRESULT hr = BuildGraph(graph, T2W(pSource), T2W(pDest), dwBitRate, dwQuality);
    if (hr==S_OK) {
        printf("Running");
        CComQIPtr<IMediaControl, &IID_IMediaControl> mediaControl(graph);
        hr = mediaControl->Run();
        CHECK_HR(hr, "Can't run the graph");
        CComQIPtr<IMediaEvent, &IID_IMediaEvent> mediaEvent(graph);
        BOOL stop = FALSE;
        while(!stop) 
        {
            long ev=0, p1=0, p2=0;
            Sleep(500);
            printf(".");
            if (mediaEvent->GetEvent(&ev, &p1, &p2, 0)==S_OK)
            {
                if (ev == EC_COMPLETE || ev == EC_USERABORT)
                {
                    printf("Done!\n");
                    stop = TRUE;
                    iSuccess = 1;
                }
                else
                if (ev == EC_ERRORABORT)
                {
                    printf("An error occured: HRESULT=%x\n", p1);
                    mediaControl->Stop();
                    stop = TRUE;
                }
                mediaEvent->FreeEventParams(ev, p1, p2);
            }
        }
    }
    CoUninitialize();
    return iSuccess ? 0 : 1;
}