/***************************************************************************
         Copyright (c) Microsoft Corporation, All rights reserved.             
    This code sample is provided "AS IS" without warranty of any kind, 
    it is not recommended for use in a production environment.
***************************************************************************/

#include "project.h"

#include "dllmain.h"
#include "../vs_haskell_ui/resource.h" //IDB_BABELICONS
#include "package.h"  
#include "resource.h"  

/*---------------------------------------------------------
  globals
-----------------------------------------------------------*/
HINSTANCE  g_instance             = NULL;

ITypeInfo* g_typeInfoColorSink    = NULL;
ITypeInfo* g_typeInfoParseSink    = NULL;
ITypeInfo* g_typeInfoScope        = NULL;
ITypeInfo* g_typeInfoDeclarations = NULL;
ITypeInfo* g_typeInfoNames        = NULL;
ITypeInfo* g_typeInfoMethods      = NULL;
ITypeInfo* g_typeInfoBabelPackage = NULL;

const CLSID clsidBabelPackage = 
{ 0x6525805, 0xcbf0, 0x4726, { 0x9b, 0x50, 0x18, 0x4d, 0x2b, 0xbe, 0xa2, 0xc6 } };

HIMAGELIST g_imageList       = NULL;
HMODULE g_hResourceDllHandle = NULL;

const int  ImageWidth          = 16;		
const int  ImageGrowFactor     = 4;

/*---------------------------------------------------------
  local globals
-----------------------------------------------------------*/
static ULONG g_dllRefCount     = 0;

//registry paths
static CStr regPathVS        = STR("Software\\Microsoft\\VisualStudio\\7.1");
static CStr pathPackages     = STR("Packages");
static CStr babelDescription = STR("Visual Haskell Babel Package");

/*---------------------------------------------------------
  Prototypes
-----------------------------------------------------------*/
static HRESULT RegisterServer(bool registerIt);

static HRESULT LoadTypeInfo();
static HRESULT DoneTypeInfo();

static HRESULT LoadResources();
static HRESULT DoneResources();

/*---------------------------------------------------------
  Open vstudio root
-----------------------------------------------------------*/
HRESULT RegOpenKeyVS( out HKEY* key )
{
  return WINERROR(RegOpenKey( HKEY_LOCAL_MACHINE, regPathVS, key ));
}


/*---------------------------------------------------------
  Class Factory
-----------------------------------------------------------*/
class PackageFactory : IClassFactory
{
private:
  ULONG   m_refCount;

public:
  PackageFactory();
  ~PackageFactory();

  //IUnknown
  STDMETHODIMP QueryInterface( REFIID iid, LPVOID* obj );
  STDMETHODIMP_(ULONG) AddRef();
  STDMETHODIMP_(ULONG) Release();

  //IClassFactory
  STDMETHODIMP CreateInstance( IUnknown* outer, REFIID iid, LPVOID* obj );
  STDMETHODIMP LockServer( BOOL lockIt );
};



/*---------------------------------------------------------
  Dll entry points
-----------------------------------------------------------*/
BOOL APIENTRY DllMain( HINSTANCE instance, DWORD reason, LPVOID reserved )
{
  HRESULT hr = S_OK;

  switch (reason)
  {
  case DLL_PROCESS_ATTACH:     
    //get clsid as string
    OleChar oleCLSID[MAXCLSID];
    StringFromGUID2( clsidBabelPackage, oleCLSID, MAXCLSID);

    g_instance    = instance;
    g_dllRefCount = 0;
    hr = LoadTypeInfo();  ASSERT(SUCCEEDED(hr));
    hr = LoadResources(); ASSERT(SUCCEEDED(hr));
    break;

  case DLL_PROCESS_DETACH:
    TRACE("unload dll");
    DoneResources();
    DoneTypeInfo();
    TRACE_DONE();
    break;
  case DLL_THREAD_ATTACH:
    TRACE("dll attach thread");
    break;
  case DLL_THREAD_DETACH:
    TRACE("dll detach thread");
    break;
  }

  return TRUE;
}

STDAPI DllCanUnloadNow()
{
  return (g_dllRefCount == 0);
}


