#pragma once

#ifndef RS_PACKAGEHELPER_H_INCLUDED
#define RS_PACKAGEHELPER_H_INCLUDED

/**
    Implement method of this type in your DLL to act as Rosetta package.
    Method has to be named GetRosettaPackage
    Example (from Kernel package):

    ///// BEGIN OF EXAMPLE - COPY AFTER THIS POINT ////

    #ifdef __cplusplus    // If used by C++ code, 
    extern "C" {          // we need to export the C interface
    #endif

    __declspec(dllexport) HRESULT __cdecl GetRosettaPackage(IRsUnknownPackage **ppPackage)
    {
        CComObject<CRpKernel> *pPackage = NULL;
        HRESULT hRes = CComObject<CRpKernel>::CreateInstance(&pPackage);
        if (FAILED(hRes))
            return hRes;
        CComQIPtr<IRsUnknownPackage> spPackage = pPackage;
        if (!pPackage)
            return E_NOINTERFACE;
        return spPackage.CopyTo(ppPackage);
    }

    #ifdef __cplusplus
    }
    #endif

    ///// END OF EXAMPLE ////
**/
typedef HRESULT (__cdecl *Rosetta_GetRosettaPackage)(IRsUnknownPackage **ppPackage);

/**
Rosetta package helper is here to help you to define package and succesfully load it.
*/
class CRsRosettaPackageHelper
{

public:
    CRsRosettaPackageHelper()
    {
        m_hPackageHandle = NULL;
        m_bUnloadOnDestruct = false;
    }

    CRsRosettaPackageHelper(PCTSTR pszPathList, PCTSTR pszModuleName, bool bUnloadOnDestruct)
    {
        m_hPackageHandle = NULL;
        m_bUnloadOnDestruct = false;
        HRESULT hRes = Load(pszPathList, pszModuleName, bUnloadOnDestruct);
        _ASSERTE(SUCCEEDED(hRes));
    }

    ~CRsRosettaPackageHelper()
    {
        if (m_bUnloadOnDestruct)
        {
            HRESULT hRes = Unload();
            _ASSERTE(SUCCEEDED(hRes));
        }
    }

    static CString GetApplicationPath()
    {
        TCHAR       tchBuffer[MAX_PATH];
        TCHAR       *lpszDelimiter;
        GetModuleFileName(NULL, tchBuffer, MAX_PATH);
        if (NULL != (lpszDelimiter = _tcsrchr(tchBuffer, _T('\\'))))
            *lpszDelimiter = 0;
        CString szPath = CString(tchBuffer) + CString("\\");
        return szPath;
    }

    HRESULT Load(PCTSTR pszPathList, PCTSTR pszModuleName, bool bUnloadOnDestruct)
    {
        Unload();

        CString szPathList = CString(pszPathList) + CString(_T(";")) + GetApplicationPath();
        int position = 0;
        CString path = szPathList.Tokenize(";", position);
        while (!path.IsEmpty())
        {
            CString szPath = path;
            szPath.Trim();
            if ((szPath.Right(1) != _T("\\")) && (szPath.Right(1) != _T("/")))
                szPath+="\\";
            CString szFullModuleName = szPath + CString(pszModuleName);

            // Does the file exist (to prevent DLL injection)
            HANDLE hFile = ::CreateFile(szFullModuleName, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
            if (hFile != (HANDLE)(0) && hFile != (HANDLE)(-1))
            {
                ::CloseHandle(hFile);
                HMODULE m_hPackageHandle = ::LoadLibraryEx(szFullModuleName, NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
                if (!m_hPackageHandle)
                {
                    _ASSERTE(!"Package exists in given path but we are not able to load it");
                }
                else
                {
                    m_szPackageFile = szFullModuleName;
                    Rosetta_GetRosettaPackage Method_GetRosettaPackage = (Rosetta_GetRosettaPackage)::GetProcAddress(m_hPackageHandle, _T("GetRosettaPackage"));
                    if (!Method_GetRosettaPackage)
                    {
                        _ASSERTE(!"Package exists, can be loaded, but does not export GetRosettaPackage method!");
                    }
                    else
                    {
                        HRESULT hRes = Method_GetRosettaPackage(&m_spPackage);
                        if ((FAILED(hRes)) || (!m_spPackage))
                        {
                            _ASSERTE(!"Package failed to return IRsUnknownPackage instance!");
                        }
                        else
                        {
                            m_bUnloadOnDestruct = bUnloadOnDestruct; 
                            return S_OK;
                        }
                    }
                }
            }

            Unload();
            path = szPathList.Tokenize(";", position);
        }
        _ASSERTE(!"Package not loaded");
        return E_INVALIDARG;
    }

	HRESULT Load(PCTSTR pszFullModuleName, bool bUnloadOnDestruct)
	{
		Unload();

		CString szFullModuleName(pszFullModuleName);

		// Does the file exist (to prevent DLL injection)
		HANDLE hFile = ::CreateFile(szFullModuleName, GENERIC_READ, FILE_SHARE_READ, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
		if (hFile != (HANDLE)(0) && hFile != (HANDLE)(-1))
		{
			::CloseHandle(hFile);
			m_hPackageHandle = ::LoadLibraryEx(szFullModuleName, NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
			if (!m_hPackageHandle)
			{
				_ASSERTE(!"Package exists in given path but we are not able to load it");
			}
			else
			{
				m_szPackageFile = szFullModuleName;
				Rosetta_GetRosettaPackage Method_GetRosettaPackage = (Rosetta_GetRosettaPackage)::GetProcAddress(m_hPackageHandle, _T("GetRosettaPackage"));
				if (!Method_GetRosettaPackage)
				{
					_ASSERTE(!"Package exists, can be loaded, but does not export GetRosettaPackage method!");
				}
				else
				{
					HRESULT hRes = Method_GetRosettaPackage(&m_spPackage);
					if ((FAILED(hRes)) || (!m_spPackage))
					{
						_ASSERTE(!"Package failed to return IRsUnknownPackage instance!");
					}
					else
					{
						m_bUnloadOnDestruct = bUnloadOnDestruct; 
						return S_OK;
					}
				}
			}
		}

		Unload();

		_ASSERTE(!"Package not loaded");
		return E_INVALIDARG;
	}


	HRESULT Unload()
	{
		if (m_hPackageHandle)
		{
			m_spPackage.Release();
			BOOL result = FreeLibrary(m_hPackageHandle);
			m_szPackageFile.Empty();
			m_hPackageHandle = NULL;
			m_bUnloadOnDestruct = false;
			if (!result)
				return E_FAIL;
		}
		return S_OK;
	}

	CString GetPackageFile()
	{
        return m_szPackageFile;
    }

    HRESULT GetPackage(IRsUnknownPackage **ppPackage)
    {
        if (ppPackage == NULL)
            return E_INVALIDARG;

        if (m_spPackage == NULL)
            return E_UNEXPECTED;

        return m_spPackage.CopyTo(ppPackage);
    }

private:
    bool                            m_bUnloadOnDestruct;
    CString                         m_szPackageFile;
    HMODULE                         m_hPackageHandle;
    CComPtr<IRsUnknownPackage>      m_spPackage;
};

#endif //RS_PACKAGEHELPER_H_INCLUDED

