#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys

try:
  sys.setdefaultencoding('utf-8')
except AttributeError:
  pass

PY3 = sys.version_info[0] != 2

import os
import re
try:
  import dbus
except:
  sys.stderr.write("warning: dbus not available\n")
import subprocess
from optparse import OptionParser
from signal import signal, SIGINT, SIG_DFL
from functools import partial


signal(SIGINT, SIG_DFL)

VERSION = '0.2'
DEFAULT_TREE = "cpu"
CGROUP_PATHS = {}
CGROUP_SUBSYS = []

CONFIG_SINGLE_TASK = "single_task"


def prnt(arg):
  sys.stdout.write(PY3 and str(arg) or unicode(arg).encode("utf-8"))
  sys.stdout.write("\n")

if PY3:
  u = str
else:
  def tounicode(x):
    return unicode(x, "unicode_escape")
  u = tounicode

fp = open("/proc/cgroups", "r")
for line in fp:
    if line[0] == "#":
        continue
    chunks = line.split()
    CGROUP_SUBSYS.append(chunks[0])

fp = open("/proc/mounts", "r")
for line in fp:
    chunks = line.split()
    if chunks[2] == "cgroup":
        opts = chunks[3].split(",")
        for opt in opts:
            if opt in CGROUP_SUBSYS:
                CGROUP_PATHS[opt] = chunks[1]



class drug():
    def __init__(self, **kwargs):
        for k in kwargs:
            setattr(self, k, kwargs[k])


style = drug(**dict(
        default    = "\033[m",
        # styles
        bold       = "\033[1m",
        underline  = "\033[4m",
        blink      = "\033[5m",
        reverse    = "\033[7m",
        concealed  = "\033[8m",
        # font colors
        black      = "\033[30m",
        red        = "\033[31m",
        green      = "\033[32m",
        yellow     = "\033[33m",
        blue       = "\033[34m",
        magenta    = "\033[35m",
        cyan       = "\033[36m",
        white      = "\033[37m",
        # background colors
        on_black   = "\033[40m",
        on_red     = "\033[41m",
        on_green   = "\033[42m",
        on_yellow  = "\033[43m",
        on_blue    = "\033[44m",
        on_magenta = "\033[45m",
        on_cyan    = "\033[46m",
        on_white   = "\033[47m"))


def terminal_size():
    def ioctl_GWINSZ(fd):
        try:
            import fcntl, termios, struct, os
            cr = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ,
        '1234'))
        except:
            return None
        return cr
    cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2)
    if not cr:
        try:
            fd = os.open(os.ctermid(), os.O_RDONLY)
            cr = ioctl_GWINSZ(fd)
            os.close(fd)
        except:
            pass
    if not cr:
        try:
            cr = (os.environment['LINES'], os.environment['COLUMNS'])
        except:
            cr = (25, 80)
    return int(cr[1]), int(cr[0])


TERM_WIDTH = terminal_size()[0]

def cap_line(line, width, utf):
    if len(line) < width:
        return line
    else:
        if utf:
            return line[:width-1] + u("\u2026")
        else:
            return line[:width-3] + u("...")

def list_processes():
    rv = []
    for line in os.listdir("/proc"):
        if re.match("\d+", line):
            rv.append(line)
    return rv

def get_process_data(pid):
    try:
        path = os.path.join("/proc", pid, "status")
        if not os.path.exists(path):
            return None
        fp = open(path)
    except OSError:
        return None
    rv = {}
    for line in fp.readlines():
        chunks = line.split(":\t")
        rv[chunks[0].strip().lower()] = chunks[1].strip()
    fp.close()
    try:
        path = os.path.join("/proc", pid, "cmdline")
        if not os.path.exists(path):
            return None
        fp = open(path)
        rv["cmdline"] = fp.read().split("\0")[:-1]
        fp.close()
    except OSError:
        return None
    return rv

