ATL PersistXML implementation

// CAtlPersistXMLDlg.cpp : Implementation of CAtlPersistXMLDlg

/////////////////////////////////////////////////////
//											
// This is a part of the ATL PersistXML add-on
// 
// Copyright (c) VitalyKatasonov.  All rights reserved.
//
// free to use everywhere!
//
// any comments send to:  vkatasonov@yahoo.com
//
/////////////////////////////////////////////////////

#include "stdafx.h"
#include "CAtlPersistXMLDlg.h"
#include <CommCtrl.H>
#include <WindowsX.H>
#include "..\PersistVarXml\IPersistVarXMLImpl.h"
#include <time.h>
#include <stdio.h>

/////////////////////////////////////////////////////////////////////////////
// CAtlPersistXMLDlg


LRESULT CAtlPersistXMLDlg::OnInitDialog(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled)
{
	CoInitialize(NULL); 

	srand( (unsigned)time( NULL ) );

	HICON hi;

	m_ImageList = ImageList_Create(16,16,ILC_MASK,10,1);

	hi = ::LoadIcon(_Module.m_hInstResource,MAKEINTRESOURCE(IDI_SYSTEM));
	m_iSysten = ImageList_AddIcon(m_ImageList,hi);

	hi = ::LoadIcon(_Module.m_hInstResource,MAKEINTRESOURCE(IDI_USER));
	m_iUser = ImageList_AddIcon(m_ImageList,hi);

	hi = ::LoadIcon(_Module.m_hInstResource,MAKEINTRESOURCE(IDI_USER1));
	ImageList_AddIcon(m_ImageList,hi);

	hi = ::LoadIcon(_Module.m_hInstResource,MAKEINTRESOURCE(IDI_USER2));
	ImageList_AddIcon(m_ImageList,hi);

	hi = ::LoadIcon(_Module.m_hInstResource,MAKEINTRESOURCE(IDI_USER3));
	ImageList_AddIcon(m_ImageList,hi);

	HWND hTree = GetDlgItem(IDC_TREE);

	TreeView_SetImageList(hTree,m_ImageList,TVSIL_NORMAL);//LVSIL_SMALL);


	HRESULT hRes;
	if(SUCCEEDED(hRes= m_RootSystem.CoCreateInstance(__uuidof(MySystem))))
	{
		HWND hTree = GetDlgItem(IDC_TREE);
		m_RootSystem->BindToTree((long)hTree,(long)TVI_ROOT);

		if(SUCCEEDED(ATLIPersistFile_Load(L"AtlPersistXMLSample.ini",m_RootSystem)))
			MessageBox(_T("The system persistence was successfully restored\nfrom AtlPersistXMLSample.ini (XML format)!"),
				_T("AtlPersistXMLSample Application"));

	}

	return 1;  // Let the system set the focus
}


LRESULT CAtlPersistXMLDlg::OnAddUser(WORD wNotifyCode, WORD wID, HWND hWndCtl, BOOL& bHandled)
{
	USES_CONVERSION;
	HRESULT hRes;
	HWND hTree = GetDlgItem(IDC_TREE);

	TVITEM tvItem;
	tvItem.hItem = TreeView_GetSelection(hTree);
	tvItem.mask = TVIF_PARAM;
	if(TreeView_GetItem(hTree,&tvItem))
	{
		if(tvItem.lParam)
		{
			CComQIPtr<IMySystem> pSystem((IUnknown*)(tvItem.lParam));
			if(pSystem)
			{
				int n;
				
				TCHAR UserName[255];
				_stprintf(UserName,_T("User%d"),n = rand());
				
				TCHAR SystemName[255];
				_stprintf(SystemName,_T("System%d"),n);
				
				CComPtr<ISystemUser> pNewUser;
				if(SUCCEEDED(hRes = pNewUser.CoCreateInstance(__uuidof(SystemUser))))
				{
					pNewUser->put_Name(T2OLE(UserName));
				
					hRes = pSystem->AddUser(pNewUser);


					CComPtr<IMySystem> pSystem2;
					hRes = pNewUser->get_System((IDispatch**)&pSystem2);
					pSystem2->put_Name(T2OLE(SystemName));
					pSystem2->put_Price((float)(rand()/100.));

					CComQIPtr<IPersistVarXML> pPersistXML(pSystem);
					if(pPersistXML)
						UpdateXMLText(pPersistXML);

					::UpdateWindow(hTree);
				}
			}
		}
	}

	return 0L;
}

