#include <atlbase.h>
#include <dshow.h>
#include <stdio.h>
#include <windows.h>


/** straight from example code; add to ROT for use in graph edit. */
static HRESULT AddToRot(IUnknown *pUnkGraph, DWORD *pdwRegister) 
{
    IMoniker * pMoniker;
    IRunningObjectTable *pROT;
    if (FAILED(GetRunningObjectTable(0, &pROT))) {
        return E_FAIL;
    }
    WCHAR wsz[256];
    wsprintfW(wsz, L"FilterGraph %08x pid %08x", (DWORD*)pUnkGraph, GetCurrentProcessId());
    HRESULT hr = CreateItemMoniker(L"!", wsz, &pMoniker);
    if (SUCCEEDED(hr)) {
        hr = pROT->Register(ROTFLAGS_REGISTRATIONKEEPSALIVE, pUnkGraph,
            pMoniker, pdwRegister);
        pMoniker->Release();
    }
    pROT->Release();
    return hr;
}


/** called on exit. */
static void _cleanup (void) {

    CoUninitialize();

}


/** returns first free unconnected pin in given direction in base, or NULL. */
static IPin * _getnextpin (IBaseFilter *base, PIN_DIRECTION dir) {

    HRESULT hr;
    PIN_DIRECTION pindir;
    CComPtr<IEnumPins> ep;
    IPin *pin, *buddy, *found = NULL;

    hr = base->EnumPins(&ep);
    if (FAILED(hr))
        return NULL;

    while (ep->Next(1, &pin, NULL) == S_OK) {
        pin->QueryDirection(&pindir);
        if (pindir == dir) {
            hr = pin->ConnectedTo(&buddy);
            if (FAILED(hr)) {
                found = pin;
                break;
            }
            buddy->Release();
        }
        pin->Release();
    }

    return found;

}


/** display message and exit */
#define EXITMSG(s) { \
    MessageBox(NULL, (s), "Error", MB_ICONERROR); \
    exit(1); \
}


/** display error message and exit if FAILED(hr) */
static void _checkhr (HRESULT hr, const char *msg) {

    char buf[512], buf2[1024];

    if (FAILED(hr)) {
        AMGetErrorText(hr, buf, sizeof(buf) - 1);
        _snprintf(buf2, sizeof(buf2), "%s:\n%s", msg, buf);
        EXITMSG(buf2);
    }  

}


/** using the graph builder "gb", connect the first unconnected output pin of
 *  "from" to the first unconnected input pin of "to". "from" and "to" can't
 *  be NULL and must already have been added to the graph builder. displays
 *  message and exits on error. */
static void _connectpins (IGraphBuilder *gb,
                          IBaseFilter *from,
                          IBaseFilter *to,
                          const char *what)
{

    HRESULT hr;
    CComPtr<IPin> fpin, tpin;
    char errstr[512];

    _snprintf(errstr, sizeof(errstr), "Error connecting %s", what);

    if (!(fpin = _getnextpin(from, PINDIR_OUTPUT)))
        EXITMSG("Error connecting pins (no free source output pin)");

    if (!(tpin = _getnextpin(to, PINDIR_INPUT)))
        EXITMSG("Error connecting pins (no free destination input pin)");

    hr = gb->Connect(fpin, tpin);
    _checkhr(hr, errstr);

}



