// Applet.cpp --- COM stuff
//
// Copyright (C) 2004, 2007 Raymond Penners <raymond@dotsphinx.com>
// All rights reserved.
//
// 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 3 of the License, 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 this program.  If not, see <http://www.gnu.org/licenses/>.
//
// $Id$

#include <ole2.h>
#include <comcat.h>
#include <olectl.h>
#include "ClsFact.h"
#include <shlwapi.h>

#pragma data_seg(".text")
#define INITGUID
#include <initguid.h>
#include <shlguid.h>
#include "Guid.h"

#include "Debug.h"

#pragma data_seg()

extern "C" BOOL WINAPI DllMain(HINSTANCE, DWORD, LPVOID);
BOOL RegisterServer(CLSID, LPTSTR, BOOL reg);
BOOL RegisterComCat(CLSID, CATID, BOOL reg);

HINSTANCE g_hInst;
UINT g_DllRefCount;

extern "C" BOOL WINAPI
DllMain(HINSTANCE hInstance, DWORD dwReason, LPVOID lpReserved)
{
  switch (dwReason)
    {
    case DLL_PROCESS_ATTACH:
      g_hInst = hInstance;
      break;

    case DLL_PROCESS_DETACH:
      break;
    }

  return TRUE;
}

STDAPI
DllCanUnloadNow()
{
  TRACE_ENTER("DllCanUnloadNow");
  TRACE_RETURN(g_DllRefCount);
  return (g_DllRefCount ? S_FALSE : S_OK);
}

STDAPI
DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID *ppReturn)
{
  *ppReturn = NULL;

  // if we don't support this classid, return the proper error code
  if (!IsEqualCLSID(rclsid, CLSID_WorkraveDeskBand))
    return CLASS_E_CLASSNOTAVAILABLE;

  // create a CClassFactory object and check it for validity
  CClassFactory *pClassFactory = new CClassFactory(rclsid);
  if (NULL == pClassFactory)
    return E_OUTOFMEMORY;

  // get the QueryInterface return for our return value
  HRESULT hResult = pClassFactory->QueryInterface(riid, ppReturn);

  // call Release to decement the ref count - creating the object set it to one
  // and QueryInterface incremented it - since its being used externally (not by
  // us), we only want the ref count to be 1
  pClassFactory->Release();

  // return the result from QueryInterface
  return hResult;
}

static void
ClearDeskBandCache()
{
  /*
  Remove the cache of the deskbands on Windows 2000. This will cause the new
  deskband to be displayed in the toolbar menu the next time the user brings it
  up. See KB article Q214842 for more information on this.
  */
  TCHAR szSubKey[MAX_PATH];
  TCHAR szCATID[MAX_PATH];
  LPWSTR pwszCATID;

  StringFromCLSID(CATID_DeskBand, &pwszCATID);
  if (pwszCATID)
    {
#ifdef UNICODE
      lstrcpy(szCATID, pwszCATID);
#else
      WideCharToMultiByte(CP_ACP, 0, pwszCATID, -1, szCATID, ARRAYSIZE(szCATID), NULL, NULL);
#endif

      // free the string
      CoTaskMemFree(pwszCATID);

      wsprintf(szSubKey, TEXT("Component Categories\\%s\\Enum"), szCATID);

      RegDeleteKey(HKEY_CLASSES_ROOT, szSubKey);
    }
}

STDAPI
DllUnregisterServer()
{
  RegisterComCat(CLSID_WorkraveDeskBand, CATID_DeskBand, FALSE);
  RegisterServer(CLSID_WorkraveDeskBand, TEXT("Workrave"), FALSE);
  ClearDeskBandCache();
  return S_OK;
}

STDAPI
DllRegisterServer()
{
  // Register the desk band object.
  if (!RegisterServer(CLSID_WorkraveDeskBand, TEXT("Workrave"), TRUE))
    return SELFREG_E_CLASS;

  // Register the component categories for the desk band object.
  if (!RegisterComCat(CLSID_WorkraveDeskBand, CATID_DeskBand, TRUE))
    return SELFREG_E_CLASS;

  ClearDeskBandCache();

  return S_OK;
}

typedef struct
{
  HKEY hRootKey;
  LPTSTR szSubKey; // TCHAR szSubKey[MAX_PATH];
  LPTSTR lpszValueName;
  LPTSTR szData; // TCHAR szData[MAX_PATH];
} DOREGSTRUCT, *LPDOREGSTRUCT;

