#include <OleCtl.h>
#include <stdio.h>

ITypeLib *g_pTypeLib = NULL;
HMODULE g_hModule = NULL;

#ifdef _WIN64
  #pragma comment(linker, "/export:DllGetClassObject,PRIVATE")
  #pragma comment(linker, "/export:DllRegisterServer,PRIVATE")
  #pragma comment(linker, "/export:DllUnregisterServer,PRIVATE")
#else
  #pragma comment(linker, "/export:DllGetClassObject=_DllGetClassObject@12,PRIVATE")
  #pragma comment(linker, "/export:DllRegisterServer=_DllRegisterServer@0,PRIVATE")
  #pragma comment(linker, "/export:DllUnregisterServer=_DllUnregisterServer@0,PRIVATE")
#endif

const GUID IID_SimpleDispatch = {0x2bb79939, 0xee89, 0x4ae0, 0xbf, 0x7d, 0xe7, 0xfb, 0x17, 0x5a, 0x87, 0xcf};

// Base class for COM objects created by COMDemo

struct SimpleDispatch : IDispatch
{
  unsigned long m_refs;

  SimpleDispatch() : m_refs(1) {}
  virtual ~SimpleDispatch() {}

// IUnknown methods ...
  HRESULT __stdcall QueryInterface(REFIID riid, void **ppv)
  {
    if (riid == IID_IUnknown || riid == IID_IDispatch || riid == IID_SimpleDispatch || riid == IID_This())
    {
      *ppv = this;
      AddRef();
      return NOERROR;
    }
    *ppv = NULL;
    return ResultFromScode(E_NOINTERFACE);
  }
  unsigned long __stdcall AddRef() {return m_refs++;}
  unsigned long __stdcall Release()
  {
    if (--m_refs > 0) return m_refs;
    delete this;
    return 0;
  }

// IDispatch methods. These handle late-bound calls ...
  HRESULT __stdcall GetTypeInfoCount(unsigned int *pctinfo)
  {
    *pctinfo = 1;
    return NOERROR;
  }

  HRESULT __stdcall GetTypeInfo(unsigned int itinfo, LCID lcid, ITypeInfo **pptinfo)
  {
    return g_pTypeLib->GetTypeInfoOfGuid(IID_This(), pptinfo);
  }

  HRESULT __stdcall GetIDsOfNames(REFIID riid, LPOLESTR *rgszNames,
      unsigned int cNames, LCID lcid, DISPID *rgdispid)
  {
    ITypeInfo *TI;
    HRESULT hRes;
    
    if (hRes = GetTypeInfo(0, 0, &TI)) return hRes;
    hRes = DispGetIDsOfNames(TI, rgszNames, cNames, rgdispid);
    TI->Release();
    return hRes;
  }

  HRESULT __stdcall Invoke(DISPID dispidMember, REFIID riid,
    LCID lcid, unsigned short wFlags, DISPPARAMS *pdispparams,
    VARIANT *pvarResult, EXCEPINFO *pexcepinfo, unsigned int *puArgErr)
  {
    ITypeInfo *TI;
    HRESULT hRes;
    
    if (hRes = GetTypeInfo(0, 0, &TI)) return hRes;
    hRes = DispInvoke(this, TI, dispidMember, wFlags, pdispparams,
        pvarResult, pexcepinfo, puArgErr);
    TI->Release();
    return hRes;
  }

// Extra virtual function is to obtain the IID of the object ...
  virtual const GUID &IID_This() const {return IID_SimpleDispatch;}
};

struct ClassFactoryWithClsid : IClassFactory
{
  GUID m_clsid;
  ClassFactoryWithClsid *m_nextCF;

  ClassFactoryWithClsid(const GUID &clsid);

// These are only ever static objects, so the memory management functions can be implemented trivially ...
  unsigned long __stdcall AddRef() {return 1;}
  unsigned long __stdcall Release() {return 1;}

  HRESULT __stdcall LockServer(BOOL) {return NOERROR;}

} *g_firstCF = NULL;