STDAPI DllGetClassObject( REFCLSID clsid, REFIID iid, LPVOID* obj )
{
  OUTARG(obj);

  if (clsid == clsidBabelPackage)
  {
    PackageFactory*  factory  = new PackageFactory();
    if (!factory) return E_OUTOFMEMORY;

    HRESULT hr = factory->QueryInterface(iid,obj);
    RELEASE(factory);
    return hr;
  }
  else 
    return CLASS_E_CLASSNOTAVAILABLE;
}


STDAPI DllRegisterServer()
{
  return RegisterServer(true);
}


STDAPI DllUnregisterServer()
{
  return RegisterServer(false);
}

STDAPI DllInstall(BOOL bInstall, LPCWSTR pszCmdLine)
{
	int len1, len2;
	Str newRegPathVS;

	if (!pszCmdLine)
		return S_OK;

	len1 = strLen(regPathVS);
	len2 = strLen(pszCmdLine);

	newRegPathVS = (Str) malloc((len1+len2+1)*sizeof(Char));
	if (!newRegPathVS)
		return HRESULT_FROM_WIN32(ERROR_NOT_ENOUGH_MEMORY);

	strCopy(newRegPathVS,        regPathVS,  len1);
	strCopy(newRegPathVS+len1-3, pszCmdLine, len2);

	regPathVS = newRegPathVS;
	return (bInstall ? DllRegisterServer() : DllUnregisterServer());
};

/*---------------------------------------------------------
  Class Factory
-----------------------------------------------------------*/
PackageFactory::PackageFactory()
{
  TRACE_CREATE("PackageFactory", this );
  m_refCount = 1;
}

PackageFactory::~PackageFactory()
{
  TRACE_DESTROY(this);
}

//IUnknown
STDMETHODIMP PackageFactory::QueryInterface( REFIID iid, LPVOID* obj )
{
  OUTARG(obj);

  if (iid == IID_IUnknown || iid == IID_IClassFactory)
  {
    *obj = static_cast<IClassFactory*>(this);
  }
  else
    return E_NOINTERFACE;

  AddRef();
  return S_OK;
}

STDMETHODIMP_(ULONG) PackageFactory::AddRef()
{
  return IncRefCount(&m_refCount);
}

STDMETHODIMP_(ULONG) PackageFactory::Release()
{
  if (DecRefCount(&m_refCount) == 0)
  {
    delete this;
    return 0;
  }
  else
    return m_refCount;
}

//IClassFactory
STDMETHODIMP PackageFactory::CreateInstance( IUnknown* outer, REFIID iid, LPVOID* obj )
{
  TRACE("PackageFactory::CreateInstance");
  OUTARG(obj);
  HRESULT hr;

  if (outer) return CLASS_E_NOAGGREGATION;

  Package* package = new Package();
  if (!package) return E_OUTOFMEMORY;

  // Package init must be performed during the IVsPackage::SetSite call
  // because the package should open the registry root for the appid(devenv)
  // instance that loaded this package
  // hr = package->Init();
  // if (FAILED(hr)) { RELEASE(package); return hr; }

  hr = package->QueryInterface( iid, obj );
  RELEASE(package);

  return hr;
}

STDMETHODIMP PackageFactory::LockServer( BOOL lockIt )
{
  if (lockIt)
  {
    IncRefCount( &g_dllRefCount );
  }
  else
  {
    DecRefCount( &g_dllRefCount );
  }

  return S_OK;
}


/*---------------------------------------------------------
  Load resources
-----------------------------------------------------------*/
static HRESULT LoadResources()
{
  Char path[MAXPATH+1];
  DWORD dwLen = ::GetModuleFileName(g_instance, path, MAXPATH);
  if (dwLen == 0)
    return HRESULT_FROM_WIN32(::GetLastError());
  if (dwLen == MAXPATH)
    return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
   
  //change the file name
  Str fileName = strFindCharRev( path, OLE('\\') );
  if (!fileName) return E_FAIL;
    
  fileName[1] = 0;
  strCat( path, STR("1033\\vs_haskell_ui.dll"), MAXPATH);

  //and try to load from the pathname
  g_hResourceDllHandle = LoadLibrary(Str2COle(path));
  if (!g_hResourceDllHandle) return HRESULT_FROM_WIN32(::GetLastError());

  g_imageList = ImageList_LoadBitmap(g_hResourceDllHandle, MAKEINTRESOURCE( IDB_BABEL_ICONS ), 
			                          ImageWidth, ImageGrowFactor, RGB( 0, 255, 0 ));
  if (!g_imageList) return E_FAIL;

  return S_OK;
}

