/**************************************************************************
 * 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 <streams.h>
#include <Dmodshow.h>
#include "DShowTool.h"

#define SAFE_RELEASE(p) { if((p) != NULL) { (p)->Release(); (p)=NULL; } }
#define SAFE_DELETEMEDIATYPE(p) { if(p!=NULL) { DeleteMediaType(p); (p)=NULL; } }

HRESULT CDShowTool::Connect(IGraphBuilder *pGraphBuilder, IBaseFilter *pFilter1, IBaseFilter *pFilter2)
{
   CheckPointer(pGraphBuilder, E_POINTER);
   CheckPointer(pFilter1, E_POINTER);
   CheckPointer(pFilter2, E_POINTER);

   HRESULT hr;
   IPin *pInPin = NULL;
   IPin *pOutPin = NULL;
   int nInputPinCount, nOutputPinCount;

   nOutputPinCount = this->GetPinCount(pFilter1, PINDIR_OUTPUT);
   if(nOutputPinCount == 0) return E_FAIL;

   nInputPinCount  = this->GetPinCount(pFilter2, PINDIR_INPUT);
   if(nInputPinCount == 0) return E_FAIL;

   for(int i=0;i<nOutputPinCount;i++)
   {
      for(int j=0;j<nInputPinCount;j++)
      {
         pOutPin = this->GetPin(pFilter1, PINDIR_OUTPUT, i);
         if(pOutPin == NULL) return E_FAIL;

         pInPin = this->GetPin(pFilter2, PINDIR_INPUT, j);
         if(pInPin == NULL)
         {
            SAFE_RELEASE(pOutPin);
            return E_FAIL;;
         }

         hr = pGraphBuilder->Connect(pOutPin, pInPin);

         SAFE_RELEASE(pInPin);
         SAFE_RELEASE(pOutPin);

         if (hr == S_OK)
           return S_OK;

         if (hr == VFW_S_PARTIAL_RENDER)
           return VFW_S_PARTIAL_RENDER;
      }
   }

   SAFE_RELEASE(pInPin);
   SAFE_RELEASE(pOutPin);

   return VFW_E_CANNOT_CONNECT;
}
HRESULT CDShowTool::ConnectDirect(IGraphBuilder *pGraphBuilder, IBaseFilter *pFilter1, IBaseFilter *pFilter2, AM_MEDIA_TYPE *pAmMediaType)
{
   CheckPointer(pGraphBuilder, E_POINTER);
   CheckPointer(pFilter1, E_POINTER);
   CheckPointer(pFilter2, E_POINTER);


   HRESULT hr;
   int nInputPinCount, nOutputPinCount;
   IPin *pInPin = NULL;
   IPin *pOutPin = NULL;

   nOutputPinCount = this->GetPinCount(pFilter1, PINDIR_OUTPUT);
   if(nOutputPinCount == 0) return E_FAIL;

   nInputPinCount  = this->GetPinCount(pFilter2, PINDIR_INPUT);
   if(nInputPinCount == 0) return E_FAIL;

   for(int i=0;i<nOutputPinCount;i++)
   {
      for(int j=0;j<nInputPinCount;j++)
      {
         pOutPin = this->GetPin(pFilter1, PINDIR_OUTPUT, i);
         if(pOutPin == NULL) return E_FAIL;

         pInPin = this->GetPin(pFilter2, PINDIR_INPUT, j);
         if(pInPin == NULL)
         {
            SAFE_RELEASE(pOutPin);
            return E_FAIL;
         }

         hr = pGraphBuilder->ConnectDirect(pOutPin, pInPin, pAmMediaType);

         SAFE_RELEASE(pInPin);
         SAFE_RELEASE(pOutPin);

         if (hr == S_OK)
           return S_OK;
      }
   }

   SAFE_RELEASE(pInPin);
   SAFE_RELEASE(pOutPin);

   return VFW_E_CANNOT_CONNECT;
}
HRESULT CDShowTool::RenderDirectByFilter(IGraphBuilder *pGraphBuilder, int num, ...)
{
   CheckPointer(pGraphBuilder, E_POINTER);

   HRESULT hr;
   IBaseFilter *pFilter[2];

   va_list list;
   va_start(list, num);

   pFilter[0] = va_arg(list, IBaseFilter *);


   for(int i=1;i<num;i++)
   {
      pFilter[i % 2] = va_arg(list, IBaseFilter *);

      hr = ConnectDirect(pGraphBuilder, pFilter[(i-1) % 2], pFilter[i % 2], NULL);
      if(FAILED(hr))
         goto end;
   }

   hr = S_OK;

end:

   va_end(list);

   return hr;
}
AM_MEDIA_TYPE *CDShowTool::GetMediaType(IPin *pPin, int nFormatNumber)
{
   CheckPointer(pPin, 0);

   int nCount;
   HRESULT hr;
   AM_MEDIA_TYPE *pAmMediaType = NULL;
   IEnumMediaTypes *pEnumMediaTypes = NULL;

   nCount = 0;

   hr = pPin->EnumMediaTypes(&pEnumMediaTypes);
   if(pEnumMediaTypes == NULL)
      goto error;

   pEnumMediaTypes->Reset();

   while(true)
   {
      hr = pEnumMediaTypes->Next(1, &pAmMediaType, NULL);
      if(hr != S_OK)
         goto error;

      if(nCount == nFormatNumber)
         return pAmMediaType;

      nCount++;
   }

error:

   SAFE_DELETEMEDIATYPE(pAmMediaType);
   SAFE_RELEASE(pEnumMediaTypes);

   return NULL;
}
int CDShowTool::GetPinCount(IBaseFilter *pFilter, PIN_DIRECTION pinDirection)
{
   CheckPointer(pFilter, 0);

   HRESULT hr;
   int nPinCount;
   IEnumPins *pEnumPins = NULL;
   IPin *pPin = NULL;

   nPinCount = 0;

   hr = pFilter->EnumPins(&pEnumPins);
   if (hr != NO_ERROR) return 0;

   hr = pEnumPins->Reset();

   while(pEnumPins->Next(1, &pPin, 0) == S_OK)
   {
      PIN_DIRECTION dir;
      pPin->QueryDirection(&dir);

      if (dir == pinDirection)
         nPinCount++;

      SAFE_RELEASE(pPin);
   }

   SAFE_RELEASE(pPin);
   SAFE_RELEASE(pEnumPins);

   return nPinCount;
}
IPin* CDShowTool::GetPin(IBaseFilter *pFilter, PIN_DIRECTION pinDirection, int nPinNumber)
{
   CheckPointer(pFilter, NULL);

   int hr, count;
   IEnumPins *pEnumPins = NULL;
   IPin *pPin = NULL;

   count = 0;

   hr = pFilter->EnumPins(&pEnumPins);
   if (hr != NO_ERROR) return NULL;

   pEnumPins->Reset();

   while(pEnumPins->Next(1, &pPin, 0) == S_OK)
   {
      PIN_DIRECTION dir;
      pPin->QueryDirection(&dir);

      if (dir == pinDirection)
      {
         if(count == nPinNumber)
         {
            SAFE_RELEASE(pEnumPins);
            return pPin;
         }

         count++;
      }

      SAFE_RELEASE(pPin);
   }

   SAFE_RELEASE(pPin);
   SAFE_RELEASE(pEnumPins);

   return NULL;
}
IPin* CDShowTool::GetPin(IBaseFilter *pFilter, PIN_DIRECTION pinDirection, GUID majortype)
{
   CheckPointer(pFilter, NULL);

   int hr;
   ULONG cFetched;

   IPin *pPin = NULL;
   AM_MEDIA_TYPE *pAm = NULL;
   IEnumPins *pEnumPins = NULL;
   IEnumMediaTypes *pEnumMediaTypes = NULL;

   hr = pFilter->EnumPins(&pEnumPins);
   if (hr != NO_ERROR) return NULL;

   pEnumPins->Reset();

   while(pEnumPins->Next(1, &pPin, 0) == S_OK)
   {
      PIN_DIRECTION dir;
      pPin->QueryDirection(&dir);

      if (dir == pinDirection)
      {
         pPin->EnumMediaTypes(&pEnumMediaTypes);
         if(pEnumMediaTypes == NULL)
            goto error;

         pEnumMediaTypes->Reset();

         while(pEnumMediaTypes->Next(1, &pAm, &cFetched) == S_OK)
         {
            if(pAm->majortype == majortype)
            {
               SAFE_DELETEMEDIATYPE(pAm);
               SAFE_RELEASE(pEnumMediaTypes);
               SAFE_RELEASE(pEnumPins);

               return pPin;
            }

            SAFE_DELETEMEDIATYPE(pAm);
         }

         SAFE_RELEASE(pEnumMediaTypes);
      }
   }

error:

   SAFE_DELETEMEDIATYPE(pAm);
   SAFE_RELEASE(pEnumMediaTypes);
   SAFE_RELEASE(pPin);
   SAFE_RELEASE(pEnumPins);

   return NULL;
}
HRESULT CDShowTool::Disconnect(IGraphBuilder *pGraphBuilder, IBaseFilter *pFilter1, IBaseFilter *pFilter2)
{
   CheckPointer(pGraphBuilder, E_POINTER);
   CheckPointer(pFilter1, E_POINTER);
   CheckPointer(pFilter2, E_POINTER);

   HRESULT hr;
   IPin *pOutPin = NULL;
   IPin *pInPin = NULL;

   pOutPin = GetPin(pFilter1, PINDIR_OUTPUT, 0);
   pInPin = GetPin(pFilter2, PINDIR_INPUT, 0);

   hr = pGraphBuilder->Disconnect(pOutPin);
   if(hr == S_OK)
   {
      hr = pGraphBuilder->Disconnect(pInPin);
   }

   SAFE_RELEASE(pOutPin);
   SAFE_RELEASE(pInPin);

   return hr;
}
IBaseFilter* CDShowTool::AddDmoFilter(IGraphBuilder *pGraphBuilder, CLSID clsid_Category, CLSID clsid)
{
   CheckPointer(pGraphBuilder, NULL);

   HRESULT hr;
   IBaseFilter *pDmoFilter = NULL;
   IDMOWrapperFilter *pDmoWrapper = NULL;

   hr = CoCreateInstance(CLSID_DMOWrapperFilter, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, reinterpret_cast<void**>(&pDmoFilter));
   if(FAILED(hr))
      return NULL;

   hr = pDmoFilter->QueryInterface(IID_IDMOWrapperFilter, reinterpret_cast<void**>(&pDmoWrapper));
   if(FAILED(hr))
   {
      SAFE_RELEASE(pDmoFilter);
      return NULL;
   }

   hr = pDmoWrapper->Init(clsid, clsid_Category); 
   SAFE_RELEASE(pDmoWrapper);

   if(FAILED(hr))
   {
      SAFE_RELEASE(pDmoFilter);
      return NULL;
   }

   hr = pGraphBuilder->AddFilter(pDmoFilter, NULL);
   if(FAILED(hr))
   {
      SAFE_RELEASE(pDmoFilter);
      return NULL;
   }

   return pDmoFilter;
}
IBaseFilter* CDShowTool::AddFilter(IGraphBuilder *pGraphBuilder, CLSID clsid)
{
   CheckPointer(pGraphBuilder, NULL);

   HRESULT hr;
   IBaseFilter *pBaseFilter = NULL;

   hr = ::CoCreateInstance(clsid, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, (void**)&pBaseFilter);
   if(pBaseFilter == NULL) return NULL;

   hr = pGraphBuilder->AddFilter(pBaseFilter, NULL);
   if(hr != NO_ERROR)
   {
      SAFE_RELEASE(pBaseFilter);
      return NULL;
   }

   return pBaseFilter;
}
IBaseFilter* CDShowTool::AddSourceFilter(IGraphBuilder *pGraphBuilder, CLSID clsid, WCHAR *wpFileName)
{
   CheckPointer(pGraphBuilder, NULL);

   HRESULT hr;
   IBaseFilter *pBaseFilter = NULL;
   IFileSourceFilter *pFileSourceFilter = NULL;

   hr = ::CoCreateInstance(clsid, NULL, CLSCTX_INPROC_SERVER, IID_IBaseFilter, (void**)&pBaseFilter);
   if(pBaseFilter == NULL) return NULL;

   pBaseFilter->QueryInterface(IID_IFileSourceFilter, (void**)&pFileSourceFilter);
   if(pFileSourceFilter == NULL)
   {
      SAFE_RELEASE(pBaseFilter);
      return NULL;
   }

   hr = pGraphBuilder->AddFilter(pBaseFilter, NULL);
   if(hr != NO_ERROR)
   {
      SAFE_RELEASE(pBaseFilter);
      return NULL;
   }

   hr = pFileSourceFilter->Load(wpFileName, NULL);
   SAFE_RELEASE(pFileSourceFilter);
   if(FAILED(hr))
   {
      pGraphBuilder->RemoveFilter(pBaseFilter);
      SAFE_RELEASE(pBaseFilter);
      return NULL;
   }

   return pBaseFilter;
}
IBaseFilter* CDShowTool::AddWriteFilter(IGraphBuilder *pGraphBuilder, CLSID clsid, WCHAR *wpFileName)
{
   CheckPointer(pGraphBuilder, NULL);

   HRESULT hr;
   IBaseFilter *pBaseFilter = NULL;
   IFileSinkFilter *pFileSinkFilter = NULL;

   pBaseFilter = this->AddFilter(pGraphBuilder, clsid);
   if(pBaseFilter == NULL)
      return NULL;

   pBaseFilter->QueryInterface(IID_IFileSinkFilter, (void**)&pFileSinkFilter);
   if(pFileSinkFilter == NULL)
   {
      pGraphBuilder->RemoveFilter(pBaseFilter);
      SAFE_RELEASE(pBaseFilter);
      return NULL;
   }

   hr = pFileSinkFilter->SetFileName(wpFileName, NULL);
   if(FAILED(hr))
   {
      SAFE_RELEASE(pFileSinkFilter);
      pGraphBuilder->RemoveFilter(pBaseFilter);
      SAFE_RELEASE(pBaseFilter);
      return NULL;
   }

   SAFE_RELEASE(pFileSinkFilter);

   return pBaseFilter;
}