ClassFactoryWithClsid::ClassFactoryWithClsid(const GUID &clsid) : m_clsid(clsid)
{
// Create a linked list of ClassFactory objects ...
  m_nextCF = g_firstCF;
  g_firstCF = this;
}

template<class C> struct SimpleClassFactory : ClassFactoryWithClsid
{
  SimpleClassFactory(const GUID &clsid) : ClassFactoryWithClsid(clsid) {}

  HRESULT __stdcall QueryInterface(REFIID riid, void **ppv)
  {
    if (riid == IID_IClassFactory)
    {
      *ppv = this;
      return NOERROR;
    }
    if (riid == IID_IUnknown || riid == IID_IDispatch || riid == IID_SimpleDispatch)
    {
      *ppv = new C;
      return NOERROR;
    }
    return E_NOINTERFACE;
  }

  HRESULT __stdcall CreateInstance(IUnknown *pUnkOuter, REFIID riid, void **ppvObject)
  {
    C *obj = new C;
    HRESULT hRes = obj->QueryInterface(riid, ppvObject);
    obj->Release();
    return hRes;
  }
};

const GUID IID_Library = {0xc7e9002b, 0x9e7f, 0x43b5, 0x97, 0x1d, 0xe2, 0x53, 0x9e, 0x60, 0x39, 0xc2};

// DllMain is the first entry point, and is used just to load the Type Library, which must be registered ...

BOOL __stdcall DllMain(HANDLE hModule, DWORD reason, void *)
{
  switch(reason)
  {
  case DLL_PROCESS_ATTACH:
    g_hModule = (HMODULE)hModule;
    if (!wcsstr(GetCommandLine(), L"regsvr32"))
    {
      if (LoadRegTypeLib(IID_Library, 1, 0, NULL, &g_pTypeLib))
        MessageBox(NULL, L"COMDemo is not registered.", L"COMDemo", MB_ICONEXCLAMATION);
    }
    return TRUE;
  
  case DLL_PROCESS_DETACH:;
    if (g_pTypeLib) g_pTypeLib->Release();
    return TRUE;
  }
  return TRUE;
}

/* DllGetClassObject is called when an object is created via "CreateObject" or "New". We attempt to 
   match the CLSID to the appropriate ClassFactory, passing back the ClassFactory pointer ... */

HRESULT __stdcall DllGetClassObject(REFCLSID clsid, REFIID riid, LPVOID *ppv)
{
  for (ClassFactoryWithClsid *CF = g_firstCF; CF; CF = CF->m_nextCF)
    if (CF->m_clsid == clsid) return CF->QueryInterface(riid, ppv);
  return CLASS_E_CLASSNOTAVAILABLE;
}

char g_errorText[1024];

void ThrowError(const char *formatString, ...)
{
  va_list args;
  va_start(args, formatString);
  _vsnprintf_s(g_errorText, sizeof(g_errorText), sizeof(g_errorText), formatString, args);
  throw (char *)g_errorText;
}

// HRESULT is the generic COM error code. Here we attempt to convert it into a string ...

void ThrowError(HRESULT hRes)
{
  if (!hRes) return;
  if (!FormatMessageA(FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_MAX_WIDTH_MASK, NULL, hRes,
    MAKELANGID(LANG_NEUTRAL, SUBLANG_SYS_DEFAULT), g_errorText, sizeof(g_errorText), NULL))
      sprintf_s(g_errorText, sizeof(g_errorText), "Error with HRESULT 0x%lx - FormatMessage failed.", hRes);
  throw (char *)g_errorText;
}

void GUIDAsText(char *buff, const GUID &iid)
{
  sprintf_s(buff, 39, "{%08X-%04X-%04X-%02X%02X-%02X%02X%02X%02X%02X%02X}",
    iid.Data1, iid.Data2, iid.Data3, iid.Data4[0], iid.Data4[1], iid.Data4[2], 
    iid.Data4[3], iid.Data4[4], iid.Data4[5], iid.Data4[6], iid.Data4[7]);
}

inline size_t wcslen1(const wchar_t *wstr) {return wstr? wcslen(wstr): 0;}