static HRESULT DoneResources()
{
  if (g_imageList) 
  {
    ImageList_Destroy( g_imageList );
    g_imageList = NULL;
  }

  if (g_hResourceDllHandle)
  {
	  FreeLibrary(g_hResourceDllHandle);
	  g_hResourceDllHandle = NULL;
  }

  return S_OK;
}

/*---------------------------------------------------------
  Load the BabelService type library
-----------------------------------------------------------*/
static HRESULT LoadTypeInfo()
{
  HRESULT hr;

  //load it only once
  if (g_typeInfoColorSink) return S_OK;

  //load the library
  ITypeLib* lib;
  hr = LoadRegTypeLib( LIBID_BabelServiceLib, 1, 0, LANG_NEUTRAL, &lib );

  //if not registered, load from filepath
  if (FAILED(hr))  
  {
    USECONV;
	Char path[MAXPATH+1];
    DWORD dwLen = ::GetModuleFileName(g_instance, path, MAXPATH);
    if (dwLen == 0)
      return HRESULT_FROM_WIN32(::GetLastError());
    if (dwLen == MAXPATH)
      return HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
   
    //change the file name
    Str fileName = strFindCharRev( path, OLE('\\') );
    if (!fileName) return E_FAIL;
    
    fileName[1] = 0;
    strCat( path, STR("..\\babelservice.tlb"), MAXPATH );

    //and try to load from the pathname
    hr = LoadTypeLibEx( Str2COle(path), REGKIND_REGISTER, &lib );
    if (FAILED(hr)) return hr;
  }


  //get the typeinfo 
  if (!g_typeInfoColorSink)
  {
    hr = lib->GetTypeInfoOfGuid( IID_IColorSink, &g_typeInfoColorSink );
    if (FAILED(hr)) { RELEASE(lib); return hr; }
  }

  if (!g_typeInfoParseSink)
  {
    hr = lib->GetTypeInfoOfGuid( IID_IParseSink, &g_typeInfoParseSink );
    if (FAILED(hr)) { RELEASE(lib); return hr; }
  }

  if (!g_typeInfoScope)
  {
    hr = lib->GetTypeInfoOfGuid( IID_IScope, &g_typeInfoScope );
    if (FAILED(hr)) { RELEASE(lib); return hr; }
  }

  if (!g_typeInfoDeclarations)
  {
    hr = lib->GetTypeInfoOfGuid( IID_IDeclarations, &g_typeInfoDeclarations );
    if (FAILED(hr)) { RELEASE(lib); return hr; }
  }

  if (!g_typeInfoNames)
  {
    hr = lib->GetTypeInfoOfGuid( IID_INames, &g_typeInfoNames );
    if (FAILED(hr)) { RELEASE(lib); return hr; }
  }

  if (!g_typeInfoMethods)
  {
    hr = lib->GetTypeInfoOfGuid( IID_IMethods, &g_typeInfoMethods );
    if (FAILED(hr)) { RELEASE(lib); return hr; }
  }

  if (!g_typeInfoBabelPackage)
  {
    hr = lib->GetTypeInfoOfGuid( IID_IBabelPackage, &g_typeInfoBabelPackage );
    if (FAILED(hr)) { RELEASE(lib); return hr; }
  }

  RELEASE(lib);
  return S_OK;
}

static HRESULT DoneTypeInfo()
{
  RELEASE(g_typeInfoBabelPackage);
  RELEASE(g_typeInfoMethods);
  RELEASE(g_typeInfoNames);
  RELEASE(g_typeInfoDeclarations);
  RELEASE(g_typeInfoColorSink);
  RELEASE(g_typeInfoParseSink);
  RELEASE(g_typeInfoScope);
  return S_OK;
}

