/*
 *
 *
 *
 */

#include <windows.h>
#include <stdio.h>
#include "getopt.h"

#define ENVVARS_ROOT_SUBKEY_SYSTEM HKEY_LOCAL_MACHINE
#define ENVVARS_PATH_SUBKEY_SYSTEM \
	"SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment"

#define ENVVARS_ROOT_SUBKEY_USER HKEY_CURRENT_USER
#define ENVVARS_PATH_SUBKEY_USER "Environment"

#define ENVVARS_ROOT_SUBKEY_DEFAULT HKEY_USERS
#define ENVVARS_PATH_SUBKEY_DEFAULT ".DEFAULT\\Environment"

#define LEN_MAX_STDIN_SET 4096

char* NameVersion = "renv version 1.0";

HKEY EnvKeyRoot = ENVVARS_ROOT_SUBKEY_USER;
char* EnvNamePath = ENVVARS_PATH_SUBKEY_USER;
int FDoDelete = 0;
int FIsExpanded = 0;

void PrintError(char* msg)
{
	fputs(msg, stderr);
	fputc('\n', stderr);
}

void PrintError1(char* fmt, void* arg)
{
	fprintf(stderr, fmt, arg);
	fputc('\n', stderr);
}

int PrintVariable(HKEY envset, char* name)
{
	DWORD error;
	char* val;
	DWORD typ;
	DWORD sz = 0;

	if (error = RegQueryValueEx(envset, name, NULL, &typ, NULL, &sz)) {
		SetLastError(error);
		if (error == ERROR_FILE_NOT_FOUND) {
			PrintError("No such a variable.");
		} else if (error == ERROR_NOT_ENOUGH_MEMORY) {
			PrintError("Not enough memory.");
		} else if (error == ERROR_ACCESS_DENIED) {
			PrintError("Access denied.");
		}
		return 0;
	}

	if (!(val = (char*)malloc((int)sz))) {
		SetLastError(ERROR_NOT_ENOUGH_MEMORY);
		PrintError("Not enough memory.");
		return 0;
	}

	if (error = RegQueryValueEx(envset, name, NULL, &typ, val, &sz)) {
		SetLastError(error);
		if (error == ERROR_FILE_NOT_FOUND) {
			PrintError("No such a variable.");
		} else if (error == ERROR_NOT_ENOUGH_MEMORY) {
			PrintError("Not enough memory.");
		} else if (error == ERROR_ACCESS_DENIED) {
			PrintError("Access denied.");
		}
		free(val);
		return 0;
	}

	switch (typ) {
	case REG_EXPAND_SZ:
	case REG_SZ:
		fputs(val, stdout);
		break;
	default:;
	}

	free(val);

	return 1;
}

int PrintAllVariables(HKEY envset)
{
	DWORD error;
	DWORD index;
	DWORD cntKey, lenKey, lenCls;
	DWORD cntVal, lenVal, lenDat, lenSd;
	FILETIME ftWrt;
	char* name;

	if (error = RegQueryInfoKey(envset, NULL, NULL, NULL,
								&cntKey, &lenKey, &lenCls,
								&cntVal, &lenVal, &lenDat,
								&lenSd, &ftWrt)) {
		SetLastError(error);
		if (error == ERROR_NOT_ENOUGH_MEMORY) {
			PrintError("Not enough memory.");
		} else if (error == ERROR_ACCESS_DENIED) {
			PrintError("Access denied.");
		}
		return 0;
	}

	if (!(name = malloc(lenVal + 1))) {
		SetLastError(ERROR_NOT_ENOUGH_MEMORY);
		PrintError("Not enough memory.");
		return 0;
	}

	for (index = 0;; index++) {
		int szName = lenVal + 1;
		if (error = RegEnumValue(envset, index, name, &szName,
								 NULL, NULL, NULL, NULL)) {
			SetLastError(error);
			if (error == ERROR_NO_MORE_ITEMS) break;
			else if (error == ERROR_NOT_ENOUGH_MEMORY) {
				PrintError("Not enough memory.");
			} else if (error == ERROR_ACCESS_DENIED) {
				PrintError("Access denied.");
			}
			free(name);
			return 0;
		}
		fputs(name, stdout);
		fputc(',', stdout);
		PrintVariable(envset, name);
		fputc('\n', stdout);
	}

	free(name);

	return 1;
}

int SetVariable(HKEY envset, char* name, char* val)
{
	DWORD error;
	char* real_val = NULL;

	if (!strcmp(val, "-")) {
		if (!(val = real_val = malloc(LEN_MAX_STDIN_SET))) {
			SetLastError(ERROR_NOT_ENOUGH_MEMORY);
			PrintError("Not enough memory.");
			return 0;
		}
		if (!fgets(real_val, LEN_MAX_STDIN_SET, stdin)) {
			SetLastError(ERROR_INVALID_PARAMETER);
			PrintError("Cannot read value from standard input.");
			free(real_val);
			return 0;
		}
		if (real_val[strlen(real_val) - 1] == '\n') {
			real_val[strlen(real_val) - 1] = '\0';
		}
	}

	if (error = RegSetValueEx(envset, name, 0,
							  FIsExpanded ? REG_EXPAND_SZ : REG_SZ,
							  (LPVOID)val, strlen(val) + 1)) {
		SetLastError(error);
		if (error == ERROR_NOT_ENOUGH_MEMORY) {
			PrintError("Not enough memory.");
		} else if (error == ERROR_ACCESS_DENIED) {
			PrintError("Access denied.");
		}
		if (real_val) free(real_val);
		return 0;
	}

	if (real_val) free(real_val);

	return 1;
}

