#!/usr/bin/env perl
# cpuctl - hotplug CPUs on Linux
use feature "say";

sub err {
	say STDERR "error: @_";
	our $err; ++$err;
}

sub uniq (@) { my %seen; grep {!$seen{$_}++} @_; }

sub range_split {
	my ($str) = @_;
	my @items;
	for (split /,/, $str) {
		if (/^(\d+)-(\d+)$/) {
			push @items, $1..$2;
		} elsif (/^\d+$/) {
			push @items, $_;
		} else {
			err "not a number: $_";
		}
	}
	return @items;
}

sub put {
	my ($file, $str) = @_;
	if (open my $fh, ">", $file) {
		say $fh $str;
		close $fh;
	} else {
		err "cannot open $file: $!";
	}
}

sub get {
	my ($file) = @_;
	if (open my $fh, "<", $file) {
		chomp(my $str = <$fh>);
		close $fh;
		return $str;
	} else {
		err "cannot open $file: $!";
	}
}

sub cpu {
	"/sys/devices/system/cpu/cpu".shift;
}

sub core_id {
	my $path = cpu(shift)."/topology/core_id";
	-e $path ? int get($path) : 0;
}

sub can_offline {
	my $path = cpu(shift)."/online";
	-e $path ? $path : undef;
}

sub is_online {
	my $path = can_offline(shift);
	$path ? int get($path) : 1;
}

sub put_online {
	my $path = can_offline(shift);
	put($path, shift) if $path;
}

sub all_cpus {
	grep {-e cpu($_)} 0..127;
}

sub all_cores {
	uniq map {core_id($_)} all_cpus();
}

sub core_cpus {
	my ($core) = @_;
	grep {core_id($_) == $core} all_cpus();
}

my $cmd = shift(@ARGV);

for ($cmd) {
	if (!defined $_) {
		for (all_cpus()) {
			my $core = core_id($_);
			my $status = is_online($_) ? "online" : "offline";
			$status .= " (fixed)" if !can_offline($_);
			say "cpu $_ (core $core): $status";
		}
	}
	elsif ($_ eq "get-cpus") {
		my @cpus = all_cpus();
		say "@cpus";
	}
	elsif ($_ eq "get-online-cpus") {
		my @cpus = grep {is_online($_)} all_cpus();
		say "@cpus";
	}
	elsif ($_ eq "online") {
		my @cpus = @ARGV ? map {range_split($_)} @ARGV : all_cpus();
		@cpus = grep {can_offline($_)} @cpus;
		if (@cpus) {
			say "Onlining CPUs: @cpus";
			put_online($_, "1") for @cpus;
		} else {
			err "no hotpluggable CPUs given";
		}
	}
	elsif ($_ eq "offline") {
		my @cpus = @ARGV ? map {range_split($_)} @ARGV : all_cpus();
		@cpus = grep {can_offline($_)} @cpus;
		if (@cpus) {
			say "Offlining CPUs: @cpus";
			put_online($_, "0") for @cpus;
		} else {
			err "no hotpluggable CPUs given";
		}
	}
	elsif ($_ eq "offline-core") {
		my @cores = @ARGV ? map {range_split($_)} @ARGV
			          : grep {$_ > 0} all_cores();
		if (@cores) {
			for my $core (@cores) {
				my @cpus = core_cpus($core);
				if (@cpus) {
					say "Offlining CPUs: @cpus (core $core)";
					put_online($_, "0") for @cpus;
				} else {
					err "no hotpluggable CPUs in core $core";
				}
			}
		} else {
			err "no cores given";
		}
	}
	elsif ($_ eq "smt") {
		my $path = "/sys/devices/system/cpu/smt";
		if (@ARGV) {
			if ($ARGV[0] =~ /^(on|yes|true|enable)$/) {
				say "Enabling SMT";
				put("$path/control", "on");
			}
			elsif ($ARGV[0] =~ /^(off|no|false|disable)$/) {
				say "Disabling SMT";
				put("$path/control", "off");
			}
			else {
				err "bad parameter '$ARGV[0]'";
			}
		} else {
			say "enabled: ".get($path."/control");
			say "active: ".(get($path."/active") ? "yes" : "no");
		}
	}
	elsif ($_ eq "governor") {
		if (@ARGV) {
			system("cpupower", "frequency-set", "-g", @ARGV);
		} else {
			system("cpupower", "frequency-info");
		}
	}
	elsif ($_ eq "gpu-governor") {
		my $path = "/sys/class/drm/card0/device/power_dpm_state";
		if (@ARGV) {
			put($path, $ARGV[1]);
		} else {
			say "current: ".get($path);
			say "available: balanced, battery";
		}
	}
}

exit !!$err;
