#! /usr/bin/python3
# -*- python -*-
# -*- coding: utf-8 -*-
#   tuna - Application Tuning GUI
#   Copyright (C) 2008, 2009, 2010, 2011 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.

""" tuna - Application Tuning GUI """

import os
import sys
import errno
import re
import getopt
import fnmatch
import gettext
import locale
from functools import reduce
import ethtool
import tuna.tuna_sched as tuna_sched
import procfs
from tuna import tuna, sysfs

try:
    import inet_diag
    have_inet_diag = True
except:
    have_inet_diag = False

# FIXME: ETOOMANYGLOBALS, we need a class!

nr_cpus = None
ps = None
irqs = None
version = "0.15"


def usage():
    print(_('Usage: tuna [OPTIONS]'))
    fmt = '\t%-40s %s'
    print(fmt % ('-h, --help',                  _('Give this help list')))
    print(fmt % ('-a, --config_file_apply=profilename',
                 _('Apply changes described in profile')))
    print(fmt % ('-l, --config_file_list',
                 _('List preloaded profiles')))
    print(fmt % ('-g, --gui',                   _('Start the GUI')))
    print(fmt % ('-G, --cgroup',
                 _('Display the processes with the type of cgroups they are in')))
    print(fmt % ('-c, --cpus=' + _('CPU-LIST'), _('%(cpulist)s affected by commands') %
                 {"cpulist": _('CPU-LIST')}))
    print(fmt % ('-C, --affect_children',
                 _('Operation will affect children threads')))
    print(fmt % ('-f, --filter',
                 _('Display filter the selected entities')))
    print(fmt % ('-i, --isolate',               _('Move all threads away from %(cpulist)s') %
                 {"cpulist": _('CPU-LIST')}))
    print(fmt % ('-I, --include',               _('Allow all threads to run on %(cpulist)s') %
                 {"cpulist": _('CPU-LIST')}))
    print(fmt % ('-K, --no_kthreads',
                 _('Operations will not affect kernel threads')))
    print(fmt % ('-m, --move',                  _('Move selected entities to %(cpulist)s') %
                 {"cpulist": _('CPU-LIST')}))
    print(fmt % ('-N, --nohz_full',
                 _('CPUs in nohz_full= kernel command line will be affected by operations')))
    if have_inet_diag:
        print(fmt % ('-n, --show_sockets',
                     _('Show network sockets in use by threads')))
    print(fmt % ('-p, --priority=[' +
                 _('POLICY') + ':]' +
                 _('RTPRIO'),                   _('Set thread scheduler tunables: %(policy)s and %(rtprio)s') %
                 {"policy": _('POLICY'), "rtprio": _('RTPRIO')}))
    print(fmt % ('-P, --show_threads',          _('Show thread list')))
    print(fmt % ('-Q, --show_irqs',             _('Show IRQ list')))
    print(fmt % ('-q, --irqs=' + _('IRQ-LIST'), _('%(irqlist)s affected by commands') %
                 {"irqlist": _('IRQ-LIST')}))
    print(fmt % ('-r, --run=' + _('COMMAND'),   _('fork a new process and run the %(command)s') %
                 {"command": _('COMMAND')}))
    print(fmt % ('-s, --save=' + _('FILENAME'), _('Save kthreads sched tunables to %(filename)s') %
                 {"filename": _('FILENAME')}))
    print(fmt % ('-S, --sockets=' +
                 _('CPU-SOCKET-LIST'),          _('%(cpusocketlist)s affected by commands') %
                 {"cpusocketlist": _('CPU-SOCKET-LIST')}))
    print(fmt % ('-t, --threads=' +
                 _('THREAD-LIST'),              _('%(threadlist)s affected by commands') %
                 {"threadlist": _('THREAD-LIST')}))
    print(fmt % ('-U, --no_uthreads',
                 _('Operations will not affect user threads')))
    print(fmt % ('-v, --version',               _('Show version')))
    print(fmt % ('-W, --what_is',
                 _('Provides help about selected entities')))
    print(fmt % ('-x, --spread',                _('Spread selected entities over %(cpulist)s') %
                 {"cpulist": _('CPU-LIST')}))


