#include <windows.h>
#include <iostream>
#include <sstream>
#include <cstdlib>
#include <string>
#include <vector>
#include <algorithm>
#include <nsis/pluginapi.h> // nsis plugin

using namespace std;

extern "C" void __declspec(dllexport) ChangePath ( HWND hwndParent, int string_size, char *variables, stack_t **stacktop, extra_parameters *extra );

int  ModifyPath( const string &action, const string &type, string &element );
bool FindInPath( string &path, const string &search );
bool CleanPath ( string &path, const string &delpath = "" );

// GLOBALS
HINSTANCE   g_hInstance;
HWND        g_hwndParent;

static const string         ActionAddBefore("ADDBEFORE");
static const string         ActionAddAfter ("ADDAFTER");
static const string         ActionDelete   ("DELETE");

static const string         TypeSystem    ("SYSTEM");
static const string         TypeUser      ("USER");

// Parameters from NSIS
// ACTION : Action to perform 
//          ADDBEFORE : adds element to front of current path
//          ADDAFTER  : adds element to end of current path
//          DELETE    : deletes element from current path
// TYPE   : Specifies the "SYSTEM" or "USER" path environment to modify
// ELEMENT: The element to add/remove from path
// Returnes
// 0 on success
// 1 on failure
// 2 on invalid parameters
void __declspec(dllexport) ChangePath( HWND hwndParent, int string_size, char *variables, stack_t **stacktop, extra_parameters *extra )
{
    string  action, element, type;
    char    dummy[8152];

    g_hwndParent=hwndParent;

    EXDLL_INIT();
    
    if(popstring(dummy))
    {
        pushstring("2");
        return;
    }
    action = dummy;
    transform(action.begin(), action.end(), action.begin(), ::toupper);    
    if(popstring(dummy))
    {
        pushstring("2");
        return;
    }
    type = dummy;
    transform(type.begin(), type.end(), type.begin(), ::toupper);
    if(popstring(dummy))
    {
        pushstring("2");
        return;
    }
    element = dummy;
   
    if ( action != ActionAddBefore && action != ActionAddAfter && action != ActionDelete )
    {
        pushstring("2");
        return;
    }
    
    if ( type != TypeSystem && type != TypeUser )
    {
        pushstring("2");
        return;
    }

    // now we are good to good
    if ( ModifyPath(action, type, element) )
        pushstring("0");
    else
        pushstring("1");

    return;
}


int ModifyPath(const string &action, const string &type, string &element)
{
    string                      path;
    string::reverse_iterator    rit;    
    HKEY                        hKey;
    LONG                        res;
    char                        val[8152];
    DWORD                       size = sizeof(val);
    bool                        pathChanged = false;

    // Open the registry key for system or user enviroment
    if ( type == TypeSystem )
    {
        res = RegOpenKeyEx(HKEY_LOCAL_MACHINE, TEXT("SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment"), 0, KEY_ALL_ACCESS, &hKey);
    }
    else
    {
        res = RegOpenKeyEx(HKEY_CURRENT_USER, TEXT("Environment"), 0, KEY_ALL_ACCESS, &hKey);
    }
    if ( res != ERROR_SUCCESS )
        return 0;

    // Load the Path subkey
    res = RegQueryValueEx(hKey, TEXT("Path"), NULL, NULL, (LPBYTE)&val, &size);
    if ( res == ERROR_FILE_NOT_FOUND )
    {   // if not found set val to empty string
        val[0] = '\0';
    }
    else if ( res != ERROR_SUCCESS )
    {
        RegCloseKey(hKey);
        return 0;
    }
    path = val;

    // Remove any trailing ; from element
    for ( rit = element.rbegin(); *rit == ';'; rit = element.rbegin() )
        element.resize(element.length()-1);

    if ( action == ActionAddBefore && ! FindInPath(path, element) )
    {
        // make sure new dir to add has a ; on the end so we can insert it to the start of path 
        element += ';';
        path.insert(0, element);
        // Clean up any empty fields
        CleanPath(path);
        pathChanged = true;
    }
    else if ( action == ActionAddAfter && ! FindInPath(path, element) )
    {
        // append ; followed by new dir, clean path will remove empty fields so dont need to check for ; on the end
        path += ';';
        path += element;
        // Clean up any empty fields
        CleanPath(path);
        pathChanged = true;
    }
    else if ( action == ActionDelete )
    {
        // remove element from path if it exists
        pathChanged = CleanPath(path, element);
    }

    // Now write the path
    res = ERROR_SUCCESS;
    if ( pathChanged )
        res = RegSetValueEx(hKey, TEXT("Path"), NULL, REG_EXPAND_SZ, (LPBYTE)path.c_str(), path.length()+1);

    RegCloseKey(hKey);

    return ( res != ERROR_SUCCESS ? 0 : 1);
}


bool FindInPath( string &path, const string &search )
{
    string                      element;
    istringstream               spath(path);

    //Split each path element on ; and discard any empty ones
    while ( getline(spath, element, ';') )
    {
        if ( element == search )
            return true;
    }
    
    return false;
}


bool CleanPath( string &path, const string &delpath )
{
    string                      newpath("");
    string                      element;
    istringstream               spath(path);
    string::reverse_iterator    it;
    bool                        changed = false;

    //Split each path element on ; and discard any empty ones
    while ( getline(spath, element, ';') )
    {
        if ( ! element.empty() && ( delpath.empty() || element != delpath ))
        {
            newpath += element + ";";
        }
        else
        {
            changed = true;
        }
    }
    
    // Remove any trailing ;
    for ( it = newpath.rbegin(); *it == ';'; it = newpath.rbegin() )
    {
        newpath.resize(newpath.length()-1);
        changed = true;
    }

    if ( changed )
        path = newpath;
    
    return changed;
}


BOOL WINAPI DllMain(HANDLE hInst, ULONG ul_reason_for_call, LPVOID lpReserved)
{
    g_hInstance = (HINSTANCE)hInst;
    return TRUE;
}