LRESULT CAtlPersistXMLDlg::OnDeleteUser(WORD wNotifyCode, WORD wID, HWND hWndCtl, BOOL& bHandled)
{
	HRESULT hRes;
	HWND hTree = GetDlgItem(IDC_TREE);

	TVITEM tvItem;
	tvItem.hItem = TreeView_GetSelection(hTree);
	tvItem.mask = TVIF_PARAM;
	if(TreeView_GetItem(hTree,&tvItem))
	{
		if(tvItem.lParam)
		{
			CComQIPtr<ISystemUser> pUser((IUnknown*)(tvItem.lParam));
			if(pUser)
			{
				tvItem.hItem = TreeView_GetParent(hTree,tvItem.hItem);
				if(TreeView_GetItem(hTree,&tvItem))
				{
					CComQIPtr<IMySystem> pSystem((IUnknown*)(tvItem.lParam));
					if(pSystem)
						hRes = pSystem->DeleteUser(pUser);
				}
			}
		}
	}

	return 0L;
}


void CAtlPersistXMLDlg::UpdateXMLText(IPersistVarXML* pPersistXML)
{
	if(pPersistXML)
	{
		VARIANT VarOutput;
		V_VT(&VarOutput)= VT_VARIANT|VT_BYREF;

		CComVariant varBSTR(L"");
		V_VARIANTREF(&VarOutput) = &varBSTR;

		pPersistXML->SaveXML(VarOutput);
		if(V_VT(&varBSTR)==VT_BSTR)
		{
			USES_CONVERSION;
			SetDlgItemText(IDC_XML,OLE2T(V_BSTR(&varBSTR)));
		}
	}
}

LRESULT CAtlPersistXMLDlg::OnSelchangedTree(int idCtrl, LPNMHDR pnmh, BOOL& bHandled)
{
	NMTREEVIEW* pTreeView = (NMTREEVIEW*)pnmh; 
	CComQIPtr<IPersistVarXML> pPersistXML((IUnknown*)(pTreeView->itemNew.lParam));
	UpdateXMLText(pPersistXML);

	HWND hAddUserBtn = GetDlgItem(IDADDUSER);
	HWND hDeleteUserBtn = GetDlgItem(IDDELETEUSER);

	CComQIPtr<IMySystem> pSystem(pPersistXML);
	if(pSystem)
	{
		::EnableWindow(hAddUserBtn, TRUE);
		::EnableWindow(hDeleteUserBtn, FALSE);
	}
	else
	{
		::EnableWindow(hAddUserBtn, FALSE);
		::EnableWindow(hDeleteUserBtn, TRUE);
	}
	return 0L;
}


LRESULT CAtlPersistXMLDlg::OnOK(WORD wNotifyCode, WORD wID, HWND hWndCtl, BOOL& bHandled)
{
	EndDialog(wID);
	return 0;
}


LRESULT CAtlPersistXMLDlg::OnDestroy(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled)
{
	if(SUCCEEDED(ATLIPersistFile_Save(L"AtlPersistXMLSample.ini",m_RootSystem)))
		MessageBox(_T("The system persistence was successfully saved\nin AtlPersistXMLSample.ini (XML format)!"),
		_T("AtlPersistXMLSample Application"));

	ImageList_Destroy(m_ImageList);
	CoUninitialize();
	return 0L;
}



