///////////////////////////////////////////////////////////////////////////
//  component.cpp - Computes prime numbers.  Provides a custom class     //
//  ver 1.0         object that takes the starting prime as an argument  //
//                                                                       //
//  Eddon and Eddon, Inside Distributed COM, Microsoft Press, 1998       //
//                                                                       //
//  Modifed by Jim Fawcett, Syracuse University, jfawcett@twcny.rr.com   //
//  CSE791 - Distributed Objects, Spring 2002                            //
///////////////////////////////////////////////////////////////////////////

#define _WIN32_DCOM
#include <windows.h>
#include <iostream>    // For cout
#include "registry.h"  // For registry functions
#include "Component\component.h" // Generated by MIDL
using namespace std;

HANDLE g_hEvent;

class CPrime : public IPrime
{
public:
  // IUnknown
  ULONG __stdcall AddRef();
  ULONG __stdcall Release();
  HRESULT __stdcall QueryInterface(REFIID riid, void** ppv);

  // IPrime
  HRESULT __stdcall GetNextPrime(int* retval);

  CPrime(int start_number);
  ~CPrime();

private:
  long m_cRef;
  int m_start_number;
};

void InitiateComponentShutdown()
{
  cout << "InitiateComponentShutdown()" << endl;
  SetEvent(g_hEvent);
}

CPrime::CPrime(int start_number) : m_cRef(1),  m_start_number(start_number)
{
  CoAddRefServerProcess();
}

CPrime::~CPrime()
{
  cout << "Component: CPrime::~CPrime()" << endl;
  if(CoReleaseServerProcess() == 0)
    InitiateComponentShutdown();
}

ULONG CPrime::AddRef()
{
  cout << "Component: CPrime::AddRef() m_cRef = " << m_cRef + 1 << endl;
  return InterlockedIncrement(&m_cRef);
}

ULONG CPrime::Release()
{
  cout << "Component: CPrime::Release() m_cRef = " << m_cRef - 1 << endl;
  unsigned cRef = InterlockedDecrement(&m_cRef);
  if(cRef != 0)
    return cRef;
  delete this;
  return 0;
}

HRESULT CPrime::QueryInterface(REFIID riid, void** ppv)
{
  if(riid == IID_IUnknown)
  {
    cout << "Component: CPrime::QueryInterface() for IUnknown" << endl;
    *ppv = (IUnknown*)this;
  }
  else if(riid == IID_IPrime)
  {
    cout << "Component: CInsideDCOM::QueryInterface() for IPrime" << endl;
    *ppv = (IPrime*)this;
  }
  else 
  {
    *ppv = NULL;
    return E_NOINTERFACE;
  }
  AddRef();
  return S_OK;
}

bool IsPrime(int TestNumber)
{
  int HalfNumber = TestNumber / 2 + 1;
  for(int count = 2; count < HalfNumber; count++)
    if(TestNumber % count == 0)
      return false;
  return true;
}

HRESULT CPrime::GetNextPrime(int* retval)
{
  cout << "Component: CPrime::GetNextPrime()" << endl;

  while(!IsPrime(++m_start_number));
  *retval = m_start_number;
  return S_OK;
}

class CPrimeFactory : public IPrimeFactory, public IExternalConnection
{
public:
  // IUnknown
  ULONG __stdcall AddRef();
  ULONG __stdcall Release();
  HRESULT __stdcall QueryInterface(REFIID riid, void** ppv);

  // IPrimeFactory
  HRESULT __stdcall CreatePrime(int start_number, IPrime** ppPrime);

  // IExternalConnection
  DWORD __stdcall AddConnection(DWORD extconn, DWORD dwreserved);
  DWORD __stdcall ReleaseConnection(
    DWORD extconn, DWORD dwreserved, BOOL fLastReleaseCloses
  );

  CPrimeFactory() : m_cRef(1) { }
  ~CPrimeFactory() { cout << "Factory destructor" << endl; }

private:
  long m_cRef;
};

ULONG CPrimeFactory::AddRef()
{
  cout << "Component: CPrimeFactory::AddRef() m_cRef = " 
       << m_cRef + 1 << endl;
  return InterlockedIncrement(&m_cRef);
}

ULONG CPrimeFactory::Release()
{
  unsigned cRef = InterlockedDecrement(&m_cRef);
  cout << "Component: CPrimeFactory::Release() m_cRef = " 
       << cRef << endl;
  if(cRef != 0)
    return cRef;
  delete this;
  return 0;
}