BOOL
RegisterServer(CLSID clsid, LPTSTR lpszTitle, BOOL reg)
{
  int i;
  HKEY hKey;
  LRESULT lResult;
  DWORD dwDisp;
  TCHAR szSubKey[MAX_PATH];
  TCHAR szCLSID[MAX_PATH];
  TCHAR szModule[MAX_PATH];
  LPWSTR pwsz;

  // get the CLSID in string form
  StringFromIID(clsid, &pwsz);

  if (pwsz)
    {
#ifdef UNICODE
      lstrcpy(szCLSID, pwsz);
#else
      WideCharToMultiByte(CP_ACP, 0, pwsz, -1, szCLSID, ARRAYSIZE(szCLSID), NULL, NULL);
#endif

      // free the string
      CoTaskMemFree(pwsz);
    }

  // get this app's path and file name
  GetModuleFileName(g_hInst, szModule, ARRAYSIZE(szModule));

  DOREGSTRUCT ClsidEntries[] = {HKEY_CLASSES_ROOT,
                                TEXT("CLSID\\%s"),
                                NULL,
                                lpszTitle,
                                HKEY_CLASSES_ROOT,
                                TEXT("CLSID\\%s\\InprocServer32"),
                                NULL,
                                szModule,
                                HKEY_CLASSES_ROOT,
                                TEXT("CLSID\\%s\\InprocServer32"),
                                TEXT("ThreadingModel"),
                                TEXT("Apartment")};

  if (reg)
    {
      // register the CLSID entries
      for (i = 0; i < sizeof(ClsidEntries) / sizeof(ClsidEntries[0]); i++)
        {
          // create the sub key string - for this case, insert the file extension
          wsprintf(szSubKey, ClsidEntries[i].szSubKey, szCLSID);

          lResult =
            RegCreateKeyEx(ClsidEntries[i].hRootKey, szSubKey, 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hKey, &dwDisp);

          if (NOERROR == lResult)
            {
              TCHAR szData[MAX_PATH];

              // if necessary, create the value string
              wsprintf(szData, ClsidEntries[i].szData, szModule);

              lResult = RegSetValueEx(hKey, ClsidEntries[i].lpszValueName, 0, REG_SZ, (LPBYTE)szData, lstrlen(szData) + 1);

              RegCloseKey(hKey);
            }
          else
            return FALSE;
        }
    }
  else
    {
      // create the sub key string - for this case, insert the file extension
      wsprintf(szSubKey, ClsidEntries[0].szSubKey, szCLSID);

      SHDeleteKey(ClsidEntries[0].hRootKey, szSubKey);
    }

  // If running on NT, register the extension as approved.
  OSVERSIONINFO osvi;

  osvi.dwOSVersionInfoSize = sizeof(osvi);
  GetVersionEx(&osvi);

  if (VER_PLATFORM_WIN32_NT == osvi.dwPlatformId)
    {
      lstrcpy(szSubKey, TEXT("Software\\Microsoft\\Windows\\CurrentVersion\\Shell Extensions\\Approved"));

      lResult = RegCreateKeyEx(HKEY_LOCAL_MACHINE, szSubKey, 0, NULL, REG_OPTION_NON_VOLATILE, KEY_WRITE, NULL, &hKey, &dwDisp);

      if (NOERROR == lResult)
        {
          TCHAR szData[MAX_PATH];

          // Create the value string.
          lstrcpy(szData, lpszTitle);

          if (reg)
            {
              lResult = RegSetValueEx(hKey, szCLSID, 0, REG_SZ, (LPBYTE)szData, (lstrlen(szData) + 1) * sizeof(TCHAR));
            }
          else
            {
              lResult = RegDeleteValue(hKey, szCLSID);
            }

          RegCloseKey(hKey);
        }
      else
        return FALSE;
    }

  return TRUE;
}

BOOL
RegisterComCat(CLSID clsid, CATID CatID, BOOL reg)
{
  ICatRegister *pcr;
  HRESULT hr = S_OK;

  CoInitialize(NULL);

  hr = CoCreateInstance(CLSID_StdComponentCategoriesMgr, NULL, CLSCTX_INPROC_SERVER, IID_ICatRegister, (LPVOID *)&pcr);

  if (SUCCEEDED(hr))
    {
      if (reg)
        {
          hr = pcr->RegisterClassImplCategories(clsid, 1, &CatID);
        }
      else
        {
          hr = pcr->UnRegisterClassImplCategories(clsid, 1, &CatID);
        }

      pcr->Release();
    }

  CoUninitialize();

  return SUCCEEDED(hr);
}
