#include <stdio.h>
#include <unistd.h>
#include <getopt.h>
#include <string.h>
#include <sys/types.h>
#include <sched.h>
#include <sys/time.h>

/*
 * Gross hack, just in case glibc isn't up to date yet.
 * (it wasn't when SCHED_IDLE was added to the kernel)
 */
#ifndef SCHED_IDLE
#define SCHED_IDLE 3
#endif

static char *usage = "Usage: %s -n <priority> -p <pid> -d FF|TS|RR|ID [command...]\n";

static char progname[128];

main(int argc, char *argv[])
	{

	int minpri, maxpri;
	int ch;
	char *p;
	struct timeval interval;
	static struct sched_param param;
	pid_t pid = 0;
	int do_priority = 0;
	int do_policy = 0;
	int policy = 0;
	int priority = 0;
	int args = 0;
	extern char *optarg;
	extern int optind;
	
	(void) strcpy(progname, argv[0]);

	while ((ch = getopt(argc, argv, "p:d:n:")) != EOF)	{

		switch (ch)	{
			case 'p':
				if (sscanf(optarg, "%d", &pid) != 1)	{
					fprintf(stderr,"%s: Bad pid.\n", progname);
					fprintf(stderr, usage, progname);
					exit(3);
		 		}
				args++;
#ifdef DEBUG
				fprintf(stderr, "pid = %d.\n", pid);
#endif

			break;

			case 'd':

				if (!strcmp(optarg, "TS"))	{
					policy = SCHED_OTHER;
		 		}
		 		else
		 		if (!strcmp(optarg, "RR"))	{
		 			policy = SCHED_RR;
		 		}
				else
				if (!strcmp(optarg, "ID"))
					policy = SCHED_IDLE;
		 		else
		 		if (!strcmp(optarg, "FF") || !strcmp(optarg, "FIFO"))	{
		 			policy = SCHED_FIFO;
		 		}
		 		else	{
					fprintf(stderr,"%s: Bad scheduling domain '%s'.\n", progname, optarg);
					fprintf(stderr, usage, progname);
					exit(4);
		 		}
				args++;
				do_policy++;
#ifdef DEBUG
				fprintf(stderr,"policy = %d.\n", policy);
#endif

			break;

			case 'n':
					/* we have one digit, let's get more */
				if (sscanf(optarg, "%d", &priority) != 1)	{
					fprintf(stderr,"%s: Bad priority value.\n", progname);
					fprintf(stderr, usage, progname);
					exit(4);
		 		}
				do_priority++;
				args++;
#ifdef DEBUG
				fprintf(stderr, "priority = %d.\n", priority);
#endif

			break;

			default:
				fprintf(stderr,"%s: Unrecognized option '%c'.\n", progname, ch);
				fprintf(stderr, usage, progname);
				exit(6);

			break;
		}

	}

	argc -= optind;
	argv += optind;

	if (args == 0)	{
		fprintf(stderr, usage, progname);
		exit(7);
	}


	if ((maxpri = sched_get_priority_max(SCHED_FIFO)) == -1)	{
		perror("sched_get_priority_max fails");
		exit(8);
	}

	if ((minpri = sched_get_priority_min(SCHED_FIFO)) == -1)	{
		perror("sched_get_priority_min fails");
		exit(10);
	}

	if (do_priority)	{
		if (priority > maxpri)	{
			fprintf(stderr,"%s: maximum priority allowed is %d.\n", progname, maxpri);
			exit(9);
 		}

		if (priority < minpri)	{
			fprintf(stderr,"%s: minimum priority allowed is %d.\n", progname, minpri);
			exit(11);
 		}
	}

	if (!do_policy)	{
		if ((policy = sched_getscheduler(pid)) == -1)	{
			policy = SCHED_OTHER;
		}
	}

	if (!do_priority)	{
		if (sched_getparam(pid, &param) == 0)	{ /* will fail for TS! */
			priority = param.sched_priority;
 		}
 		else	{
 			priority = minpri;
 		}
 		if (!do_policy && pid)	{  /* print and exit */
 			static char *pol;  /* default */
 			if (policy == SCHED_OTHER) pol = "TS";
 			if (policy == SCHED_FIFO) pol = "FIFO";
 			if (policy == SCHED_RR) pol = "RR";

 			fprintf(stdout,"%s: scheduler %s priority %d.\n", progname, pol, priority);

			exit(0);
 		}
	}
	if (policy == SCHED_OTHER || policy == SCHED_IDLE)
		priority = 0;

	param.sched_priority = priority;

	if (sched_setscheduler(pid, policy, &param) == -1)	{
		perror("sched_setscheduler fails");
		exit(13);
	}

	if (pid != 0)	{  /* non-local target => exit */
		exit(0);
	}

	if (!argv[0])	{
		fprintf(stderr,"%s: no valid command args.\n", progname);
		fprintf(stderr, usage, progname);

		exit(14);  /* don't exec if target was not us or no args  */
	}

		/* the big leap! */
		/* hal: dave, will I dream? */

	if (!execvp(argv[0], &argv[0]))	{
		perror("rtrun exec fails");
		exit(-1);
	}
}