#include "URIReg.h"
#include <QtGlobal>
#include <QCoreApplication>
#include <QSettings>
#include <QDir>
#include <QDebug>

URIReg URIReg::instance;

URIReg::URIReg() {

}

URIReg::~URIReg() {

}

URIReg& URIReg::getInstance() {
	return instance;
}

bool URIReg::reg(const QString& protocolName) {
#ifdef Q_OS_WIN
	// 获取当前工作目录
	QString appPath = QCoreApplication::applicationFilePath();

#if 0
	// 提升权限
	if (!runAsAdmin())
		return false;
#endif

#if 1
	if (!isRunningAsAdmin())
		return false;
#endif

	// 创建注册表项
#if 0  // 因应用操作注册表的权限问题，改成HKEY_CURRENT_USER
	QSettings settings("HKEY_CLASSES_ROOT\\" + protocolName, QSettings::NativeFormat);
#else
	QSettings settings("HKEY_CURRENT_USER\\Software\\Classes\\" + protocolName, QSettings::NativeFormat);
#endif 

	// 设置协议描述
	settings.setValue("Default", QString("URL:%1 Protocol").arg(protocolName));
	settings.sync();
	if (QSettings::NoError != settings.status())
		return false;

	// 标识这是一个URL协议
	settings.setValue("URL Protocol", "");
	settings.sync();
	if (QSettings::NoError != settings.status())
		return false;

	// 创建命令子项
	settings.beginGroup("shell");
	settings.beginGroup("open");
	settings.beginGroup("command");

	// 设置启动命令（注意转义 %1）
	QString command = QString("\"%1\" \"%2\"").arg(appPath).arg("%1");
	settings.setValue("Default", command);

	settings.endGroup();
	settings.endGroup();
	settings.endGroup();

	QSettings::Status s = settings.status();

	// 验证写入结果
	if (QSettings::NoError != settings.status())
		return false;

	return true;
#else
	// TODO:非Win平台下未实现
	assert(0);
	return false;
#endif
}

#ifdef Q_OS_WIN
#include <windows.h>
#include <shellapi.h>

bool URIReg::runAsAdmin() {
	wchar_t path[MAX_PATH];
	GetModuleFileNameW(nullptr, path, MAX_PATH);

	SHELLEXECUTEINFOW sei = { sizeof(sei) };
	sei.lpVerb = L"runas";
	sei.lpFile = path;
	sei.hwnd = nullptr;
	sei.nShow = SW_NORMAL;

	if (!ShellExecuteExW(&sei)) {
		qWarning() << QStringLiteral("提权失败:") << GetLastError();
		return false;
	}
	return true;
}

// 检查是否以管理员权限运行
bool URIReg::isRunningAsAdmin() {
	HANDLE hToken;
	BOOL isElevated = FALSE;
	if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken)) {
		TOKEN_ELEVATION elevation;
		DWORD cbSize = sizeof(TOKEN_ELEVATION);
		GetTokenInformation(hToken, TokenElevation, &elevation, cbSize, &cbSize);
		CloseHandle(hToken);
		isElevated = elevation.TokenIsElevated;
	}
	return isElevated != FALSE;
}

#include <aclapi.h>
#include <sddl.h> 

void URIReg::fixRegistryPermissions(const QString& keyPath) {
	HKEY hKey;
	LONG result = RegOpenKeyExW(
		HKEY_CURRENT_USER,
		keyPath.toStdWString().c_str(),
		0,                 // 选项，必须为0
		KEY_READ | KEY_WRITE,  // 读写权限,
		&hKey
	);

	if (result == ERROR_SUCCESS) {
		// 重置为默认权限
		PSECURITY_DESCRIPTOR sd;
		if (ConvertStringSecurityDescriptorToSecurityDescriptorW(
			L"D:(A;;KA;;;SY)(A;;KA;;;BA)(A;;KR;;;IU)",  // 标准用户权限
			SDDL_REVISION_1,
			&sd,
			nullptr)) {
			RegSetKeySecurity(hKey, DACL_SECURITY_INFORMATION, sd);
			LocalFree(sd);
		}
		RegCloseKey(hKey);
	}
}

bool URIReg::writeUrlProtocol(const QString& protocol) {
	HKEY hKey;
	DWORD disp;

	QString fullPath = QString("Software\\Classes\\") + protocol;

	LONG res = RegCreateKeyExW(
		HKEY_CURRENT_USER,
		fullPath.toStdWString().c_str(),
		0,
		nullptr,
		REG_OPTION_NON_VOLATILE,
		KEY_WRITE,
		nullptr,
		&hKey,
		&disp
	);

	if (res == ERROR_SUCCESS) {
		// 写入协议标识
		DWORD zero = 0;
		RegSetValueExW(hKey, L"URL Protocol", 0, REG_SZ, (BYTE*)L"", 2);

		// 写入描述
		std::wstring desc = L"URL:" + protocol.toStdWString() + L" Protocol";
		RegSetValueExW(hKey, nullptr, 0, REG_SZ, (BYTE*)desc.c_str(), (desc.size() + 1) * sizeof(wchar_t));

		RegCloseKey(hKey);
		return true;
	}
	return false;
}

#endif // Q_OS_WIN