/*
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
    
    (c) 2011 Emanuele Oriani - ema at fastwebnet dot it
*/

#include <iostream>
#include <cerrno>
#include <cstdlib>
#include <cstdio>

static pid_t	pid = -1;
static int	cpu = -1;
static bool	reset = false;
static bool	cpunum = false;

// To get the number of available CPUs we run this 
// getaffinity BEFORE setting any affinity mask.
// We will get all the available CPUs.
static int num_cpus(void) {
	cpu_set_t cpu_set;
	CPU_ZERO(&cpu_set);
	const int max_cpus = 8*sizeof(cpu_set);
	if (0 == sched_getaffinity(getpid(), sizeof(cpu_set_t), &cpu_set)) {
		for(int i = 0; i < max_cpus; ++i)
			if (!CPU_ISSET(i, &cpu_set)) return i;
	}
	return max_cpus;
}

const int NCPUS = num_cpus();

static void usage(const char *img) {
	std::cerr << img << ", a simple program to set CPU affinity on Linux." << std::endl;
	std::cerr << "(c) 2011 Emanuele Oriani - ema at fastwebnet dot it" << std::endl;
	std::cerr << "\nUsage:" << std::endl;
	std::cerr << "      " << img << " [-p pid] [-c cpu] [-r]" << std::endl;
	std::cerr << "      p: specify a process pid" << std::endl;
	std::cerr << "      c: specify a cpu id (0-n)" << std::endl;
	std::cerr << "      r: reset (enable all cpus)" << std::endl;
	std::cerr << "      g: print cpus number" << std::endl;
}

static void opts(int argc, char *argv[]) {
	extern char	*optarg;
	int 		res = -1;
	while(EOF != (res = getopt(argc, argv, "p:c:rhg"))) {
		switch((char)res) {
			case 'p':
				pid = (pid_t)std::atoi(optarg);
				break;
			case 'c':
				cpu = std::atoi(optarg);
				break;
			case 'r':
				reset = true;
				break;
			case 'g':
				cpunum = true;
				break;
			case 'h':
			default:
				usage(argv[0]);
				std::exit(-2);
				break;
		}
	} 
}

int main(int argc, char *argv[]) {
	opts(argc, argv);
	cpu_set_t cpu_set;
	if (cpunum) std::cout << "Available cpus: " << NCPUS << '.' << std::endl;
	if (-1 == pid) pid = getpid();
	if (reset || (-1 != cpu)) {
		CPU_ZERO(&cpu_set);
		if (reset) {
			for(int i = 0; i < NCPUS; ++i)
				CPU_SET(i, &cpu_set);
		}
		else CPU_SET(cpu, &cpu_set);
		if (-1 == sched_setaffinity(pid, sizeof(cpu_set_t), &cpu_set)) {
			std::cerr << "Error : ";
			switch(errno) {
			case EFAULT:
				std::cerr << "A supplied memory address was invalid." << std::endl;
				std::exit(-1);
				break;
			case EINVAL:
				std::cerr << "The affinity bitmask mask contains no processors that are physically on the system, or cpusetsize is smaller than the size of the affinity mask used by the kernel." << std::endl;
				std::exit(-1);
				break;
			case EPERM:
				std::cerr << "The calling process does not have appropriate privileges." << std::endl;
				std::exit(-1);
				break;
			case ESRCH:
				std::cerr << "The process whose ID is " << pid << " could not be found." << std::endl;
				std::exit(-1);
				break;
			default:
				std::cerr << "Unknown error." << std::endl;
				std::exit(-1);
				break;
			}
		} 
	}

	CPU_ZERO(&cpu_set);
	if (0 == sched_getaffinity(pid, sizeof(cpu_set_t), &cpu_set)) {
		std::cout << "Report:" << std::endl;
		for(int i = 0; i < NCPUS; ++i)
			if (CPU_ISSET(i, &cpu_set)) std::cout << "Cpu " << i << " is set." << std::endl;
	}
}