// Force Intel GPUs into max power state on macOS
// Based off decompilation of Apple's MTLReplayer GTPMService.xpc
// clang++ -std=c++17 -O3 -framework IOKit IntelForceHighPower.cpp -o IntelForceHighPower

#include <cstdio>
#include <optional>
#include <IOKit/IOKitLib.h>

enum AGPMType {
	AGPMTypeIG = 2,
};

enum AGPMSelector {
	AGPMSelectorBeginCommands   = 0x1c85,
	AGPMSelectorEndCommands     = 0x1c86,
	AGPMSelectorGetPowerState   = 0x1c88,
	AGPMSelectorSetPowerState   = 0x1c89,
	AGPMSelectorGetControlState = 0x1c8a,
	AGPMSelectorSetControlState = 0x1c8b,
	AGPMSelectorGetType         = 0x1c91,
};

static io_connect_t FindService(const char* name, uint32_t type, bool (*match)(io_connect_t))
{
	CFMutableDictionaryRef dic = IOServiceMatching(name);
	io_iterator_t iter;
	if (IOServiceGetMatchingServices(0, dic, &iter) != kIOReturnSuccess)
		return 0;
	io_connect_t output = 0;
	while (io_object_t obj = IOIteratorNext(iter)) {
		io_connect_t con;
		if (IOServiceOpen(obj, mach_task_self(), type, &con) == kIOReturnSuccess) {
			if (!match || match(con))
				output = con;
			else
				IOServiceClose(con);
		}
		IOObjectRelease(obj);
		if (output)
			break;
	}
	IOObjectRelease(iter);
	return output;
}

static std::optional<uint64_t> CallGetter(io_connect_t service, AGPMSelector method)
{
	uint64_t value;
	uint32_t cnt = 1;
	if (IOConnectCallScalarMethod(service, method, nullptr, 0, &value, &cnt) != kIOReturnSuccess || cnt != 1)
		return std::nullopt;
	return value;
}

static bool IsIGService(io_connect_t service)
{
	auto res = CallGetter(service, AGPMSelectorGetType);
	return res.has_value() && *res == AGPMTypeIG;
}

static bool BeginCommands(io_connect_t service)
{
	return IOConnectCallScalarMethod(service, AGPMSelectorBeginCommands, nullptr, 0, nullptr, nullptr) == kIOReturnSuccess;
}

static bool EndCommands(io_connect_t service)
{
	return IOConnectCallScalarMethod(service, AGPMSelectorEndCommands, nullptr, 0, nullptr, nullptr) == kIOReturnSuccess;
}

static bool SetControlState(io_connect_t service, bool forced)
{
	uint64_t input = forced;
	uint64_t output;
	uint32_t cnt = 1;
	return IOConnectCallScalarMethod(service, AGPMSelectorSetControlState, &input, 1, &output, &cnt) == kIOReturnSuccess && cnt == 1;
}

static bool SetPowerState(io_connect_t service, uint32_t state)
{
	uint64_t input = state;
	return IOConnectCallScalarMethod(service, AGPMSelectorSetPowerState, &input, 1, nullptr, nullptr) == kIOReturnSuccess;
}

static void printUsage(const char* argv0) {
	fprintf(stderr, "Usage: %s (enable|disable)\n", argv0);
}

int main(int argc, const char * argv[]) {
	bool enable = false;
	if (argc <= 1) {
		printUsage(argv[0]);
		return EXIT_FAILURE;
	}
	if (0 == strcasecmp(argv[1], "enable")) {
		enable = true;
	} else if (0 == strcasecmp(argv[1], "disable")) {
		enable = false;
	} else {
		printUsage(argv[0]);
		return EXIT_FAILURE;
	}

	io_connect_t service = FindService("AGPM", 0, IsIGService);
	if (!service) {
		fprintf(stderr, "Failed to get AGPM service\n");
		return EXIT_FAILURE;
	}
	if (!BeginCommands(service)) {
		fprintf(stderr, "Failed to begin commands\n");
		return EXIT_FAILURE;
	}
	if (!SetControlState(service, enable)) {
		fprintf(stderr, "Failed to set control state\n");
		return EXIT_FAILURE;
	}
	// 0 is highest power, higher values are lower power
	if (enable && !SetPowerState(service, 0)) {
		fprintf(stderr, "Failed to set power state\n");
		return EXIT_FAILURE;
	}
	if (!EndCommands(service)) {
		fprintf(stderr, "Failed to end commands\n");
		return EXIT_FAILURE;
	}
	IOServiceClose(service);
	return EXIT_SUCCESS;
}