def get_nr_cpus():
    global nr_cpus
    if nr_cpus:
        return nr_cpus
    nr_cpus = procfs.cpuinfo().nr_cpus
    return nr_cpus


nics = None


def get_nics():
    global nics
    if nics:
        return nics
    nics = ethtool.get_active_devices()
    return nics


def thread_help(tid):
    global ps
    if not ps:
        ps = procfs.pidstats()

    if tid not in ps:
        print("tuna: " + _("thread %d doesn't exists!") % tid)
        return

    pinfo = ps[tid]
    cmdline = procfs.process_cmdline(pinfo)
    help, title = tuna.kthread_help_plain_text(tid, cmdline)
    print("%s\n\n%s" % (title, _(help)))


def save(cpu_list, thread_list, filename):
    kthreads = tuna.get_kthread_sched_tunings()
    for name in list(kthreads.keys()):
        kt = kthreads[name]
        if (cpu_list and not set(kt.affinity).intersection(set(cpu_list))) or \
           (thread_list and kt.pid not in thread_list):
            del kthreads[name]
    tuna.generate_rtgroups(filename, kthreads, get_nr_cpus())


def ps_show_header(has_ctxt_switch_info, cgroups=False):
    print("%7s %6s %5s %7s       %s" %
          (" ", " ", " ", _("thread"),
           has_ctxt_switch_info and "ctxt_switches" or ""))
    print("%7s %6s %5s %7s%s %15s" % ("pid", "SCHED_", "rtpri", "affinity",
                                      has_ctxt_switch_info and " %9s %12s" % (
                                          "voluntary", "nonvoluntary")
                                      or "", "cmd"), end=' ')
    if cgroups:
        print(" %7s" % ("cgroup"))
    else:
        print("")


def ps_show_sockets(pid, ps, inodes, inode_re, indent=0):
    header_printed = False
    dirname = "/proc/%s/fd" % pid
    try:
        filenames = os.listdir(dirname)
    except:  # Process died
        return
    sindent = " " * indent
    for filename in filenames:
        pathname = os.path.join(dirname, filename)
        try:
            linkto = os.readlink(pathname)
        except:  # Process died
            continue
        inode_match = inode_re.match(linkto)
        if not inode_match:
            continue
        inode = int(inode_match.group(1))
        if inode not in inodes:
            continue
        if not header_printed:
            print("%s%-10s %-6s %-6s %15s:%-5s %15s:%-5s" %
                  (sindent, "State", "Recv-Q", "Send-Q",
                   "Local Address", "Port",
                   "Peer Address", "Port"))
            header_printed = True
        s = inodes[inode]
        print("%s%-10s %-6d %-6d %15s:%-5d %15s:%-5d" %
              (sindent, s.state(),
               s.receive_queue(), s.write_queue(),
               s.saddr(), s.sport(), s.daddr(), s.dport()))


def format_affinity(affinity):
    if len(affinity) <= 4:
        return ",".join(str(a) for a in affinity)

    return ",".join(str(hex(a)) for a in procfs.hexbitmask(affinity, get_nr_cpus()))