def flags_to_string(flags):
    out = ""
    for flag in flags:
        if isinstance(flag, (tuple, list)):
            flag = flag[1]
        out += "[%s" %flag["name"]
        for k,v in flag.iteritems():
            if k == "name":
                continue
            if isinstance(v, (int, long)):
                if v != 0:
                    out += " %s=%s" %(unicode(k), unicode(v))
            elif isinstance(v, (str, unicode)):
                if v != "":
                    out += " %s=%s" %(unicode(k), unicode(v))
            else:
                out += " %s=%s" %(unicode(k), unicode(v))
        out += "]"
    return out

def print_process(pid, flags=None):
    data = get_process_data(pid)
    out = "%6s %s" %(data["pid"], data["name"].ljust(25))
    if flags:
        out += flags_to_string(flags)
    prnt(out)

class Session:

    def __init__(self, dbus_loop=None, silent=True):
        self.dbus_loop = dbus_loop
        self.bus = dbus.SystemBus(mainloop=self.dbus_loop)
        self.silent = silent
        self.sets = {
            'config':self.set_config,
            'active':self.set_active,
            }
        self.gets = {
            'config':self.get_config,
            'desc':self.get_desc,
            'systemflags':self.get_systemflags,
            'activelistlength':self.get_activelistlength,
            }

    def init_system(self):
        if hasattr(self, 'system'): return True
        try:
            self.system = self.bus.get_object(
                'org.quamquam.ulatencyd',
                '/org/quamquam/ulatencyd/System',
                follow_name_owner_changes=bool(self.dbus_loop))
        except dbus.exceptions.DBusException as e:
            if self.silent:
                return False
            raise
        return True

    def test_connection(self):
        try:
            self.system = self.bus.get_object(
                'org.quamquam.ulatencyd',
                '/org/quamquam/ulatencyd/System')
            return True
        except dbus.exceptions.DBusException as e:
            print(e)


    def init_user(self):
        if hasattr(self, 'user'): return True
        try:
            self.user = self.bus.get_object(
                'org.quamquam.ulatencyd',
                '/org/quamquam/ulatencyd/User',
                follow_name_owner_changes=bool(self.dbus_loop))
        except dbus.exceptions.DBusException as e:
            if self.silent:
                return False
            raise
        return True

    def config_list(self):
        if not self.init_system():
          return []

        system = dbus.Interface(self.system, 'org.quamquam.ulatencyd.System')
        try: return list(map(unicode, system.listSchedulerConfigs()))
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return None
            raise

    def set(self, key, *value):
        if key not in self.sets:
            prnt("error: key '{0}' not available".format(key))
            sys.exit(1)
        return self.sets[key](*value)

    def set_config(self, config=None, timeout=600):
        if not self.init_system():
            return False
        configs = self.config_list()
        if config is None: return ",".join(configs)
        system = dbus.Interface(self.system, 'org.quamquam.ulatencyd.System')
        try: return system.setSchedulerConfig(config, timeout=timeout)
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return None
            raise

    def set_active(self, pid=None):
        if pid is None:
            prnt("error: no pid given")
            sys.exit(1)
        if not self.init_user():
            return False
        user = dbus.Interface(self.user, 'org.quamquam.ulatencyd.User')
        try: return user.setActive(pid)
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return None
            raise

    def get(self, key, *value):
        if key not in self.gets:
            prnt("error: key '{0}' not available".format(key))
            sys.exit(1)
        return self.gets[key](*value)

    def get_config(self):
        if not self.init_system():
          return ""
        properties=dbus.Interface(self.system,'org.freedesktop.DBus.Properties')
        try:
            return unicode(properties.Get(
                'org.quamquam.ulatencyd.System', 'config'))
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return None
            raise

    def get_desc(self, config=None):
        if not self.init_system():
            return None

        configs = self.config_list()
        if config is None: return ",".join(configs)
        if config not in configs:
            prnt("error: config '{0}' not available".format(config))
            sys.exit(1)
        system = dbus.Interface(self.system, 'org.quamquam.ulatencyd.System')
        try: return system.getSchedulerConfigDescription(config)
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return None
            raise

    def get_systemflags(self):
        if not self.init_system():
            return []
        system = dbus.Interface(self.system, 'org.quamquam.ulatencyd.System')
        try:
            return list(map(unicode, system.listSystemFlags()))
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return None
            raise

    def get_processflags(self, pid=None, recrusive=True):
        if not self.init_system():
            return []
        if pid is None:
            prnt("error: no pid given")
            sys.exit(1)
        pid = int(pid)
        system = dbus.Interface(self.system, 'org.quamquam.ulatencyd.System')
        try:
            rv = []
            data = system.listFlags(pid, recrusive)
            for flag in data:
              rv.append([flag[0],dict((unicode(x),v) for x,v in flag[1].iteritems())])
            return rv
        except dbus.exceptions.DBusException as e:
            if self.silent:
                return []
            raise


    def get_activelistlength(self):
        if not self.init_user():
            return None

        properties = dbus.Interface(self.user,'org.freedesktop.DBus.Properties')
        try:
            return unicode(properties.Get(
                'org.quamquam.ulatencyd.User', 'activeListLength'))
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return []
            raise

    def add_flag(self, pid, **values):
        if not self.init_system():
          return ""
        system = dbus.Interface(self.system, 'org.quamquam.ulatencyd.System')
        try:
            rv = []
            data = system.addFlag(pid, 0, values["name"],
                                          unicode(values.get("reason", "")),
                                          long(values.get("timeout", 0)),
                                          int(values.get("priority", 0)),
                                          long(values.get("value", 0)),
                                          long(values.get("treshold", 0)),
                                          bool(values.get("inherit", False)))
        except dbus.exceptions.DBusException as e:
            prnt(e)
            if self.silent:
                return False
            raise e


