//
// lnk - Shell link tool.
//

#include <windows.h>
#include <shlobj.h>

#include <string.h>
#include <assert.h>

#include "lnk.h"

ShellLinkData::ShellLinkData() :
	_path(NULL),
	_linkw(NULL),
	_desc(NULL),
	_arg(NULL),
	_dir(NULL),
	_iconp(NULL),
	_iconx(0),
	_show(0),
	_hotkey(0)
{
}

ShellLinkData::~ShellLinkData()
{
	delete[] _path;
	delete[] _linkw;
	delete[] _desc;
	delete[] _arg;
	delete[] _dir;
	delete[] _iconp;
}

inline
void OleCheck(HRESULT hOleResult)
{
	if (!SUCCEEDED(hOleResult)) throw hOleResult;
}

BOOL ShellLinkData::Init(LPCSTR path, LPCSTR link)
{
	assert(path);
	if (!link) link = path;

	LPSTR path_new = NULL;
	LPWSTR linkw_new = NULL;
	try {
		// Copy the file and link names.
		//
		path_new = new CHAR[::lstrlenA(path) + 1];
		::lstrcpyA(path_new, path);

		linkw_new = new WCHAR[::lstrlenA(link) + 1];
		::MultiByteToWideChar(CP_ACP, 0, link, ::lstrlenA(link), linkw_new, ::lstrlenA(link) + 1);
	}
	catch (...) {
		delete[] path_new;
		delete[] linkw_new;
		::SetLastError(ERROR_NOT_ENOUGH_MEMORY);
		return FALSE;
	}

	delete[] _path;
	delete[] _linkw;
	_path = path_new;
	_linkw = linkw_new;

	return TRUE;
}

BOOL ShellLinkData::SetArguments(LPCSTR arg)
{
	LPSTR arg_new = NULL;
	try {
		if (arg) {
			arg_new = new CHAR[::lstrlenA(arg) + 1];
			::lstrcpyA(arg_new, arg);
		}
	}
	catch (...) {
		delete[] arg_new;
		::SetLastError(ERROR_NOT_ENOUGH_MEMORY);
		return FALSE;
	}

	delete[] _arg;
	_arg = arg_new;

	return TRUE;
}

BOOL ShellLinkData::SetDirectory(LPCSTR dir)
{
	LPSTR dir_new = NULL;
	try {
		if (dir) {
			dir_new = new CHAR[::lstrlenA(dir) + 1];
			::lstrcpyA(dir_new, dir);
		}
	}
	catch (...) {
		delete[] dir_new;
		::SetLastError(ERROR_NOT_ENOUGH_MEMORY);
		return FALSE;
	}

	delete[] _dir;
	_dir = dir_new;

	return TRUE;
}

BOOL ShellLinkData::SetDescription(LPCSTR desc)
{
	LPSTR desc_new = NULL;
	try {
		if (desc) {
			desc_new = new CHAR[::lstrlenA(desc) + 1];
			::lstrcpyA(desc_new, desc);
		}
	}
	catch (...) {
		delete[] desc_new;
		::SetLastError(ERROR_NOT_ENOUGH_MEMORY);
		return FALSE;
	}

	delete[] _desc;
	_desc = desc_new;

	return TRUE;
}

BOOL ShellLinkData::SetIcon(LPCSTR iconpath, int iconindex)
{
	LPSTR iconp_new = NULL;
	int iconx_new = 0;
	try {
		if (!iconp_new) {
			iconp_new = new CHAR[::lstrlenA(iconpath) + 1];
			::lstrcpyA(iconp_new, iconpath);
			iconx_new = iconindex;
		}
	}
	catch (...) {
		delete[] iconp_new;
		::SetLastError(ERROR_NOT_ENOUGH_MEMORY);
		return FALSE;
	}

	delete[] _iconp;
	_iconp = iconp_new;
	_iconx = iconx_new;

	return TRUE;
}

BOOL ShellLinkData::SetShowCommand(int show)
{
	_show = show;
	return TRUE;
}

BOOL ShellLinkData::SetHotkey(WORD hotkey)
{
	_hotkey = hotkey;
	return TRUE;
}

BOOL ShellLinkData::Register()
{
	if (!_path || !_linkw) {
		::SetLastError(ERROR_INVALID_DATA);
		return FALSE;
	}

	// Interfaces to be used. (must released after use)
	//
	IShellLink* piLink = NULL;
	IPersistFile* piFile = NULL;
	try {
		// Initialize the OLE.
		//
		::OleCheck(::CoInitialize(NULL));

		// Create a shell link instance.
		//
		::OleCheck(::CoCreateInstance(CLSID_ShellLink,
									  NULL,
									  CLSCTX_INPROC_SERVER,
									  IID_IShellLink,
									  (void**)&piLink));

		// Set the properties.
		//
		piLink->SetPath(_path);
		if (_arg) piLink->SetArguments(_arg);
		if (_desc) piLink->SetDescription(_desc);
		if (_iconp) piLink->SetIconLocation(_iconp, _iconx);
		if (_dir) piLink->SetWorkingDirectory(_dir);
		piLink->SetShowCmd(_show);
		piLink->SetHotkey(_hotkey);

		// Save the instance.
		//
		::OleCheck(piLink->QueryInterface(IID_IPersistFile,
										  (void**)&piFile));
		::OleCheck(piFile->Save(_linkw, TRUE));

		throw S_OK;
	}
	catch (HRESULT hOleResult) {
		// Free the interfaces.
		//
		if (piLink) piLink->Release();
		if (piFile) piFile->Release();

		// Uninitialize the OLE.
		//
		::CoUninitialize();
		if (hOleResult != S_OK) {
			::SetLastError(hOleResult);
			return FALSE;
		}
	}

	return TRUE;
}