def ps_show_thread(pid, affect_children, ps, has_ctxt_switch_info, sock_inodes,
                   sock_inode_re, cgroups):
    global irqs
    try:
        affinity = format_affinity(os.sched_getaffinity(pid))
    except OSError as e:
        if e.args[0] == errno.ESRCH:
            return
        raise e

    sched = tuna_sched.sched_str(os.sched_getscheduler(pid))[6:]
    rtprio = int(ps[pid]["stat"]["rt_priority"])
    cgout = ps[pid]["cgroups"]
    cmd = ps[pid]["stat"]["comm"]
    users = ""
    if tuna.is_irq_thread(cmd):
        try:
            if not irqs:
                irqs = procfs.interrupts()
            if cmd[:4] == "IRQ-":
                users = irqs[tuna.irq_thread_number(cmd)]["users"]
                for u in users:
                    if u in get_nics():
                        users[users.index(u)] = "%s(%s)" % (
                            u, ethtool.get_module(u))
                users = ",".join(users)
            else:
                u = cmd[cmd.find('-') + 1:]
                if u in get_nics():
                    users = ethtool.get_module(u)
        except:
            users = "Not found in /proc/interrupts!"

    ctxt_switch_info = ""
    if has_ctxt_switch_info:
        voluntary_ctxt_switches = int(
            ps[pid]["status"]["voluntary_ctxt_switches"])
        nonvoluntary_ctxt_switches = int(
            ps[pid]["status"]["nonvoluntary_ctxt_switches"])
        ctxt_switch_info = " %9d %12s" % (voluntary_ctxt_switches,
                                          nonvoluntary_ctxt_switches)

    if affect_children:
        print(" %-5d " % pid, end=' ')
    else:
        print("  %-5d" % pid, end=' ')
    print("%6s %5d %8s%s %15s %s" % (sched, rtprio, affinity,
                                     ctxt_switch_info, cmd, users), end=' ')
    if cgroups:
        print(" %9s" % cgout, end=' ')
    print("")
    if sock_inodes:
        ps_show_sockets(pid, ps, sock_inodes, sock_inode_re,
                        affect_children and 3 or 4)
    if affect_children and "threads" in ps[pid]:
        for tid in list(ps[pid]["threads"].keys()):
            ps_show_thread(tid, False, ps[pid]["threads"],
                           has_ctxt_switch_info,
                           sock_inodes, sock_inode_re, cgroups)


def ps_show(ps, affect_children, thread_list, cpu_list,
            irq_list_numbers, show_uthreads, show_kthreads,
            has_ctxt_switch_info, sock_inodes, sock_inode_re, cgroups):

    ps_list = []
    for pid in list(ps.keys()):
        iskth = tuna.iskthread(pid)
        if not show_uthreads and not iskth:
            continue
        if not show_kthreads and iskth:
            continue
        in_irq_list = False
        if irq_list_numbers:
            if tuna.is_hardirq_handler(ps, pid):
                try:
                    irq = int(ps[pid]["stat"]["comm"][4:])
                    if irq not in irq_list_numbers:
                        if not thread_list:
                            continue
                    else:
                        in_irq_list = True
                except:
                    pass
            elif not thread_list:
                continue
        if not in_irq_list and thread_list and pid not in thread_list:
            continue
        try:
            affinity = os.sched_getaffinity(pid)
        except OSError as e:
            if e.args[0] == errno.ESRCH:
                continue
            raise e
        if cpu_list and not set(cpu_list).intersection(set(affinity)):
            continue
        ps_list.append(pid)

    ps_list.sort()

    for pid in ps_list:
        ps_show_thread(pid, affect_children, ps, has_ctxt_switch_info,
                       sock_inodes, sock_inode_re, cgroups)


def load_socktype(socktype, inodes):
    idiag = inet_diag.create(socktype=socktype)
    while True:
        try:
            s = idiag.get()
        except:
            break
        inodes[s.inode()] = s


def load_sockets():
    inodes = {}
    for socktype in (inet_diag.TCPDIAG_GETSOCK, inet_diag.DCCPDIAG_GETSOCK):
        load_socktype(socktype, inodes)
    return inodes


def do_ps(thread_list, cpu_list, irq_list, show_uthreads, show_kthreads,
          affect_children, show_sockets, cgroups):
    ps = procfs.pidstats()
    if affect_children:
        ps.reload_threads()

    sock_inodes = None
    sock_inode_re = None
    if show_sockets:
        sock_inodes = load_sockets()
        sock_inode_re = re.compile(r"socket:\[(\d+)\]")

    has_ctxt_switch_info = "voluntary_ctxt_switches" in ps[1]["status"]
    try:
        if sys.stdout.isatty():
            ps_show_header(has_ctxt_switch_info, cgroups)
        ps_show(ps, affect_children, thread_list,
                cpu_list, irq_list, show_uthreads, show_kthreads,
                has_ctxt_switch_info, sock_inodes, sock_inode_re, cgroups)
    except IOError:
        # 'tuna -P | head' for instance
        pass