class Tree(object):
    """Displays a cgroup tree"""

    boxes = {
      True: drug(
        start     = u("\u2500\u252c\u00bb"),
        branch    = u("\u251c"),
        last      = u("\u2514"),
        end       = u("\u00ab"),
        subtree   = u("\u2502"),
        file      = u(" ") ),
      False: drug(
        start     = u("# "),
        branch    = u("+-"),
        last      = u("+-"),
        end       = u(" #"),
        subtree   = u("|"),
        file      = u(" ") ),
    }

    def __init__(self, tree=DEFAULT_TREE, processes=True, show_all=True,
                 utf = True, color = False, flags = False, cmdline = False):
        self.tree = tree
        self.processes = processes
        self.show_all = show_all
        self.utf = utf
        self.flags = flags
        self.cmdline = cmdline
        if flags:
            self.session = Session()
        if color:
            self.color = drug(
                process = u("{0.yellow}{1}{0.default} {2}"),
                subtree = u("{0.bold}{0.black}{1}{0.default}"),
                tree    = u("{0.bold}{0.black}{1}{2}{3}{0.default}{0.green}") +
                          u("{5}{0.bold}{0.black}{4}{0.default}"),
                file    = u("{0.bold}{0.black}{1}{2}{3}{0.default}{4}{5}{0.magenta}{6}{0.default}") )
        else:
            self.color = drug(
                process = u("{1} {2}"),
                subtree = u("{1}"),
                tree    = u("{1}{2}{3}{5}{4}"),
                file    = u("{1}{2}{3}{4}{5}{6}") )
        path = self.get_path()
        if path:
            self.cache = drug(children = [], path = path, state = None,
                name = os.path.basename(os.path.abspath(path)))
        else:
            prnt("error: tree does not exist")
            return None

    def processes_(self, b):
        self.processes = b

    def show_all_(self, b):
        self.show_all = b

    def get_process(self, pid):
        data = get_process_data(pid)
        if not data:
            return None
        if not self.show_all and data and data["pid"] != data["tgid"]:
            return None
        if self.cmdline:
            cmdline = u(" ").join(data.get("cmdline", []))
            return data["pid"], cmdline or data.get("name", "")
        return data["pid"], data.get("name", "")

    def generate(self, dir):
        tree = drug(children = [], path = dir, state = "",
            name = os.path.basename(os.path.abspath(dir)))
        files = []
        if self.processes and os.path.exists(os.path.join(dir, "tasks")):
            try:
                fp = open(os.path.join(dir, "tasks"), "r")
                for line in fp.readlines():
                    proc = self.get_process(line.strip())
                    if proc:
                        files.append(proc)
                fp.close()
            except IOError:
                pass
        for x in os.listdir(dir):
            if os.path.isdir(dir + os.sep + x):
                files.append((None, x))

        for pid, fname in files:
            path = os.path.join(dir, fname)
            if os.path.isdir(path):
                tree.children.append(self.generate(path))
            else:
                tree.children.append(drug(
                    state = "", pid = pid, name = fname, path = path))
        return tree

    def get_path(self):
        path = CGROUP_PATHS.get(self.tree, None)
        if not path or not os.path.exists(path):
            return None
        return path


    def update(self):
        path = self.get_path()
        if not path: return
        tree = self.generate(path)

        def rec(old, new):
            if old.state is None:
                old.state = "new"
            else: old.state = ""
            if new.name != old.name:
                old.name = new.name
                if not old.state:
                    old.state = "update"
            if hasattr(old, 'children'):
                children = dict([(child.name, child) for child in new.children])
                doomed = []
                for n, child in enumerate(old.children):
                    if child.state == "delete":
                        doomed.append(n)
                        continue
                    if child.name in children:
                        rec(child, children[child.name])
                        del children[child.name]
                    else:
                        child.state = "delete"
                for n in sorted(doomed, reverse = True):
                    del old.children[n]
                for child in children.values():
                    old.children.append(child)
                    child.state = "new"

        rec(self.cache, tree)
        return self.cache

    def display(self):
        path = self.get_path()
        if not path:
          return None
        prnt(path)
        if not os.path.exists(path):
            prnt("error: tree does not exist")
            sys.exit(1)

        def rec(branch, padding, islast=None):
            if hasattr(branch, 'pid'):
                if self.flags:
                    fpad = " "*(max(28-len(branch.name)-len(str(branch.pid)), 3))
                    flags = flags_to_string(self.session.get_processflags(branch.pid))
                else:
                    flags = ""
                    fpad = ""
                process = self.color.process.format(
                    style, branch.pid, cap_line(branch.name, TERM_WIDTH - len(padding) - 13 - len(flags), self.utf))
                prnt(self.color.file.format(style, padding,
                    islast or self.boxes[self.utf].branch,
                    self.boxes[self.utf].file, process, fpad, flags))
                return
            prnt(self.color.tree.format(style, padding,
                islast or self.boxes[self.utf].branch,
                self.boxes[self.utf].start,
                self.boxes[self.utf].end,
                branch.name))
            count = 0
            for child in branch.children:
                count += 1
                last = None
                p = padding
                if count == len(branch.children):
                    last = self.boxes[self.utf].last
                    if os.path.isdir(branch.path):
                        if islast:
                            p = p + u(" ")
                        else:
                            p = p + self.boxes[self.utf].subtree
                    p = p + u(" ")
                else:
                    if islast:
                        p = p + u(" ")
                    else:
                        p = p + self.boxes[self.utf].subtree
                    if os.path.isdir(branch.path):
                        p = p + u(" ")
                rec(child, p, last)

        rec(self.generate(path), u(""), self.boxes[self.utf].last)

    def clear(self):
        """clears a cgroup mountpoint so it can be removed"""

        path = self.get_path()
        if not path:
            prnt("subsystem not mounted")
            return

        try:
            rootfp = open(os.path.join(path, "tasks"), "w")
        except IOError:
            prnt("need to run as root")
            sys.exit(1)

        def move(pid):
            rootfp.write("%s\n" %pid)
            try:
                rootfp.flush()
            except IOError:
                pass

        def clear(root):
            prnt("clear %s" %root)
            try:
                fp = open(os.path.join(root, "tasks"), "r")
                for line in fp:
                    move(line.strip())
            except IOError:
                pass

        for root, dirs, files in os.walk(path, topdown=False):
            if root == path:
                continue
            for retry in xrange(5):
                clear(root)
                try:
                    os.rmdir(root)
                    break
                except OSError:
                    pass
            else:
                prnt("could not clear %s" %root)
        prnt("cleared %s" %path)