int DeleteVariable(HKEY envset, char* name)
{
	DWORD error;

	if (error = RegDeleteValue(envset, name)) {
		SetLastError(error);
		if (error == ERROR_FILE_NOT_FOUND) {
			PrintError("No such a variable.");
		} else if (error == ERROR_NOT_ENOUGH_MEMORY) {
			PrintError("Not enough memory.");
		} else if (error == ERROR_ACCESS_DENIED) {
			PrintError("Access denied.");
		}
		return 0;
	}

	return 1;
}

int ProcessVariable(char* name, char* val)
{
	HKEY envset;
	DWORD error;
	int result;

	if (error = RegOpenKeyEx(EnvKeyRoot, EnvNamePath, 0,
							  KEY_READ | KEY_WRITE,
							  &envset)) {
		SetLastError(error);
		if (error == ERROR_PATH_NOT_FOUND) {
			PrintError("Registry is corrupt.");
		} else if (error == ERROR_NOT_ENOUGH_MEMORY) {
			PrintError("Not enough memory.");
		} else if (error == ERROR_ACCESS_DENIED) {
			PrintError("Access denied.");
		}
		return 0;
	}

	if (!name) {
		result = PrintAllVariables(envset);
	} else if (FDoDelete) {
		result = DeleteVariable(envset, name);
	} else if (!val) {
		result = PrintVariable(envset, name);
	} else {
		result = SetVariable(envset, name, val);
	}

	RegCloseKey(envset);

	return result;
}

enum {
	LONGOPT_DELETE,
	LONGOPT_EXPANDED,
	LONGOPT_SYSTEM,
	LONGOPT_USER,
	LONGOPT_DEFAULT,
	LONGOPT_HELP,
	LONGOPT_VERSION,
	LONGOPT_MAX
};
#define SHORTOPT_DELETE 'd'
#define SHORTOPT_EXPANDED 'e'
#define SHORTOPT_SYSTEM 'S'
#define SHORTOPT_USER 'U'
#define SHORTOPT_DEFAULT 'D'
#define SHORTOPT_OPTIONSET "deSUD"
static struct option pLongOptions[] = {
	{ "delete", 0, 0, 0},
	{ "expanded", 0, 0, 0},
	{ "system", 0, 0, 0},
	{ "user", 0, 0, 0},
	{ "default", 0, 0, 0},
	{ "help", 0, 0, 0},
	{ "version", 0, 0, 0}
};

void PrintVersion()
{
	fprintf(stderr, "%s\n", NameVersion);
}

void PrintHelp()
{
	fputs(
		"usage: renv [OPTION] NAME [VALUE-TO-SET]\n"
		"Set/Print an environment variable.\n\n"
		"  -d, --delete   delete the variable\n"
		"  -e, --expanded set the variable as REG_EXPAND_SZ\n"
		"  -D, --default  set/print/delete the default user's variable\n"
		"  -S, --system   set/print/delete the system's variable\n"
		"  -U, --user     set/print/delete the current user's variable (default)\n"
		"  --help         display this help and exit\n"
		"  --version      output version information and exit\n", stderr);
}

int main (int argc, char **argv)
{
	int optchar;

	for (;;) {
		int option_index = 0;
		if ((optchar = getopt_long(argc, argv, SHORTOPT_OPTIONSET,
								   pLongOptions,
								   &option_index)) == EOF) break;
		switch (optchar) {
		case 0:
			switch (option_index) {
			case LONGOPT_DELETE:
				FDoDelete = 1;
				break;
			case LONGOPT_EXPANDED:
				FIsExpanded = 1;
				break;
			case LONGOPT_SYSTEM:
				EnvKeyRoot = ENVVARS_ROOT_SUBKEY_SYSTEM;
				EnvNamePath = ENVVARS_PATH_SUBKEY_SYSTEM;
				break;
			case LONGOPT_USER:
				EnvKeyRoot = ENVVARS_ROOT_SUBKEY_USER;
				EnvNamePath = ENVVARS_PATH_SUBKEY_USER;
				break;
			case LONGOPT_DEFAULT:
				EnvKeyRoot = ENVVARS_ROOT_SUBKEY_DEFAULT;
				EnvNamePath = ENVVARS_PATH_SUBKEY_DEFAULT;
				break;
			case LONGOPT_HELP:
				PrintHelp();
				return 2;
			case LONGOPT_VERSION:
				PrintVersion();
				return 2;
			default:;
			}
		case SHORTOPT_DELETE:
			FDoDelete = 1;
			break;
		case SHORTOPT_EXPANDED:
			FIsExpanded = 1;
			break;
		case SHORTOPT_SYSTEM:
			EnvKeyRoot = ENVVARS_ROOT_SUBKEY_SYSTEM;
			EnvNamePath = ENVVARS_PATH_SUBKEY_SYSTEM;
			break;
		case SHORTOPT_USER:
			EnvKeyRoot = ENVVARS_ROOT_SUBKEY_USER;
			EnvNamePath = ENVVARS_PATH_SUBKEY_USER;
			break;
		case SHORTOPT_DEFAULT:
			EnvKeyRoot = ENVVARS_ROOT_SUBKEY_DEFAULT;
			EnvNamePath = ENVVARS_PATH_SUBKEY_DEFAULT;
			break;
		default:
			PrintError1("Invalid option: %c", (void*)optchar);
			PrintHelp();
			return 2;
		}
    }

	if (optind >= argc) {
		if (!ProcessVariable(NULL, NULL)) {
			return 1;
		}
		return 0;
	}

	if (optind < argc - 2 ||
		(FDoDelete && optind == argc - 2)) {
		PrintError("Too many arguments.");
		PrintHelp();
		return 2;
	}

	if (!ProcessVariable(argv[optind], argv[optind + 1])) {
		return 1;
	}

	return 0;
}