def find_drivers_by_users(users):
    nics = get_nics()
    drivers = []
    for u in users:
        try:
            idx = u.index('-')
            u = u[:idx]
        except:
            pass
        if u in nics:
            driver = ethtool.get_module(u)
            if driver not in drivers:
                drivers.append(driver)

    return drivers


def show_irqs(irq_list, cpu_list):
    global irqs
    if not irqs:
        irqs = procfs.interrupts()

    if sys.stdout.isatty():
        print("%4s %-16s %8s" % ("#", _("users"), _("affinity"),))
    sorted_irqs = []
    for k in list(irqs.keys()):
        try:
            irqn = int(k)
            affinity = irqs[irqn]["affinity"]
        except:
            continue
        if irq_list and irqn not in irq_list:
            continue

        if cpu_list and not set(cpu_list).intersection(set(affinity)):
            continue
        sorted_irqs.append(irqn)

    sorted_irqs.sort()
    for irq in sorted_irqs:
        affinity = format_affinity(irqs[irq]["affinity"])
        users = irqs[irq]["users"]
        print("%4d %-16s %8s" % (irq, ",".join(users), affinity), end=' ')
        drivers = find_drivers_by_users(users)
        if drivers:
            print(" %s" % ",".join(drivers))
        else:
            print()


def do_list_op(op, current_list, op_list):
    if not current_list:
        current_list = []
    if op == '+':
        return list(set(current_list + op_list))
    if op == '-':
        return list(set(current_list) - set(op_list))
    return list(set(op_list))


def thread_mapper(s):
    global ps
    try:
        return [int(s), ]
    except:
        pass

    ps = procfs.pidstats()

    try:
        return ps.find_by_regex(re.compile(fnmatch.translate(s)))
    except:
        return ps.find_by_name(s)


def irq_mapper(s):
    global irqs
    try:
        return [int(s), ]
    except:
        pass
    if not irqs:
        irqs = procfs.interrupts()

    irq_list_str = irqs.find_by_user_regex(re.compile(fnmatch.translate(s)))
    irq_list = []
    for i in irq_list_str:
        try:
            irq_list.append(int(i))
        except:
            pass

    return irq_list


def pick_op(argument):
    if argument == "":
        return (None, argument)
    if argument[0] in ('+', '-'):
        return (argument[0], argument[1:])
    return (None, argument)


def i18n_init():
    (app, localedir) = ('tuna', '/usr/share/locale')
    locale.setlocale(locale.LC_ALL, '')
    gettext.bindtextdomain(app, localedir)
    gettext.textdomain(app)
    gettext.install(app, localedir)


def apply_config(filename):
    from tuna.config import Config
    config = Config()
    if os.path.exists(filename):
        config.config['root'] = os.getcwd() + "/"
        filename = os.path.basename(filename)
    else:
        if not os.path.exists(config.config['root']+filename):
            print(filename + _(" not found!"))
            exit(-1)
    if config.loadTuna(filename):
        exit(1)
    ctrl = 0
    values = {}
    values['toapply'] = {}
    for index in range(len(config.ctlParams)):
        for opt in config.ctlParams[index]:
            values['toapply'][ctrl] = {}
            values['toapply'][ctrl]['label'] = opt
            values['toapply'][ctrl]['value'] = config.ctlParams[index][opt]
            ctrl = ctrl + 1
    config.applyChanges(values)


def list_config():
    from tuna.config import Config
    config = Config()
    print(_("Preloaded config files:"))
    for value in config.populate():
        print(value)
    sys.exit(1)