class Gui:

    def __init__(self):
        pass

    def run(self):
        return 0


class QtGui(Gui):

    def __init__(self):
        try:
            from PyQt4.Qt import Qt, QApplication, QSystemTrayIcon, \
                                 QIcon, QPixmap, QMenu, QMainWindow, \
                                 QTreeWidget, QTreeWidgetItem, QTimer,\
                                 QColor, QMenuBar, QKeySequence, QSettings, \
                                 QActionGroup
        except ImportError:
            prnt("python qt4 bindings missing.")
            sys.exit(1)

        if dbus:
            # if DBusQtMainLoop is not available, we simply stop to work
            # correctly when ulatencyd is restarted
            try:
                from dbus.mainloop.qt import DBusQtMainLoop
                self.dbus_loop = DBusQtMainLoop()
            except ImportError as e:
                self.dbus_loop = None

        class MainWindow(QMainWindow):
            def __init__(self, app):
                self.app = app
                super(MainWindow, self).__init__()

            def closeEvent(self, event):
                self.app.toggle_window(event)
                event.ignore()


        self.QColor = QColor
        self.QTreeWidgetItem = QTreeWidgetItem
        self.app = app = QApplication(sys.argv)
        self.session = Session(dbus_loop=self.dbus_loop)
        self.settings = QSettings("ulatencyd", "qtgui")
        self.first_fill = True

        self.tree = Tree(show_all = False)
        icon = QIcon(QPixmap(self.get_logo()))
        app.trayicon = QSystemTrayIcon(icon, app)
        app.trayicon.show()
        self.window = win = MainWindow(self)

        mb = QMenuBar(win)
        win.setMenuBar(mb)
        menu = mb.addMenu("client")
        c = menu.addMenu("set config")
        c.aboutToShow.connect(partial(self.update_active, c))
        menu.addSeparator()

        quit = menu.addAction("quit")
        quit.setShortcut(QKeySequence("Ctrl+Q"))
        quit.triggered.connect(app.quit)
        menu = mb.addMenu("settings")

        tree_menu = menu.addMenu("tree")

        self.config_action = {}
        self.config_tree = {}

        def connect(con):
            return lambda: self.switch_tree(con)
        tree_group = QActionGroup(tree_menu)
        for subsys in CGROUP_PATHS:
            prnt(subsys)
            self.config_tree[subsys] = c = tree_menu.addAction(subsys)
            c.setCheckable(True)
            c.setActionGroup(tree_group)
            c.setChecked(self.tree.tree == subsys)
            c.triggered.connect(connect(subsys))

        processes = menu.addAction("show processes")
        processes.setCheckable(True)
        processes.setChecked(self.tree.processes)
        processes.triggered.connect(lambda:
            self.tree.processes_(processes.isChecked()) or self.update_tree())
        show_all = menu.addAction("show threads")
        show_all.setCheckable(True)
        show_all.setChecked(self.tree.show_all)
        show_all.triggered.connect(lambda:
            self.tree.show_all_(show_all.isChecked()) or self.update_tree())

        win.ui = ui = drug(
            treeview = QTreeWidget(win),
            timer = QTimer(),
        )
        win.setCentralWidget(ui.treeview)
        ui.treeview.setHeaderLabels(["process", "pid"])
        ui.treeview.setAlternatingRowColors(True)
        ui.treeview.indexOfChild = ui.treeview.indexOfTopLevelItem
        ui.treeview.child = ui.treeview.topLevelItem
        ui.timer.setInterval(5000)
        ui.timer.timeout.connect(self.update_tree)


        self.cur = self.session.get_config()

        self.app.menu = QMenu()
        self.config_menu = menu = self.app.menu.addMenu("set config")
        menu.aboutToShow.connect(partial(self.update_active, menu))
        menu = self.app.menu
        menu.addSeparator()
        quit = menu.addAction("Quit")

        self.doomed_entries = []
        self.update_tree()
        app.trayicon.activated.connect(self.toggle_window)
        app.trayicon.setContextMenu(menu)
        quit.triggered.connect(self.quit)
        prnt("gui started.")

    def update_active(self, menu):
        self.cur = self.session.get_config()
        configs = self.session.config_list()

        self.config_action = {}
        menu.clear()

        if not configs:
            return

        def connect(con):
            return lambda: self.switch_config(con)
        for config in configs:
            self.config_action[config] = c = menu.addAction(config)
            c.setCheckable(True)
            c.setChecked(self.cur == config)
            c.setToolTip(unicode(self.session.get_desc(config)))
            c.triggered.connect(connect(config))

    def switch_tree(self, tree):
        self.tree = Tree(tree=tree, show_all = self.tree.show_all)
        self.doomed_entries = []
        self.window.ui.treeview.clear()
        self.update_tree()

    def quit(self):
        if self.window.isVisible():
            self.settings.setValue("main_window_geometry", self.window.saveGeometry());
            self.settings.setValue("main_window_state", self.window.saveState());
        self.settings.sync()
        self.app.quit()

    def run(self):
        return self.app.exec_()

    def toggle_window(self, reason):
        if reason != self.app.trayicon.Context:
            b = not self.window.isVisible()
            if b:
                self.update_tree()
                self.window.restoreGeometry(self.settings.value("main_window_geometry").toByteArray());
                self.window.restoreState(self.settings.value("main_window_state").toByteArray());
                self.window.ui.timer.start()
            else:
                self.settings.setValue("main_window_geometry", self.window.saveGeometry());
                self.settings.setValue("main_window_state", self.window.saveState());
                self.window.ui.timer.stop()
            self.window.setVisible(b)

    def update_tree(self):
        QTreeWidgetItem = self.QTreeWidgetItem
        tree = self.tree.update()

        if not tree:
            return

        for parent, child in self.doomed_entries:
            i = parent.indexOfChild(child)
            parent.takeChild(i)
        self.doomed_entries = []

        def add(branch, parent):
            background = self.QColor(0, 200, 0, 100)
            pid = u("")
            if hasattr(branch, 'pid'):
                pid = branch.pid
            branch.item = item = QTreeWidgetItem(parent, [branch.name, pid])
            item.setBackgroundColor(0, background)
            item.setBackgroundColor(1, background)
            item.setExpanded(True)
            if hasattr(branch, 'children'):
                for child in branch.children:
                    if hasattr(child, 'children'):
                        add(child, item)
                for child in branch.children:
                    if not hasattr(child, 'children'):
                        add(child, item)

        def update(branch, parent):
            background = self.QColor("transparent")
            dead = self.QColor(200, 0, 0, 100)
            if branch.state == "update":
                i = parent.indexOfChild(branch.item)
                if i == -1:
                    branch.state = "delete" # FIXME
                    prnt("FIXME")
                else:
                    parent.child(i).setText(0, branch.name)
            branch.item.setBackgroundColor(0, background)
            branch.item.setBackgroundColor(1, background)
            if hasattr(branch, 'children'):
                for child in branch.children:
                    if child.state == "delete":
                        self.doomed_entries.append((branch.item, child.item))
                        child.item.setBackgroundColor(0, dead)
                        child.item.setBackgroundColor(1, dead)
                    elif child.state == "new":
                        add(child, branch.item)
                    else:
                        update(child, branch.item)

        if tree.state == "new":
            add(tree, self.window.ui.treeview)
        else:
            update(tree, self.window.ui.treeview)

        if self.first_fill:
            self.window.ui.treeview.resizeColumnToContents(0)
            self.first_fill = False

    def switch_config(self, new):
        if self.session.set_config(new, timeout=60*60*24):
            self.config_action[self.cur].setChecked(False)
            self.config_action[new].setChecked(True)

    def get_logo(self):
        header = ["16 16 5 1"]
        chars = ["#","x","o","."]
        color = ["black",
            "gray25",
            "gray50",
            "white"]
        background = [
            "                ",
            "      xxxx      ",
            "    xx####xx    ",
            "   x##xxxx##x   ",
            "  x#xx    xx#x  ",
            "  x#x      x#x  ",
            " x#x        x#x ",
            " x#x        x#x ",
            " x#x        x#x ",
            " x#x        x#x ",
            "  x#x      x#x  ",
            "  x#xx    xx#x  ",
            "   x##xxxx##x   ",
            "    xx####xx    ",
            "      xxxx      ",
            "                "]
        foreground = [
            "o    o",
            "o    o",
            "o.  .o",
            "o....o",
            "oooo.o",
            "o.o   ",
            "ooo   "]
        color = ["  c None"] + [char+" c "+col for char,col in zip(chars,color)]
        for y, line in enumerate(foreground):
            for x, char in enumerate(line):
                if char != " ":
                    s = background[y+4]
                    background[y+4] = s[:x+5] + char + s[x+6:]
        return header + color + background