// A single-byte string as an automatic object. Can be constructed from a BSTR ...

struct StrWrap
{
  char *m_str;
  StrWrap() : m_str(NULL) {}
  StrWrap(BSTR bstr)
  {
    if (!bstr || !bstr[0]) m_str = NULL;
    else
    {
      int nWide = (int)wcslen(bstr);
      int n = WideCharToMultiByte(CP_ACP, WC_COMPOSITECHECK, bstr, nWide, NULL, 0, NULL, NULL);
      m_str = new char[n + 1];
      WideCharToMultiByte(CP_ACP, WC_COMPOSITECHECK, bstr, nWide, m_str, n, NULL, NULL);
      m_str[n] = '\0';
    }
  }
  ~StrWrap() {delete [] m_str;}
};

// A BSTR as an automatic object. This is a cut-down version of _bstr_t ...

struct BSTRWrap
{
  BSTR m_bstr;

  BSTRWrap() : m_bstr(NULL) {}
  BSTRWrap(const wchar_t *bstr)
  {
    m_bstr = SysAllocString(bstr);
  }

  ~BSTRWrap() {SysFreeString(m_bstr);}

  BSTRWrap &operator=(const wchar_t *bstr)
  {
    SysFreeString(m_bstr);
    m_bstr = SysAllocString(bstr);
    return *this;
  }

  operator wchar_t *() {return m_bstr;}
};

template <class C> class SmartPtr
{
public:
  C *m_obj;

  SmartPtr() : m_obj(NULL) {}
  ~SmartPtr() {if (m_obj) m_obj->Release();}
};

struct TLIBATTRWrap
{
  TLIBATTR *m_typeLibAttr;
  TLIBATTRWrap() : m_typeLibAttr(NULL) {}
  ~TLIBATTRWrap() {g_pTypeLib->ReleaseTLibAttr(m_typeLibAttr);}
};

struct TYPEATTRWrap
{
  ITypeInfo *m_typeInfo;
  TYPEATTR *m_typeAttr;

  TYPEATTRWrap(ITypeInfo *typeInfo) : m_typeInfo(typeInfo), m_typeAttr(NULL)
  {
    m_typeInfo->GetTypeAttr(&m_typeAttr);
  }

  ~TYPEATTRWrap()
  {
    m_typeInfo->ReleaseTypeAttr(m_typeAttr);
  }
};

// Class to assist handling the System Registry ...

struct RegDB
{
  HKEY m_hKey;

  RegDB() : m_hKey(NULL) {}
  RegDB(const char *vendor, const char *product, const char *section) : m_hKey(NULL)
  {
    try {Open(vendor, product, section);}
    catch (...) {Close(); throw;}
  }

  ~RegDB() {Close();}

  void Open(const char *vendor, const char *product, const char *section)
  {
    Close();
    char buff[1024];
    sprintf_s(buff, sizeof(buff), "Software\\%s\\%s", vendor, product);
    if (section)
    {
      strcat_s(buff, sizeof(buff), "\\");
      strcat_s(buff, sizeof(buff), section);
    }
    DWORD dType;
    ThrowError(RegCreateKeyExA(HKEY_LOCAL_MACHINE, buff, 0, "", 0, KEY_WRITE, NULL, &m_hKey, &dType));
  }

  void Close()
  {
    if (m_hKey) RegCloseKey(m_hKey);
    m_hKey = NULL;
  }

  void Write(const char *field, const char *buff) const    // write string
  {
    ThrowError(RegSetValueExA(m_hKey, field, 0, REG_SZ, (unsigned char *)(buff? buff: ""), (DWORD)(buff? strlen(buff) + 1: 1)));
  }
};

void WriteDefaultValue(const char *key, const char *subKey, const char *value)
{
  RegDB regDB("Classes", key, subKey);
  regDB.Write(NULL, value);
}

void WriteDefaultValue(const char *key, const char *subKey)
{
  RegDB regDB("Classes", key, subKey);
}

