/*
 * Copyright (c) 2000, 2002, Jason Tishler.
 *
 *     This program is free software; you can redistribute it and/or modify
 *     it under the terms of the GNU General Public License as published by
 *     the Free Software Foundation; either version 2 of the License, or
 *     (at your option) any later version.
 *
 *     A copy of the GNU General Public License can be found at
 *     http://www.gnu.org/
 *
 * Written by Jason Tishler  <jason@tishler.net>
 *
 * $Id: mksc.cpp,v 1.3 2002/03/02 16:12:11 jt Exp $
 */

#include <windows.h>
#include <iostream>
#include <string>
#include <shlguid.h>
#include <shlobj.h>

using namespace std;

DWORD MakeShortcut(
	const string& aPath,
	const string& anArguments,
	const string& aWorkingDirectory,
	const string& anIconLocation,
	int anIconIndex,
	const string& aShortcutPath);
void ParseArgs(
	int argc,
	char* argv[], 
	string& aPath,
	string& anArguments,
	string& aWorkingDirectory,
	string& anIconLocation,
	int& anIconIndex,
	string& aShortcutPath);
void PrintUsage();
string CollapseBackslashes(const char* aString);

int
main(int argc, char* argv[])
{
	string aPath;
	string anArguments;
	string aWorkingDirectory;
	string anIconLocation;
	int anIconIndex = 0;
	string aShortcutPath;

	// Parse command line arguments
	ParseArgs(
		argc,
		argv,
		aPath,
		anArguments,
		aWorkingDirectory,
		anIconLocation,
		anIconIndex,
		aShortcutPath);

	// Create shortcut
	DWORD aStatus = MakeShortcut(
		aPath,
		anArguments,
		aWorkingDirectory,
		anIconLocation,
		anIconIndex,
		aShortcutPath);
	if (aStatus != ERROR_SUCCESS)
	{
		cerr << "MakeShortcut failed with error = " << aStatus << endl;
	}

	return (aStatus == ERROR_SUCCESS) ? 0 : 2;
}

DWORD
MakeShortcut(
	const string& aPath,
	const string& anArguments,
	const string& aWorkingDirectory,
	const string& anIconLocation,
	int anIconIndex,
	const string& aShortcutPath)
{
	// Variable definitions located here due to use of gotos
	IShellLink *aShellLink = 0;
	IPersistFile* aPersistFile = 0;
	int aStatus2 = 0;
	DWORD aReturnCode = ERROR_SUCCESS;

	// Initialize COM
	HRESULT aStatus = CoInitialize(0);
	if (FAILED(aStatus))
	{
		aReturnCode = GetLastError();
		cerr << "CoInitialize failed with error = " << aReturnCode << endl;
		goto exit;
	}

	// Create shell link instance
	aStatus = CoCreateInstance(
		CLSID_ShellLink,
		0,
		CLSCTX_INPROC_SERVER,
		IID_IShellLink,
		(void **) &aShellLink);
	if (FAILED(aStatus))
	{
		aReturnCode = GetLastError();
		cerr << "CoCreateInstance failed with error = " << aReturnCode << endl;
		goto exit;
	}

	// Set the shell link's command path
	aStatus = aShellLink->SetPath(aPath.c_str());
	if (FAILED(aStatus))
	{
		aReturnCode = aStatus;
		cerr << "IShellLink::SetPath failed with HRESULT = " << aStatus << endl;
		goto exit;
	}

	// Set the shell link's command arguments
	aStatus = aShellLink->SetArguments(anArguments.c_str());
	if (FAILED(aStatus))
	{
		aReturnCode = aStatus;
		cerr << "IShellLink::SetArguments failed with HRESULT = " << aStatus << endl;
		goto exit;
	}

	// Set the shell link's working directory, if necessary
	if (aWorkingDirectory.size() > 0)
	{
		aStatus = aShellLink->SetWorkingDirectory(aWorkingDirectory.c_str());
		if (FAILED(aStatus))
		{
			aReturnCode = aStatus;
			cerr << "IShellLink::SetWorkingDirectory failed with HRESULT = " << aStatus << endl;
			goto exit;
		}
	}

	// Set the shell link's icon, if necessary
	if (anIconLocation.size() > 0)
	{
		aStatus = aShellLink->SetIconLocation(anIconLocation.c_str(), anIconIndex);
		if (FAILED(aStatus))
		{
			aReturnCode = aStatus;
			cerr << "IShellLink::SetIconLocation failed with HRESULT = " << aStatus << endl;
			goto exit;
		}
	}

	// Get the shell link's persist file interface
	aStatus = aShellLink->QueryInterface(
		IID_IPersistFile,
		(void **) &aPersistFile);
	if (FAILED(aStatus))
	{
		aReturnCode = GetLastError();
		cerr << "IShellLink::QueryInterface failed with error = " << aReturnCode << endl;
		goto exit;
	}

	// Convert the shortcut path to Unicode
	WCHAR aWideShortcutPath[MAX_PATH];
	aStatus2 = MultiByteToWideChar(
		CP_ACP,
		0,
		aShortcutPath.c_str(),
		-1,
		aWideShortcutPath,
		MAX_PATH);
	if (aStatus2 == 0)
	{
		aReturnCode = GetLastError();
		cerr << "MultiByteToWideChar failed with error = " << aReturnCode << endl;
		goto exit;
	}

	// Save the shell link to the filesytem
	aStatus = aPersistFile->Save(aWideShortcutPath, TRUE);
	if (FAILED(aStatus))
	{
		aReturnCode = GetLastError();
		cerr << "IPersistFile::Save failed with error = " << aReturnCode << endl;
		goto exit;
	}

exit:
	// Clean up
	if (aPersistFile)
		aPersistFile->Release();
	if (aShellLink)
		aShellLink->Release();
	CoUninitialize();

	return aReturnCode;
}