int main (int argc, char **argv) {

    USES_CONVERSION;

    LPCWSTR infn, outfn;
    HRESULT hr;
    CComPtr<IGraphBuilder> gb;
    CComPtr<IMediaControl> mc;
    CComPtr<IMediaEvent> me;
    CComPtr<IBaseFilter> input_base, tee, file_base, mux_base, vidrend_base;
    CComPtr<IFileSinkFilter2> file_output;
    long evcode;

    // -------- COMMAND LINE

    if (argc != 3)
        EXITMSG("No input/output filename specified");

    infn = A2W(argv[1]);
    outfn = A2W(argv[2]);

    //MessageBoxW(NULL, infn, L"Input Filename", MB_ICONINFORMATION);
    //MessageBoxW(NULL, outfn, L"Output Filename", MB_ICONINFORMATION);

    // -------- COM / FILTER GRAPH INITIALIZE

    hr = CoInitialize(NULL);
    _checkhr(hr, "CoInitialize failed");

    atexit(&_cleanup);
    
    hr = CoCreateInstance(CLSID_FilterGraph, NULL, CLSCTX_INPROC_SERVER, IID_IGraphBuilder, (void **)&gb);
    _checkhr(hr, "Could not create filter graph");

    DWORD reg;
    AddToRot(gb, &reg);

    hr = gb->QueryInterface(IID_IMediaControl, (void **)&mc);
    _checkhr(hr, "Could not get IMediaControl interface");

    hr = gb->QueryInterface(IID_IMediaEvent, (void **)&me);
    _checkhr(hr, "Could not get IMediaEvent interface");

    // -------- FILE SOURCE FILTER

    hr = gb->AddSourceFilter(infn, L"Source", (IBaseFilter **)&input_base);
    _checkhr(hr, "Could not read input file");

    // -------- FILE RUNS TO INFINITE PIN TEE FILTER

    hr = CoCreateInstance(CLSID_InfTee, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, (void **)&tee);
    _checkhr(hr, "CoCreateInstance(CLSID_InfTee) failed");
    
    hr = gb->AddFilter(tee, L"Tee");
    _checkhr(hr, "Failed to add tee filter to graph");

    // -------- CREATE VIDEO RENDERER

    hr = CoCreateInstance(CLSID_VideoRenderer, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, (void **)&vidrend_base);
    _checkhr(hr, "CoCreateInstance(CLSID_VideoRenderer) failed");

    hr = gb->AddFilter(vidrend_base, L"Renderer");
    _checkhr(hr, "Failed to add video renderer to graph.");

    // -------- CREATE AVI STREAM MULTIPLEXER + FILE WRITER

    hr = CoCreateInstance(CLSID_AviDest, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, (void **)&mux_base);
    _checkhr(hr, "Couldn't create AVI multiplexer");

    hr = gb->AddFilter(mux_base, L"AVI Mux");
    _checkhr(hr, "Couldn't add multiplexer to graph");

    hr = CoCreateInstance(CLSID_FileWriter, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, (void **)&file_base);
    _checkhr(hr, "Couldn't create file writer filter");

    hr = file_base->QueryInterface(IID_IFileSinkFilter2, (void **)&file_output);
    _checkhr(hr, "Couldn't get file writer sink interface");

    hr = gb->AddFilter(file_base, L"File Writer");
    _checkhr(hr, "Couldn't add file writer to graph");

    hr = file_output->SetFileName(W2COLE(outfn), NULL);
    _checkhr(hr, "Couldn't set output file name");

    hr = file_output->SetMode(AM_FILE_OVERWRITE);
    _checkhr(hr, "Couldn't set output file overwrite mode");

    // -------- CONNECT (TRY MOVING SOME OF THESE BEFORE MUXER IS ADDED ABOVE
    
    _connectpins(gb, input_base, tee, "input to tee");
    _connectpins(gb, tee, vidrend_base, "tee to renderer");
    _connectpins(gb, tee, mux_base, "tee to multiplexer");
    _connectpins(gb, mux_base, file_base, "multiplexer to file writer");

    // -------- GO

    printf("press enter\n");
    getchar(); // pause so you can check it out in graph edit
    printf("running...\n");

    hr = mc->Run();
    _checkhr(hr, "Couldn't run filter graph");

    hr = me->WaitForCompletion(INFINITE, &evcode);
    _checkhr(hr, "Couldn't wait for filter graph to finish");

    if (evcode != EC_COMPLETE)
        EXITMSG("Unexpected end of playback");

    // --------
    
    return 0;

}