/*---------------------------------------------------------
  VS Registration
-----------------------------------------------------------*/
static HRESULT RegisterServer(bool registerIt)
{
  HRESULT hr;
  HKEY    keyRoot;
  HKEY    keySub = 0;
  Char    pathSub[MAXPATH+1];
  CStr    str;
  DWORD   dwID;

  hr = RegOpenKeyVS( &keyRoot );
  if (FAILED(hr)) return hr;
  
  
  //get clsid as string
  OleChar oleCLSID[MAXCLSID];
  if (StringFromGUID2(clsidBabelPackage, oleCLSID, MAXCLSID) == 0)
  {
    hr = E_FAIL;
	goto fail;
  }

  //get module path
  Char modulePath[MAXPATH+1];
  DWORD dwLen = ::GetModuleFileName(g_instance, modulePath, MAXPATH);
  if (dwLen == 0)
  {
    hr = HRESULT_FROM_WIN32(::GetLastError());
	goto fail;
  }
  if (dwLen == MAXPATH)
  {
    hr = HRESULT_FROM_WIN32(ERROR_INSUFFICIENT_BUFFER);
	goto fail;
  }

  //package path
  strCopy( pathSub, pathPackages, MAXPATH );
  strCat ( pathSub, STR("\\"),    MAXPATH );
  oleToStrCat( pathSub, oleCLSID, MAXPATH );

  if (!registerIt)
  {
    RegDeleteKey(keyRoot, pathSub);
	UnRegisterTypeLib(LIBID_BabelServiceLib, 1, 0, LANG_NEUTRAL, SYS_WIN32);
  }
  else
  {
    hr = WINERROR(RegCreateKeyEx(keyRoot, pathSub, 0, NULL, REG_OPTION_NON_VOLATILE, 
                                 KEY_ALL_ACCESS, NULL, &keySub, NULL));
    if (FAILED(hr)) goto fail;

    hr = WINERROR(RegSetValueEx(keySub, NULL, 0, REG_SZ, (BYTE*)(babelDescription), (strLen(babelDescription)+1)*sizeof(Char)));
    if (FAILED(hr)) goto fail;

    hr = WINERROR(RegSetValueEx(keySub, STR("InprocServer32"), 0, REG_SZ, (BYTE*)(modulePath), (strLen(modulePath)+1)*sizeof(Char)));
    if (FAILED(hr)) goto fail;

	str = STR("Microsoft");
	hr = WINERROR(RegSetValueEx(keySub, STR("CompanyName"), 0, REG_SZ, (BYTE*)(str), (strLen(str)+1)*sizeof(Char)));
    if (FAILED(hr)) goto fail;

	str = STR("visual haskell babel package");
	hr = WINERROR(RegSetValueEx(keySub, STR("ProductName"), 0, REG_SZ, (BYTE*)(str), (strLen(str)+1)*sizeof(Char)));
    if (FAILED(hr)) goto fail;

	str = STR("1.0");
	hr = WINERROR(RegSetValueEx(keySub, STR("ProductVersion"), 0, REG_SZ, (BYTE*)(str), (strLen(str)+1)*sizeof(Char)));
    if (FAILED(hr)) goto fail;

	str = STR("standard");
	hr = WINERROR(RegSetValueEx(keySub, STR("MinEdition"), 0, REG_SZ, (BYTE*)(str), (strLen(str)+1)*sizeof(Char)));
    if (FAILED(hr)) goto fail;

	dwID = IDS_BABEL_PACKAGE_LOAD_KEY;
	hr = WINERROR(RegSetValueEx(keySub, STR("ID"), 0, REG_DWORD, (BYTE*)(&dwID), sizeof(dwID)));
    if (FAILED(hr)) goto fail;

    RegCloseKey(keySub);
    keySub = NULL;
  }

  RegCloseKey(keyRoot);
  keyRoot = NULL;
  return S_OK;

fail:
  if (keySub)  RegCloseKey(keySub);
  if (keyRoot) RegCloseKey(keyRoot);

  return hr;
}