def command(args, config=CONFIG_SINGLE_TASK, as_root=False):
    pid = os.getpid()
    uid = None
    if os.getuid() == 0:
        uid = (not as_root and int(os.environ.get("SUDO_UID", 0))) or os.getuid()
        #prnt("you need to run: sudo comand")
        #sys.exit(1)
    session = Session()
    old_config = session.get_config()
    if not old_config:
        prnt("ulatency: could not contact dbus daemon")
        session.test_connection()
        sys.exit(1)
    prnt("ulatency: save old config: %s" %old_config)
    session.add_flag(pid, name="cmd.config.%s" %config, inherit=True)
    if not session.set_config(config, timeout=60*60*24):
        prnt("ulatency: abort by user")
        sys.exit(2)

    prnt("ulatency: execute '%s'" %" ".join(args))

    def int_handler(ignum, frame):
        prnt("-" * 30)
        prnt("received interrupt signal")
        prnt("ulatency: restore config:%s" %old_config)
        session.set_config(old_config, timeout=60*60*24)
        sys.exit(1)

    import signal
    signal.signal(signal.SIGINT, int_handler)

    if uid is not None:
        prnt("Execute as id: %s" %uid)
        old_uid = os.getuid()
        os.seteuid(uid)
    try:
        prnt("-"*30)
        retcode = subprocess.call(args)
        prnt("-"*30)
        if retcode < 0:
            print >>sys.stderr, "Child was terminated by signal", -retcode
        else:
            print >>sys.stderr, "Child returned", retcode
    except OSError, e:
        print >>sys.stderr, "Execution failed:", e
    except KeyboardInterrupt, e:
        print "kbi"

    if uid is not None:
        os.seteuid(old_uid)
    prnt("ulatency: restore config %s" %old_config)
    session.set_config(old_config, timeout=60*60*24)