int SubKeyNames(char *subKeyNames, int nSubKeyNames, HKEY hKey)
{
  char *ptr = subKeyNames;
  ptr[0] = ptr[1] = '\0';
  int i;
  for (i = 0;; i++)
  {
    DWORD nSubKeyName = (DWORD)(subKeyNames + nSubKeyNames - 1 - ptr);
    FILETIME fileTime;
    LONG e = RegEnumKeyExA(hKey, i, ptr, &nSubKeyName, NULL, NULL, NULL, &fileTime);
    if (e == ERROR_NO_MORE_ITEMS) break;
    if (e != ERROR_MORE_DATA) ThrowError(e);
    ptr += nSubKeyName + 1;
    ptr[0] = '\0';
  }
  return i;
}

void RecursiveDeleteKey(HKEY hParentKey, const char *keyName)
{
  HKEY hKey;
  ThrowError(RegOpenKeyExA(hParentKey, keyName, 0, KEY_ALL_ACCESS, &hKey));
  char subKeyNames[1024];
  int n = SubKeyNames(subKeyNames, sizeof(subKeyNames), hKey);
  for (char *ptr = subKeyNames; n > 0; ptr += strlen(ptr) + 1, n--)
    RecursiveDeleteKey(hKey, ptr);
  ThrowError(RegCloseKey(hKey));
  ThrowError(RegDeleteKeyA(hParentKey, keyName));
}

void RegDBDeleteKey(const char *vendor, const char *product, const char *section)
{
  char subKeyName[256];
  HKEY hKey;
  if (section)
  {
    sprintf_s(subKeyName, sizeof(subKeyName), "Software\\%s\\%s", vendor, product);
    ThrowError(RegOpenKeyExA(HKEY_LOCAL_MACHINE, subKeyName, 0, KEY_ALL_ACCESS, &hKey));
    RecursiveDeleteKey(hKey, section);
  }
  else
  {
    sprintf_s(subKeyName, sizeof(subKeyName), "Software\\%s", vendor);
    ThrowError(RegOpenKeyExA(HKEY_LOCAL_MACHINE, subKeyName, 0, KEY_ALL_ACCESS, &hKey));
    RecursiveDeleteKey(hKey, product);
  }
  ThrowError(RegCloseKey(hKey));
}

// This is called by regsvr32, and fills in the registry entries for the COM objects defined here ...