def main():
    global ps

    i18n_init()
    try:
        short = "a:c:CfgGhiIKlmNp:PQq:r:s:S:t:UvWx"
        long = ["cpus=", "affect_children", "filter", "gui", "help",
                "isolate", "include", "no_kthreads", "move", "nohz_full",
                "show_sockets", "priority=", "show_threads",
                "show_irqs", "irqs=",
                "save=", "sockets=", "threads=", "no_uthreads",
                "version", "what_is", "spread", "cgroup", "config_file_apply=",
                "config_file_list=", "run="]
        if have_inet_diag:
            short += "n"
            int.append("show_sockets")
        opts, args = getopt.getopt(sys.argv[1:], short, long)
    except getopt.GetoptError as err:
        usage()
        print(str(err))
        sys.exit(2)

    run_gui = not opts
    kthreads = True
    uthreads = True
    cgroups = False
    cpu_list = None
    irq_list = None
    irq_list_str = None
    rtprio = None
    policy = None
    thread_list = []
    thread_list_str = None
    filter = False
    affect_children = False
    show_sockets = False
    p_waiting_action = False

    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            return
        if o in ("-a", "--config_file_apply"):
            apply_config(a)
        elif o in ("-l", "--config_file_list"):
            list_config()
        elif o in ("-c", "--cpus"):
            (op, a) = pick_op(a)
            try:
                op_list = tuna.cpustring_to_list(a)
            except ValueError:
                usage()
                return
            cpu_list = do_list_op(op, cpu_list, op_list)
        elif o in ("-N", "--nohz_full"):
            try:
                cpu_list = tuna.nohz_full_list()
            except:
                print("tuna: --nohz_full " +
                      _(" needs nohz_full=cpulist on the kernel command line"))
                sys.exit(2)
        elif o in ("-C", "--affect_children"):
            affect_children = True
        elif o in ("-G", "--cgroup"):
            cgroups = True
        elif o in ("-t", "--threads"):
            # The -t - will reset thread list
            if a == '-':
                thread_list = []
                thread_list_str = ''
            else:
                (op, a) = pick_op(a)
                op_list = reduce(lambda i, j: i + j,
                                 list(map(thread_mapper, a.split(","))))
                op_list = list(set(op_list))
                thread_list = do_list_op(op, thread_list, op_list)
                # Check if a process name was especified and no
                # threads was found, which would result in an empty
                # thread list, i.e. we would print all the threads
                # in the system when we should print nothing.
                if not op_list and isinstance(a, type('')):
                    thread_list_str = do_list_op(op, thread_list_str,
                                                 a.split(","))
                if not op:
                    irq_list = None
        elif o in ("-f", "--filter"):
            filter = True
        elif o in ("-g", "--gui"):
            run_gui = True
        elif o in ("-i", "--isolate"):
            if not cpu_list:
                print("tuna: --isolate " + _("requires a cpu list!"))
                sys.exit(2)
            tuna.isolate_cpus(cpu_list, get_nr_cpus())
        elif o in ("-I", "--include"):
            if not cpu_list:
                print("tuna: --include " + _("requires a cpu list!"))
                sys.exit(2)
            tuna.include_cpus(cpu_list, get_nr_cpus())
        elif o in ("-p", "--priority"):
            # Save policy and rtprio for future Actions (e.g. --run).
            (policy, rtprio) = tuna.get_policy_and_rtprio(a)
            if not thread_list:
                # For backward compatibility
                p_waiting_action = True
            else:
                try:
                    tuna.threads_set_priority(thread_list, a, affect_children)
                except OSError as err:
                    print("tuna: %s" % err)
                    sys.exit(2)
        elif o in ("-P", "--show_threads"):
            # If the user specified process names that weren't
            # resolved to pids, don't show all threads.
            if not thread_list and not irq_list:
                if thread_list_str or irq_list_str:
                    continue
            do_ps(thread_list, cpu_list, irq_list, uthreads,
                  kthreads, affect_children, show_sockets, cgroups)
        elif o in ("-Q", "--show_irqs"):
            # If the user specified IRQ names that weren't
            # resolved to IRQs, don't show all IRQs.
            if not irq_list and irq_list_str:
                continue
            show_irqs(irq_list, cpu_list)
        elif o in ("-n", "--show_sockets"):
            show_sockets = True
        elif o in ("-m", "--move", "-x", "--spread"):
            if not cpu_list:
                print("tuna: --move " + _("requires a cpu list!"))
                sys.exit(2)
            if not (thread_list or irq_list):
                print("tuna: --move " + _("requires a list of threads/irqs!"))
                sys.exit(2)

            spread = o in ("-x", "--spread")

            if thread_list:
                tuna.move_threads_to_cpu(cpu_list, thread_list, spread=spread)

            if irq_list:
                tuna.move_irqs_to_cpu(cpu_list, irq_list, spread=spread)
        elif o in ("-s", "--save"):
            save(cpu_list, thread_list, a)
        elif o in ("-S", "--sockets"):
            (op, a) = pick_op(a)
            sockets = [socket for socket in a.split(",")]

            if not cpu_list:
                cpu_list = []

            cpu_info = sysfs.cpus()
            op_list = []
            for socket in sockets:
                if socket not in cpu_info.sockets:
                    print("tuna: %s" %
                          (_("invalid socket %(socket)s sockets available: %(available)s") %
                           {"socket": socket,
                            "available": ",".join(list(cpu_info.sockets.keys()))}))
                    sys.exit(2)
                op_list += [int(cpu.name[3:])
                            for cpu in cpu_info.sockets[socket]]
            cpu_list = do_list_op(op, cpu_list, op_list)
        elif o in ("-K", "--no_kthreads"):
            kthreads = False
        elif o in ("-q", "--irqs"):
            (op, a) = pick_op(a)
            op_list = reduce(lambda i, j: i + j,
                             list(map(irq_mapper, list(set(a.split(","))))))
            irq_list = do_list_op(op, irq_list, op_list)
            # See comment above about thread_list_str
            if not op_list and isinstance(a, type('')):
                irq_list_str = do_list_op(op, irq_list_str, a.split(","))
            if not op:
                thread_list = []
            if not ps:
                ps = procfs.pidstats()
            if tuna.has_threaded_irqs(ps):
                for irq in irq_list:
                    irq_re = tuna.threaded_irq_re(irq)
                    irq_threads = ps.find_by_regex(irq_re)
                    if irq_threads:
                        # Change the affinity of the thread too
                        # as we can't rely on changing the irq
                        # affinity changing the affinity of the
                        # thread or vice versa. We need to change
                        # both.
                        thread_list += irq_threads

        elif o in ("-U", "--no_uthreads"):
            uthreads = False
        elif o in ("-v", "--version"):
            print(version)
        elif o in ("-W", "--what_is"):
            if not thread_list:
                print("tuna: --what_is " + _("requires a thread list!"))
                sys.exit(2)
            for tid in thread_list:
                thread_help(tid)
        elif o in ("-r", "--run"):
            # If -p is set, it will be consumed. So, no backward compatible
            # error handling action must be taken.
            p_waiting_action = False

            # pick_op() before run the command: to remove the prefix
            # + or - from command line.
            (op, a) = pick_op(a)

            # In order to include the new process, it must run
            # the command first, and then get the list of pids,
            tuna.run_command(a, policy, rtprio, cpu_list)

            op_list = reduce(lambda i, j: i + j,
                             list(map(thread_mapper, a.split(","))))
            op_list = list(set(op_list))
            thread_list = do_list_op(op, thread_list, op_list)

            # Check if a process name was especified and no
            # threads was found, which would result in an empty
            # thread list, i.e. we would print all the threads
            # in the system when we should print nothing.
            if not op_list and isinstance(a, type('')):
                thread_list_str = do_list_op(op, thread_list_str, a.split(","))
            if not op:
                irq_list = None

    # For backward compatibility: when -p used to be only an Action, it
    # used to exit(2) if no action was taken (i.e. if no threads_list
    # was set).
    if p_waiting_action:
        print(("tuna: -p ") + _("requires a thread list!"))
        sys.exit(2)

    if run_gui:
        try:
            from tuna import tuna_gui
        except ImportError:
            # gui packages not installed
            print(_('tuna: packages needed for the GUI missing.'))
            print(_('      Make sure xauth, pygtk2-libglade are installed.'))
            usage()
            return
        except RuntimeError:
            print("tuna: machine needs to be authorized via xhost or ssh -X?")
            return

        try:
            cpus_filtered = filter and cpu_list or []
            app = tuna_gui.main_gui(kthreads, uthreads, cpus_filtered)
            app.run()
        except KeyboardInterrupt:
            pass


if __name__ == '__main__':
    main()