void
ParseArgs(
	int argc,
	char* argv[], 
	string& aPath,
	string& anArguments,
	string& aWorkingDirectory,
	string& anIconLocation,
	int& anIconIndex,
	string& aShortcutPath)
{
	// Print usage, if appropriate
	if (argc == 1)
	{
		PrintUsage();
		exit(1);
	}

	// Process arguments
	for (int i = 1; i < argc; i += 2)
	{
		if (i + 1 >= argc)
		{
			PrintUsage();
			exit(1);
		}

		if (strcmp(argv[i], "--cmd") == 0)
		{
			aPath = CollapseBackslashes(argv[i + 1]);
		}
		else if (strcmp(argv[i], "--args") == 0)
		{
			anArguments = CollapseBackslashes(argv[i + 1]);
		}
		else if (strcmp(argv[i], "--dir") == 0)
		{
			aWorkingDirectory = CollapseBackslashes(argv[i + 1]);
		}
		else if (strcmp(argv[i], "--icon") == 0)
		{
			anIconLocation = CollapseBackslashes(argv[i + 1]);
		}
		else if (strcmp(argv[i], "--index") == 0)
		{
			anIconIndex = atoi(argv[i + 1]);
		}
		else if (strcmp(argv[i], "--shortcut") == 0)
		{
			aShortcutPath = CollapseBackslashes(argv[i + 1]);
		}
		else
		{
			PrintUsage();
			exit(1);
		}
	}

	// Check for mandatory arguments
	if (aPath.size() == 0 || anArguments.size() == 0 || aShortcutPath.size() == 0)
	{
		PrintUsage();
		exit(1);
	}
}

void
PrintUsage()
{
	cerr << "usage: mksc --cmd CommandPath --args CommandArguments" << endl;
	cerr << "    --shortcut ShortcutPath [--dir WorkingDirectory]" << endl;
	cerr << "    [--icon IconPath] [--index IconIndex]" << endl;
}

string
CollapseBackslashes(const char* aString)
{
	const string theDoubleBackslashRe = "\\\\";
	const string theSpaceRe = " ";

	string aFixedString = aString;

	// Determine if string does not contains any spaces, if so just return
	if (string::size_type i = aFixedString.find(theSpaceRe) == string::npos)
		return aFixedString;

	// Collapse double backslashes into single
	for (string::size_type i = aFixedString.find(theDoubleBackslashRe);
		i != string::npos;
		i = aFixedString.find(theDoubleBackslashRe, i))
	{
		i++;
		aFixedString.erase(i, 1);
	}

	return aFixedString;
}