HRESULT __stdcall DllRegisterServer()
{
  try
  {
    char libPath[256], buff[256];
    if (!GetModuleFileNameA(g_hModule, libPath, sizeof(libPath)))
      ThrowError("Unable to get file name for module %08X", g_hModule);
    if (!g_pTypeLib && LoadTypeLib(L"COMDemo.tlb", &g_pTypeLib))
      throw "Failed to find COMDemo.tlb";
    BSTRWrap libName, libDoc;
    char libid[39], verStr[20];
    ThrowError(g_pTypeLib->GetDocumentation(MEMBERID_NIL, &libName.m_bstr, &libDoc.m_bstr, NULL, NULL));
    TLIBATTRWrap tLA;
    ThrowError(g_pTypeLib->GetLibAttr(&tLA.m_typeLibAttr));
    GUIDAsText(libid, tLA.m_typeLibAttr->guid);
    sprintf_s(verStr, sizeof(verStr), "%x.%x", tLA.m_typeLibAttr->wMajorVerNum, tLA.m_typeLibAttr->wMinorVerNum);
    for (int i = 0;; i++)
    {
      SmartPtr<ITypeInfo> typeInfo;
      if (g_pTypeLib->GetTypeInfo(i, &typeInfo.m_obj)) break;
      BSTRWrap tiName, tiDoc;
      ThrowError(typeInfo.m_obj->GetDocumentation(MEMBERID_NIL, &tiName.m_bstr, &tiDoc.m_bstr, NULL, NULL));
      TYPEATTRWrap TAW(typeInfo.m_obj);
      if (!TAW.m_typeAttr) continue;
      if (TAW.m_typeAttr->typekind == TKIND_COCLASS)
      {
        char qualifiedName[128];
        sprintf_s(qualifiedName, sizeof(qualifiedName), "%S.%S", libName.m_bstr, tiName.m_bstr);
        StrWrap tiDocS = tiDoc.m_bstr;
        WriteDefaultValue(qualifiedName, NULL, tiDocS.m_str);
        char clsid[39];
        GUIDAsText(clsid, TAW.m_typeAttr->guid);
        WriteDefaultValue(qualifiedName, "CLSID", clsid);
        WriteDefaultValue("CLSID", clsid, tiDocS.m_str);
        sprintf_s(buff, sizeof(buff), "CLSID\\%s", clsid);
        RegDB regDB("Classes", buff, "InProcServer32");
        regDB.Write(NULL, libPath);
        regDB.Write("ThreadingModel", "Both");
        WriteDefaultValue(buff, "ProgID", qualifiedName);
        WriteDefaultValue(buff, "Programmable");
        WriteDefaultValue(buff, "TypeLib", libid);
      }
      else if (TAW.m_typeAttr->typekind == TKIND_DISPATCH)
      {
        strcpy_s(buff, sizeof(buff), "Interface\\");
        GUIDAsText(buff + strlen(buff), TAW.m_typeAttr->guid);
        StrWrap tiNameS = tiName.m_bstr;
        WriteDefaultValue(buff, NULL, tiNameS.m_str);
        WriteDefaultValue(buff, "ProxyStubClsid", "{00020424-0000-0000-C000-000000000046}");
        WriteDefaultValue(buff, "ProxyStubClsid32", "{00020424-0000-0000-C000-000000000046}");
        RegDB regDB("Classes", buff, "TypeLib");
        regDB.Write(NULL, libid);
        regDB.Write("Version", verStr);
      }
    }
    return NOERROR;
  }
  catch (char *emsg)
  {
    MessageBoxA(NULL, emsg, "DllRegisterServer", MB_ICONEXCLAMATION);
    return SELFREG_E_CLASS;
  }
}

// This is called by regsvr32 /u, and removes the registry entries for the COM objects defined here ...

HRESULT __stdcall DllUnregisterServer()
{
  try
  {
    char iid[39], libPath[256];
    if (!GetModuleFileNameA(g_hModule, libPath, sizeof(libPath)))
      ThrowError("Unable to get file name for module %08X", g_hModule);
    TLIBATTRWrap tLA;
    ThrowError(g_pTypeLib->GetLibAttr(&tLA.m_typeLibAttr));
    GUIDAsText(iid, tLA.m_typeLibAttr->guid);
    try {RegDBDeleteKey("Classes", "TypeLib", iid);}
    catch (char *) {}
    BSTRWrap libName;
    ThrowError(g_pTypeLib->GetDocumentation(MEMBERID_NIL, &libName.m_bstr, NULL, NULL, NULL));
    for (int j = 0;; j++)
    {
      SmartPtr<ITypeInfo> typeInfo;
      if (g_pTypeLib->GetTypeInfo(j, &typeInfo.m_obj)) break;
      BSTRWrap tiName;
      ThrowError(typeInfo.m_obj->GetDocumentation(MEMBERID_NIL, &tiName.m_bstr, NULL, NULL, NULL));
      TYPEATTRWrap TAW(typeInfo.m_obj);
      if (!TAW.m_typeAttr) continue;
      GUIDAsText(iid, TAW.m_typeAttr->guid);
      if (TAW.m_typeAttr->typekind == TKIND_COCLASS)
      {
        char qualifiedName[128];
        sprintf_s(qualifiedName, sizeof(qualifiedName), "%S.%S", libName.m_bstr, tiName.m_bstr);
        try {RegDBDeleteKey("Classes", qualifiedName, NULL);}
        catch (char *) {}
        try {RegDBDeleteKey("Classes", "CLSID", iid);}
        catch (char *) {}
      }
      else if (TAW.m_typeAttr->typekind == TKIND_DISPATCH)
      {
        try {RegDBDeleteKey("Classes", "Interface", iid);}
        catch (char *) {}
      }
    }
    return NOERROR;
  }
  catch (char *)
  {
    return SELFREG_E_CLASS;
  }
}

