#! /usr/bin/python
# -*- python -*-
# -*- coding: utf-8 -*-
#   Copyright (C) 2008 Red Hat Inc.
#
#   Arnaldo Carvalho de Melo <acme@redhat.com>
#
#   This application 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; version 2.
#
#   This application 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.

from __future__ import print_function
import os
import schedutils
import sys


def usage():
    print('''pchrt (python-schedutils)
usage: pchrt [options] [prio] [pid | cmd [args...]]
manipulate real-time attributes of a process
  -b, --batch                        set policy to SCHED_BATCH
  -f, --fifo                         set policy to SCHED_FIFO
  -i, --idle                         set policy to SCHED_IDLE
  -p, --pid                          operate on existing given pid
  -m, --max                          show min and max valid priorities
  -o, --other                        set policy to SCHED_OTHER
  -r, --rr                           set policy to SCHED_RR (default)
  -R, --reset-on-fork                set SCHED_RESET_ON_FORK for FIFO or RR
  -h, --help                         display this help

You must give a priority if changing policy.

Report bugs and send patches to <tuna-devel@lists.fedorahosted.org>''')
    return


def show_priority_limits(policy):
    print("%-32.32s: %d/%d" % (
        "%s min/max priority" % schedutils.schedstr(policy),
        schedutils.get_priority_min(policy),
        schedutils.get_priority_max(policy)
    ))


def show_all_priority_limits():
    for policy in (schedutils.SCHED_OTHER, schedutils.SCHED_FIFO,
                   schedutils.SCHED_RR, schedutils.SCHED_BATCH):
        show_priority_limits(policy)


def show_settings(pid):
    policy = schedutils.get_scheduler(pid)
    spolicy = schedutils.schedstr(policy)
    rtprio = schedutils.get_priority(pid)
    reset_on_fork = ""
    if policy & schedutils.SCHED_RESET_ON_FORK:
        reset_on_fork = "|SCHED_RESET_ON_FORK"
    print('''pid %d's current scheduling policy: %s%s
pid %d's current scheduling priority: %d''' % (pid, spolicy, reset_on_fork,
                                               pid, rtprio))


def valid_policy_flag(policy, policy_flag):
    if policy_flag == schedutils.SCHED_RESET_ON_FORK and \
       policy not in (schedutils.SCHED_RR, schedutils.SCHED_FIFO):
        print("SCHED_RESET_ON_FORK flag is supported for SCHED_FIFO and SCHED_RR policies only")
        return False
    return True


def change_settings(pid, policy, policy_flag, rtprio):
    try:
        schedutils.set_scheduler(pid, policy | policy_flag, rtprio)
    except SystemError as err:
        print("sched_setscheduler: %s" % err[1])
        print("failed to set pid %d's policy" % pid)


def main():

    args = sys.argv[1:]
    if not args:
        usage()
        return

    policy = schedutils.SCHED_RR
    policy_flag = 0
    while True:
        o = args.pop(0)
        try:
            priority = int(o)
            break
        except:
            pass

        if o in ("-h", "--help"):
            usage()
            return
        elif o in ("-b", "--batch"):
            policy = schedutils.SCHED_BATCH
        elif o in ("-f", "--fifo"):
            policy = schedutils.SCHED_FIFO
        elif o in ("-i", "--idle"):
            policy = schedutils.SCHED_IDLE
        elif o in ("-m", "--max"):
            show_all_priority_limits()
            return
        elif o in ("-o", "--other"):
            policy = schedutils.SCHED_OTHER
        elif o in ("-r", "--rr"):
            policy = schedutils.SCHED_RR
        elif o in ("-R", "--reset-on-fork"):
            policy_flag |= schedutils.SCHED_RESET_ON_FORK
        elif o in ("-p", "--pid"):
            if len(args) > 1:
                priority = int(args.pop(0))
                pid = int(args.pop(0))
                if not valid_policy_flag(policy, policy_flag):
                    return
                change_settings(pid, policy, policy_flag, priority)
            else:
                pid = int(args.pop(0))
                show_settings(pid)
            return
        else:
            usage()
            return

    if not valid_policy_flag(policy, policy_flag):
        return

    schedutils.set_scheduler(0, policy | policy_flag, priority)
    os.execvp(args[0], args)


if __name__ == '__main__':
    main()