// IPersistFile Help functions! :)
HRESULT ATLIPersistFile_Load(LPCOLESTR pszFileName, IUnknown* pUnk)
{
	USES_CONVERSION;
	HRESULT hRes = E_FAIL;

	LPCTSTR strFileName= OLE2CT(pszFileName);
	HANDLE hFile = CreateFile(strFileName, GENERIC_READ, 0, NULL,
							  OPEN_EXISTING,
							  FILE_ATTRIBUTE_ARCHIVE,
							  NULL);

	if (INVALID_HANDLE_VALUE == hFile)
		return HRESULT_FROM_WIN32(GetLastError());

	bool bXMLFormat = false;
	DWORD dwHiFileSize;
	DWORD dwSize = GetFileSize(hFile,&dwHiFileSize);
	if(dwSize)
	{
		HGLOBAL hGlob= GlobalAlloc(GHND, dwSize);
		LPVOID pString  = GlobalLock(hGlob);
		DWORD dwReaded;
		BOOL bRet= ReadFile(hFile,pString,dwSize,&dwReaded,NULL);

		//auto recognition XML format (can be improved)
		if( ((char*)pString)[0] == '<' || ( ((byte*)pString)[0] == 0xFF && ((byte*)pString)[1] == 0xFE && ((char*)pString)[2] =='<' ))
		{
			bXMLFormat = true;
		}
		GlobalUnlock(hGlob);
		CloseHandle(hFile);

		DWORD dwError;
		if(!bRet)
			hRes = HRESULT_FROM_WIN32(dwError = GetLastError());

		CComPtr<IStream> spStream;
		hRes = CreateStreamOnHGlobal(hGlob, TRUE, &spStream);
		if (spStream)
		{
			if(bXMLFormat)
			{
				//load XML
				CComQIPtr<IPersistVarXML> pIPersistVarXML(pUnk);
				if(pIPersistVarXML)
					hRes = pIPersistVarXML->LoadXML(CComVariant(spStream));
				else
				{
					ATLASSERT(0);
					hRes = E_FAIL;
				}
			}
			else
			{
				//load STREAM
				CComQIPtr<IPersistStreamInit> pIPersistStream(pUnk);
				if(pIPersistStream)
					hRes = pIPersistStream->Load(spStream);
				else
				{
					ATLASSERT(0);
					hRes = E_FAIL;
				}
			}
			ATLASSERT(SUCCEEDED(hRes));
			spStream.Release();
		}
	}else
		CloseHandle(hFile);
	return hRes;
}

HRESULT ATLIPersistFile_Save( LPCOLESTR pszFileName, IUnknown* pUnk)
{
	USES_CONVERSION;
	HRESULT hRes = S_OK;

	LPCTSTR strFileName= OLE2CT(pszFileName);
	HANDLE hFile = CreateFile(strFileName, GENERIC_WRITE, 0, NULL,
							  CREATE_ALWAYS,
							  FILE_ATTRIBUTE_ARCHIVE,
							  NULL);
	if (INVALID_HANDLE_VALUE == hFile)
		return HRESULT_FROM_WIN32(GetLastError());

	//Determinate Format to Save from File Name!
	int	bTryToSaveAsXML=true;
	long Size = _tcsclen(strFileName);
	if(!_tcsncicmp(&strFileName[Size-4],_T(".bin"),4))
			bTryToSaveAsXML = false;

	DWORD dwWrited;

	HGLOBAL hGlob;// = GlobalAlloc(GHND, l.LowPart);
	CComPtr<IStream> spStream;
	hRes = CreateStreamOnHGlobal(NULL/*hGlob*/, TRUE, &spStream);
	if (spStream)
	{
		if(bTryToSaveAsXML)
		{
			//try to save XML
			CComQIPtr<IPersistVarXML> pIPersistVarXML(pUnk);
			if(pIPersistVarXML)
				hRes = pIPersistVarXML->SaveXML(CComVariant(spStream));
			else
			{
				ATLASSERT(0);
				hRes = E_FAIL;
			}
		}
		
		if(FAILED(hRes) || !bTryToSaveAsXML)
		{
			//try to save STREAM
			CComQIPtr<IPersistStreamInit> pIPersistStream(pUnk);
			if(pIPersistStream)
				hRes = pIPersistStream->Save(spStream,TRUE);
			else
			{
				ATLASSERT(0);
				hRes = E_FAIL;
			}
		}
		ATLASSERT(SUCCEEDED(hRes));

		if(SUCCEEDED(hRes))
		{
			hRes = GetHGlobalFromStream(spStream,&hGlob);
			DWORD dwSize = GlobalSize(hGlob);
			LPVOID pString  = GlobalLock(hGlob);

			BOOL bRet= WriteFile(hFile,pString,dwSize,&dwWrited,NULL);

			ATLASSERT(bRet || dwWrited!=dwSize);

			GlobalUnlock(hGlob);

			DWORD dwError;
			if(!bRet)
				hRes = HRESULT_FROM_WIN32(dwError = GetLastError());
		}
		spStream.Release();
	}
	CloseHandle(hFile);
	return hRes;
}