HRESULT CPrimeFactory::QueryInterface(REFIID riid, void** ppv)
{
  if(riid == IID_IUnknown)
  {
    cout << "Component: CPrimeFactory::QueryInteface() "
         << "for IUnknown or CPrimeFactory" << endl;
    *ppv = reinterpret_cast<IUnknown*>(this);
  }
  else if(riid == IID_IPrimeFactory)
  {
    cout << "Component: CPrimeFactory::QueryInteface() "
         << "for IUnknown or CPrimeFactory" << endl;
    *ppv = (IPrimeFactory*)this;
  }
  else if(riid == IID_IExternalConnection)
  {
    cout << "Component: CPrimeFactory::QueryInteface() "
         << "for IUnknown or CPrimeFactory" << endl;
    *ppv = (IExternalConnection*)this;
  }
  else
  {
    *ppv = NULL;
    return E_NOINTERFACE;
  }
  AddRef();
  return S_OK;
}

// Since we don't use IClassFactory, we have to help COM manage object
// lifetime.  These two functions are what we need to do that

DWORD CPrimeFactory::AddConnection(DWORD extconn, DWORD dwreserved)
{
  cout << "IExternalConnection::AddConnection" << endl;
  if(extconn & EXTCONN_STRONG)
    return CoAddRefServerProcess();
  return 0;
}

DWORD CPrimeFactory::ReleaseConnection(
  DWORD extconn, DWORD dwreserved, BOOL fLastReleaseCloses
)
{
  cout << "IExternalConnection::ReleaseConnection" << endl;
  if(CoReleaseServerProcess() == 0)
  {
    InitiateComponentShutdown();
    return 0;
  }
  return 1;
}

HRESULT CPrimeFactory::CreatePrime(int start_number, IPrime** ppPrime)
{
  CPrime* pPrime = new CPrime(start_number);
  cout << "Component: CPrimeFactory::CreatePrime()" << endl;

  if(pPrime == NULL)
    return E_OUTOFMEMORY;

  HRESULT hr = pPrime->QueryInterface(IID_IPrime, (void**)ppPrime);
  pPrime->Release();
  return hr;
}

void RegisterComponent()
{
  ITypeLib* pTypeLib;
  LoadTypeLibEx(L"component.exe", REGKIND_DEFAULT, &pTypeLib);
  RegisterServer("component.exe", CLSID_Prime, "Prime Number Sample", "Component.Prime", "Component.Prime.1", NULL);
}

void CommandLineParameters(int argc, char** argv)
{
  RegisterComponent();
  if(argc < 2)
  {
    cout << "No parameter, but registered anyway" << endl;
    exit(false);
  }
  char* szToken = strtok(argv[1], "-/"); 
  if(_stricmp(szToken, "RegServer") == 0)
  {
    RegisterComponent();
    cout << "RegServer" << endl;
    exit(true);
  }
  if(_stricmp(szToken, "UnregServer") == 0)
  {
    UnRegisterTypeLib(LIBID_Component, 1, 0, LANG_NEUTRAL, SYS_WIN32);
    UnregisterServer(CLSID_Prime, "Component.Prime", "Component.Prime.1");
    cout << "UnregServer" << endl;
    exit(true);
  }
  if(_stricmp(szToken, "Embedding") != 0)
  {
    cout << "Invalid parameter" << endl;
    exit(false);
  }
}

void main(int argc, char** argv)
{
  CommandLineParameters(argc, argv);

  cout << "Component: CoInitializeEx()" << endl;
  CoInitializeEx(NULL, COINIT_MULTITHREADED);

  IPrimeFactory *pPrimeFactory = new CPrimeFactory();

  cout << "Component: CoRegisterClassObject()" << endl;
  DWORD dwRegister;
  CoRegisterClassObject(CLSID_Prime, pPrimeFactory, CLSCTX_LOCAL_SERVER, REGCLS_MULTIPLEUSE, &dwRegister);

  g_hEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
  WaitForSingleObject(g_hEvent, INFINITE);
  CloseHandle(g_hEvent);

  CoRevokeClassObject(dwRegister);
  pPrimeFactory->Release();
  CoUninitialize();

  char ch;
  cin >> ch;
}