// * * * * * * Custom class definitions start here * * * * * *

// TestObj

// GUID of "interface" definition in Type Library ...
const GUID IID_TestObj = {0x7c8721d6, 0x3d22, 0x48a1, 0xa9, 0x45, 0x5f, 0xf9, 0x81, 0x5c, 0x58, 0x07};

struct TestObj : SimpleDispatch
{
  BSTRWrap m_name;
  double m_value;

  const GUID &IID_This() const {return IID_TestObj;}

// The following methods must match the corresponding entries in the IDL file

  virtual HRESULT __stdcall get_Name(BSTRWrap &name) const
  {
    name = m_name;
    return NOERROR;
  }

  virtual HRESULT __stdcall put_Name(const wchar_t *name)
  {
    m_name = name;
    return NOERROR;
  }

  virtual HRESULT __stdcall get_Value(double &value) const
  {
    value = m_value;
    return NOERROR;
  }

  virtual HRESULT __stdcall put_Value(double value)
  {
    m_value = value;
    return NOERROR;
  }

	virtual HRESULT __stdcall Square(double &square) const
  {
    square = m_value * m_value;
    return NOERROR;
  }
};

// GUID of "coclass" definition in Type Library ...
const GUID CLSID_TestObj = {0x5fc711f1, 0xb9c7, 0x4dcc, 0x8c, 0xcc, 0xe3, 0x9f, 0x9e, 0xf, 0x75, 0x56};

SimpleClassFactory<TestObj> TestObjCF(CLSID_TestObj);

// repeat for further classes. This one supplies two worksheet-callable functions to Excel ...

const GUID IID_TestWorksheetFuncs = {0x7c2676f0, 0x142f, 0x4898, { 0xad, 0x85, 0x40, 0x4f, 0x52, 0xaf, 0x95, 0x59 } };

struct TestWorksheetFuncs : SimpleDispatch
{
  const GUID &IID_This() const {return IID_TestWorksheetFuncs;}

// The following methods must match the corresponding entries in the IDL file

  virtual HRESULT __stdcall AddTwoNumbers(VARIANTARG *Number1, VARIANTARG *Number2, VARIANT *ret) const
  {
    try
    {
      ThrowError(VariantChangeType(Number1, Number1, 0, VT_R8));
      ThrowError(VariantChangeType(Number2, Number2, 0, VT_R8));
      ret->vt = VT_R8;
      ret->dblVal = Number1->dblVal + Number2->dblVal;
    }
    catch (char *emsg)
    {
      wchar_t wbuff[80];
      swprintf_s(wbuff, 80, L"#%S!", emsg);
      ret->vt = VT_BSTR;
      ret->bstrVal = SysAllocString(wbuff);
    }
    return NOERROR;
  }

  virtual HRESULT __stdcall JoinTwoStrings(VARIANT *String1, VARIANT *String2, VARIANT *ret) const
  {
    try
    {
      ThrowError(VariantChangeType(String1, String1, 0, VT_BSTR));
      ThrowError(VariantChangeType(String2, String2, 0, VT_BSTR));
      ret->vt = VT_BSTR;
      int n1 = SysStringLen(String1->bstrVal), n2 = SysStringLen(String2->bstrVal);
      ret->bstrVal = SysAllocStringLen(String1->bstrVal, n1 + n2);
      memcpy(ret->bstrVal + n1, String2->bstrVal, n2 * sizeof(wchar_t));
    }
    catch (char *emsg)
    {
      wchar_t wbuff[80];
      swprintf_s(wbuff, 80, L"#%S!", emsg);
      ret->vt = VT_BSTR;
      ret->bstrVal = SysAllocString(wbuff);
    }
    return NOERROR;
  }
};

const GUID CLSID_TestWorksheetFuncs = {0xaba0bf16, 0x855f, 0x475a, { 0x8d, 0x8d, 0x3d, 0xa0, 0x2c, 0xfa, 0xfd, 0xac}};

SimpleClassFactory<TestWorksheetFuncs> TestWorksheetFuncsCF(CLSID_TestWorksheetFuncs);