EPILOG="""commands:
  tree [cgroup]              show a cgroup tree
  get [key]                  get daemon setting
  set [key value]            set daemon value
  flags [process]            list flags of a process or all
  sysflags                   list system flags
  mounts                     lists all mounted cgroups
  clear cgroup|ALL           clears a cgroup
  run [options]  -- [COMMAND]  run task under different scheduler config
  run-CONFIG [COMMAND]       run task under scheduler CONFIG
"""

def main():
    parser = OptionParser( usage = "usage: %prog [options] [command] [key [value]]", version="%prog " + VERSION, epilog="commands for a list of commands" )
    parser.add_option("--config-list", dest="list", action="store_true",
        default=False, help="list system configs")
    parser.add_option("-s", "--set", dest="set", action="store_true",
        default=False, help="set key value pair")
    parser.add_option("-g", "--get", dest="get", action="store_true",
        default=False, help="get value by key")
    parser.add_option("--gui", dest="gui", action="store_true",
        default=False, help="enable gui")
    parser.add_option("--no-utf", dest="utf", action="store_false",
        default=True, help="disables utf8 for tree view")
    parser.add_option("--no-color", dest="color", action="store_false",
        default=True, help="disables color")
    parser.add_option("--all", dest="all", action="store_true",
        default=False, help="show complete tree")
    parser.add_option("--flags", dest="flags", action="store_true",
        default=False, help="show flags on tree output")
    parser.add_option("--cmd", dest="cmdline", action="store_true",
        default=False, help="show cmdline")
    parser.add_option("--as-root", dest="as_root", action="store_true",
        default=False, help="run command as root")
    parser.add_option("--no-processes", dest="processes", action="store_false",
        default=True, help="hide all processes")
    parser.add_option("--scheduler", dest="scheduler", action="store",
        default=CONFIG_SINGLE_TASK, help="scheduler config to use")

    prog = os.path.basename(sys.argv[0])

    if prog[:4] in ("run-", "run_"):
        if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help"):
            prnt("Usage: %s COMMAND" %prog)
            prnt("executes program under scheduler config: %s" %prog[4:].replace("-", "_").upper())
            sys.exit(1)
        command(sys.argv[1:], config=prog[4:].replace("-", "_"))
        return

    options, args = parser.parse_args()

    # if we run without any arguments and it doesn't look like we are in a
    # terminal, we start the gui
    if prog == "ulatency-gui":
        options.gui = True

    if options.gui:
        sys.exit(QtGui().run())

    if options.list:
        prnt("\n".join(Session().config_list()))
        sys.exit(0)

    if options.set or args and args[0] == 'set':
        if len(args[1:]):
            rv = Session().set(*args[1:])
            prnt(rv)
            if rv: sys.exit(0)
            sys.exit(2)

        else: prnt(",".join(Session().sets.keys()))
        return

    if options.get or args and args[0] == 'get':
        if len(args[1:]):
            rv = Session().get(*args[1:])
            prnt(rv)
            if rv: sys.exit(0)
            sys.exit(2)
        else: prnt(",".join(Session().gets.keys()))
        return


    if args and args[0] == 'run':
        command(args[1:], options.scheduler, options.as_root)
        return

    if args and args[0][:4] == "run-":
        command(args[1:], config=args[0][4:].replace("-", "_"))
        return


    if not args or args[0] == 'tree':
        kwargs = dict(utf = options.utf, color = options.color,
            show_all = options.all, processes = options.processes,
            flags = options.flags, cmdline = options.cmdline)
        if len(args) > 1:
            tree = Tree(tree=args[1], **kwargs)
        else:
            tree = Tree(**kwargs)
        tree.display()
        return

    if args[0] == 'commands':
        parser.print_usage()
        prnt(EPILOG)
        return

    if args[0] == 'mounts':
        for v,i in CGROUP_PATHS.iteritems():
            prnt("%s %s" %(v.ljust(10, " "), i))
        return

    if args[0] == 'clear':
        if not len(args) == 2:
            prnt("not enough arguments")
            sys.exit(1)
        if args[1] == "ALL":
            for tree in CGROUP_PATHS.iterkeys():
                ct = Tree(tree=tree)
                ct.clear()
        else:
            ct = Tree(tree=args[1])
            ct.clear()
        return

    if args[0] == 'flags':
        session = Session()
        if len(args) > 1:
            print_process(args[1], flags=session.get_processflags(args[1]))
        else:
            for proc in list_processes():
                print_process(proc, flags=session.get_processflags(proc))
        return

    if args[0] == 'sysflags':
        session = Session()
        prnt(flags_to_string(session.get_systemflags()))
        return

    parser.print_usage()
    prnt(EPILOG)


if __name__ == "__main__":
    main()
