# -*- python -*-
# -*- coding: utf-8 -*-
# SPDX-License-Identifier: GPL-2.0-only

import copy
import errno
import os
import re
import sys
import shlex
import fnmatch
import platform
import tuna.new_eth as ethtool
import procfs
from procfs import utilist
from tuna import help
from tuna import tuna_sched

try:
    fntable
except NameError:
    fntable = []

def kthread_help(key):
    if '/' in key:
        key = key[:key.rfind('/')+1]
    return help.KTHREAD_HELP.get(key, " ")

def proc_sys_help(key):
    if not fntable:
        reg_match = ['[', '*', '?']
        for value in help.PROC_SYS_HELP:
            for char in reg_match:
                if char in value:
                    fntable.append(value)
    temp = help.PROC_SYS_HELP.get(key, "")
    if temp:
        return key + ":\n" + temp
    for value in fntable:
        if fnmatch.fnmatch(key, value):
            return key + ":\n" + help.PROC_SYS_HELP.get(value, "")
    return key

def kthread_help_plain_text(pid, cmdline):
    cmdline = cmdline.split(' ')[0]
    params = {'pid':pid, 'cmdline':cmdline}

    if iskthread(pid):
        title = _("Kernel Thread %(pid)d (%(cmdline)s):") % params
        help = kthread_help(cmdline)
    else:
        title = _("User Thread %(pid)d (%(cmdline)s):") % params
        help = title

    return help, title

def iskthread(pid):
    # FIXME: we should leave to the callers to handle all the exceptions,
    # in this function, so that they know that the thread vanished and
    # can act accordingly, removing entries from tree views, etc
    try:
        f = open(f"/proc/{pid}/smaps")
    except IOError:
        # Thread has vanished
        return True

    line = f.readline()
    f.close()
    if line:
        return False
    # Zombies also don't have smaps entries, so check the state:
    try:
        p = procfs.pidstat(pid)
    except:
        return True

    if p["state"] == 'Z':
        return False
    return True

def irq_thread_number(cmd):
    if cmd[:4] == "irq/":
        return cmd[4:cmd.find('-')]
    if cmd[:4] == "IRQ-":
        return cmd[4:]
    raise LookupError

def is_irq_thread(cmd):
    return cmd[:4] in ("IRQ-", "irq/")

def threaded_irq_re(irq):
    return re.compile(f"(irq/{irq}-.+|IRQ-{irq})")

# FIXME: Move to python-linux-procfs
def has_threaded_irqs(ps):
    irq_re = re.compile("(irq/[0-9]+-.+|IRQ-[0-9]+)")
    return len(ps.find_by_regex(irq_re)) > 0

def set_irq_affinity_filename(filename, bitmasklist):
    pathname = f"/proc/irq/{filename}"
    f = open(pathname, "w")
    text = ",".join(["%x" % a for a in bitmasklist])
    f.write(f"{text}\n")
    try:
        f.close()
    except IOError:
        # This happens with IRQ 0, for instance
        return False
    return True

def set_irq_affinity(irq, bitmasklist):
    return set_irq_affinity_filename("%d/smp_affinity" % irq, bitmasklist)

def cpustring_to_list(cpustr):
    """Convert a string of numbers to an integer list.

    Given a string of comma-separated numbers and number ranges,
    return a simple sorted list of the integers it represents.

    This function will throw exceptions for badly-formatted strings.

    Returns a list of integers."""

    fields = cpustr.strip().split(",")
    cpu_list = []
    for field in fields:
        ends = [int(a, 0) for a in field.split("-")]
        if len(ends) > 2:
            raise SyntaxError("Syntax error")
        if len(ends) == 2:
            cpu_list += list(range(ends[0], ends[1] + 1))
        else:
            cpu_list += [ends[0]]
    return list(set(cpu_list))

def list_to_cpustring(l):
    """Convert a list of integers into a range string.

    Consecutive values will be collapsed into ranges.

    This should not throw any exceptions as long as the list is all
    positive integers.

    Returns a string."""

    l = list(set(l))
    strings = []
    prev = -2
    while l:
        i = l.pop(0)
        if i - 1 == prev:
            while l:
                j = l.pop(0)
                if j - 1 != i:
                    l.insert(0, j)
                    break
                i = j
            t = strings.pop()
            if int(t) + 1 == i:
                strings.append("%s,%u" % (t, i))
            else:
                strings.append("%s-%u" % (t, i))
        else:
            strings.append("%u" % i)
        prev = i
    return ",".join(strings)

# FIXME: move to python-linux-procfs
def is_hardirq_handler(self, pid):
    PF_HARDIRQ = 0x08000000
    try:
        return int(self.processes[pid]["stat"]["flags"]) & \
        PF_HARDIRQ and True or False
    except:
        return False

def move_threads_to_cpu(cpus, pid_list, show_uthreads, show_kthreads, set_affinity_warning=None, spread=False):
    changed = False

    ps = procfs.pidstats()
    cpu_idx = 0
    nr_cpus = len(cpus)
    new_affinity = cpus
    last_cpu = max(cpus) + 1
    for pid in pid_list:
        iskth = iskthread(pid)
        if not show_uthreads and not iskth:
            continue
        if not show_kthreads and iskth:
            continue
        if spread:
            new_affinity = [cpus[cpu_idx]]
            cpu_idx += 1
            if cpu_idx == nr_cpus:
                cpu_idx = 0

        try:
            try:
                curr_affinity = os.sched_getaffinity(pid)
            except OSError as err:
                if err.args[0] == errno.ESRCH:
                    continue
                curr_affinity = None
                raise err
            if set(curr_affinity) != set(new_affinity):
                try:
                    os.sched_setaffinity(pid, new_affinity)
                    curr_affinity = os.sched_getaffinity(pid)
                except OSError as err:
                    if err.args[0] == errno.ESRCH:
                        continue
                    curr_affinity = None
                    raise err
                if set(curr_affinity) == set(new_affinity):
                    changed = True
                    if is_hardirq_handler(ps, pid):
                        try:
                            irq = int(ps[pid]["stat"]["comm"][4:])
                            bitmasklist = procfs.hexbitmask(new_affinity, last_cpu)
                            set_irq_affinity(irq, bitmasklist)
                        except:
                            pass
                elif set_affinity_warning:
                    set_affinity_warning(pid, new_affinity)
                else:
                    print("move_threads_to_cpu: %s " % \
                      (_("could not change %(pid)d affinity to %(new_affinity)s") % \
                       {'pid':pid, 'new_affinity':new_affinity}))

            # See if this is the thread group leader
            if pid not in ps:
                continue

            threads = procfs.pidstats(f"/proc/{pid}/task")
            for tid in list(threads.keys()):
                try:
                    curr_affinity = os.sched_getaffinity(tid)
                except OSError as err:
                    if err.args[0] == errno.ESRCH:
                        continue
                    raise err
                if set(curr_affinity) != set(new_affinity):
                    try:
                        os.sched_setaffinity(tid, new_affinity)
                        curr_affinity = os.sched_getaffinity(tid)
                    except OSError as err:
                        if err.args[0] == errno.ESRCH:
                            continue
                        raise err
                    if set(curr_affinity) == set(new_affinity):
                        changed = True
                    elif set_affinity_warning:
                        set_affinity_warning(tid, new_affinity)
                    else:
                        print("move_threads_to_cpu: %s " % \
                          (_("could not change %(pid)d affinity to %(new_affinity)s") % \
                           {'pid':pid, 'new_affinity':new_affinity}))
        except OSError as err:
            if err.args[0] == errno.ESRCH:
                # process died
                continue
            if err.args[0] == errno.EINVAL: # unmovable thread)
                print("thread %(pid)d cannot be moved as requested" %{'pid':pid}, file=sys.stderr)
                continue
            raise err
    return changed

def move_irqs_to_cpu(cpus, irq_list, spread=False):
    changed = 0
    unprocessed = []

    cpu_idx = 0
    nr_cpus = len(cpus)
    new_affinity = cpus
    last_cpu = max(cpus) + 1
    irqs = None
    ps = procfs.pidstats()
    for i in irq_list:
        try:
            irq = int(i)
        except:
            if not irqs:
                irqs = procfs.interrupts()
            irq = irqs.find_by_user(i)
            if not irq:
                unprocessed.append(i)
                continue
            try:
                irq = int(irq)
            except:
                unprocessed.append(i)
                continue

        if spread:
            new_affinity = [cpus[cpu_idx]]
            cpu_idx += 1
            if cpu_idx == nr_cpus:
                cpu_idx = 0

        bitmasklist = procfs.hexbitmask(new_affinity, last_cpu)
        set_irq_affinity(irq, bitmasklist)
        changed += 1
        pid = ps.find_by_name("IRQ-%d" % irq)
        if pid:
            pid = int(pid[0])
            try:
                os.sched_setaffinity(pid, new_affinity)
            except OSError as err:
                if err.args[0] == errno.ESRCH:
                    unprocessed.append(i)
                    changed -= 1
                    continue
                raise err

    return (changed, unprocessed)

def affinity_remove_cpus(affinity, cpus, nr_cpus):
    # If the cpu being isolated was the only one in the current affinity
    affinity = list(set(affinity) - set(cpus))
    if not affinity:
        affinity = list(range(nr_cpus))
        affinity = list(set(affinity) - set(cpus))
    return affinity


# Should be moved to python_linux_procfs.interrupts, shared with interrupts.parse_affinity, etc.
def parse_irq_affinity_filename(filename, nr_cpus):
    try:
        f = open(f"/proc/irq/{filename}")
    except IOError as err:
        if procfs.is_s390():
            print("This operation is not supported on s390", file=sys.stderr)
            print(f"tuna: {err}", file=sys.stderr)
            sys.exit(2)

    line = f.readline()
    f.close()
    return utilist.bitmasklist(line, nr_cpus)


def isolate_cpus(cpus, nr_cpus):
    fname = sys._getframe().f_code.co_name # Function name
    ps = procfs.pidstats()
    ps.reload_threads()
    previous_pid_affinities = {}
    for pid in list(ps.keys()):
        if procfs.cannot_set_affinity(ps, pid):
            continue
        try:
            affinity = os.sched_getaffinity(pid)
        except OSError as err:
            if err.args[0] == errno.ESRCH:
                continue
            if err.args[0] == errno.EINVAL:
                print("Function:", fname, ",", err.strerror, file=sys.stderr)
                sys.exit(2)
            raise err
        if set(affinity).intersection(set(cpus)):
            previous_pid_affinities[pid] = copy.copy(affinity)
            affinity = affinity_remove_cpus(affinity, cpus, nr_cpus)
            try:
                os.sched_setaffinity(pid, affinity)
            except OSError as err:
                if err.args[0] == errno.ESRCH:
                    continue
                if err.args[0] == errno.EINVAL:
                    print("Function:", fname, ",", err.strerror, file=sys.stderr)
                    sys.exit(2)
                if err.args[0] == errno.EBUSY:
                    comm = ps[pid].stat["comm"]
                    print(f'Warning: Unable to isolate pid {pid} [{comm}]')
                    continue
                raise err

        if "threads" not in ps[pid]:
            continue
        threads = ps[pid]["threads"]
        for tid in list(threads.keys()):
            if procfs.cannot_set_thread_affinity(ps, pid, tid):
                continue
            try:
                affinity = os.sched_getaffinity(tid)
            except OSError as err:
                if err.args[0] == errno.ESRCH:
                    continue
                if err.args[0] == errno.EINVAL:
                    print("Function:", fname, ",", err.strerror, file=sys.stderr)
                    sys.exit(2)
                raise err
            if set(affinity).intersection(set(cpus)):
                previous_pid_affinities[tid] = copy.copy(affinity)
                affinity = affinity_remove_cpus(affinity, cpus, nr_cpus)
                try:
                    os.sched_setaffinity(tid, affinity)
                except OSError as err:
                    if err.args[0] == errno.ESRCH:
                        continue
                    if err.args[0] == errno.EINVAL:
                        print("Function:", fname, ",", err.strerror, file=sys.stderr)
                        sys.exit(2)
                    raise err

    del ps

    # Now isolate it from IRQs too
    irqs = procfs.interrupts()
    previous_irq_affinities = {}
    for irq in list(irqs.keys()):
        # LOC, NMI, TLB, etc
        if "affinity" not in irqs[irq]:
            continue
        affinity = irqs[irq]["affinity"]
        if set(affinity).intersection(set(cpus)):
            previous_irq_affinities[irq] = copy.copy(affinity)
            affinity = affinity_remove_cpus(affinity, cpus, nr_cpus)
            set_irq_affinity(int(irq), procfs.hexbitmask(affinity, nr_cpus))

    affinity = parse_irq_affinity_filename("default_smp_affinity", nr_cpus)
    affinity = affinity_remove_cpus(affinity, cpus, nr_cpus)
    set_irq_affinity_filename("default_smp_affinity", procfs.hexbitmask(affinity, nr_cpus))

    return (previous_pid_affinities, previous_irq_affinities)

def include_cpus(cpus, nr_cpus):
    ps = procfs.pidstats()
    ps.reload_threads()
    previous_pid_affinities = {}
    for pid in list(ps.keys()):
        if procfs.cannot_set_affinity(ps, pid):
            continue
        try:
            affinity = os.sched_getaffinity(pid)
        except OSError as err:
            if err.args[0] == errno.ESRCH:
                continue
            raise err
        if affinity.intersection(set(cpus)) != set(cpus):
            previous_pid_affinities[pid] = copy.copy(affinity)
            affinity = list(affinity) + cpus
            try:
                os.sched_setaffinity(pid, affinity)
            except OSError as err:
                if err.args[0] == errno.ESRCH:
                    continue
                raise err

        if "threads" not in ps[pid]:
            continue
        threads = ps[pid]["threads"]
        for tid in list(threads.keys()):
            if procfs.cannot_set_thread_affinity(ps, pid, tid):
                continue
            try:
                affinity = os.sched_getaffinity(tid)
            except OSError as err:
                if err.args[0] == errno.ESRCH:
                    continue
                raise err
            if affinity.intersection(set(cpus)) != set(cpus):
                previous_pid_affinities[tid] = copy.copy(affinity)
                affinity = list(affinity) + cpus
                try:
                    os.sched_setaffinity(tid, affinity)
                except OSError as err:
                    if err.args[0] == errno.ESRCH:
                        continue
                    raise err

    del ps

    # Now include it in IRQs too
    irqs = procfs.interrupts()
    previous_irq_affinities = {}
    for irq in list(irqs.keys()):
        # LOC, NMI, TLB, etc
        if "affinity" not in irqs[irq]:
            continue
        affinity = irqs[irq]["affinity"]
        if set(affinity).intersection(set(cpus)) != set(cpus):
            previous_irq_affinities[irq] = copy.copy(affinity)
            affinity = list(set(affinity + cpus))
            set_irq_affinity(int(irq), procfs.hexbitmask(affinity, nr_cpus))

    affinity = parse_irq_affinity_filename("default_smp_affinity", nr_cpus)
    affinity = list(set(affinity + cpus))
    set_irq_affinity_filename("default_smp_affinity", procfs.hexbitmask(affinity, nr_cpus))

    return (previous_pid_affinities, previous_irq_affinities)

def get_irq_users(irqs, irq, nics=None):
    if not nics:
        nics = ethtool.get_active_devices()
    users = irqs[irq]["users"]
    for u in users:
        if u in nics:
            try:
                users[users.index(u)] = "%s(%s)" % (u, ethtool.get_module(u))
            except IOError:
                # Old kernel, doesn't implement ETHTOOL_GDRVINFO
                pass
    return users

def get_irq_affinity_text(irqs, irq):
    affinity_list = irqs[irq]["affinity"]
    try:
        return list_to_cpustring(affinity_list)
    except:
        # needs root prio to read /proc/irq/<NUM>/smp_affinity
        return ""

def get_policy_and_rtprio(parm):
    parms = parm.split(":")
    rtprio = 0
    policy = None

    try:
        cp = tuna_sched.Policy(parms[0])
    except:
        if parms[0].isdigit():
            rtprio = int(parms[0])
        else:
            raise ValueError
    else:
        policy = cp.get_policy()
        if len(parms) > 1:
            rtprio = int(parms[1])
        elif cp.is_rt():
            rtprio = 1

    return (policy, rtprio)

def thread_filtered(tid, cpus_filtered, show_kthreads, show_uthreads):
    if cpus_filtered:
        try:
            affinity = os.sched_getaffinity(tid)
        except OSError as err:
            if err.args[0] == errno.ESRCH:
                return False
            raise err

        if set(cpus_filtered + list(affinity)) == set(cpus_filtered):
            return True

    if not (show_kthreads and show_uthreads):
        kthread = iskthread(tid)
        if ((not show_kthreads) and kthread) or \
           ((not show_uthreads) and not kthread):
            return True

    return False

def irq_filtered(irq, irqs, cpus_filtered, is_root):
    if cpus_filtered and is_root:
        affinity = irqs[irq]["affinity"]
        if set(cpus_filtered + affinity) == set(cpus_filtered):
            return True

    return False

def thread_set_priority(tid, policy, rtprio):
    param = os.sched_param(rtprio)
    if not policy and policy != 0:
        policy = os.sched_getscheduler(tid)
    os.sched_setscheduler(tid, policy, param)

def threads_set_priority(tids, parm, affect_children=False):

    (policy, rtprio) = parm

    for tid in tids:
        try:
            thread_set_priority(tid, policy, rtprio)
        except OSError as err:
            if err.args[0] == errno.ESRCH:
                continue
            raise err
        if affect_children:
            for child in [int(a) for a in os.listdir("/proc/%d/task" % tid)]:
                if child != tid:
                    try:
                        thread_set_priority(child, policy, rtprio)
                    except OSError as err:
                        if err.args[0] == errno.ESRCH:
                            continue
                        raise err

class sched_tunings:
    def __init__(self, name, pid, policy, rtprio, affinity, percpu):
        self.name = name
        self.pid = pid
        self.policy = policy
        self.rtprio = int(rtprio)
        self.affinity = affinity
        self.percpu = percpu

def get_kthread_sched_tunings(proc=None):
    if not proc:
        proc = procfs.pidstats()

    kthreads = {}
    for pid in list(proc.keys()):
        name = proc[pid]["stat"]["comm"]
        # Trying to set the priority of the migration threads will
        # fail, at least on 3.6.0-rc1 and doesn't make sense anyway
        # and this function is only used to save those priorities
        # to reset them using tools like rtctl, skip those to
        # avoid sched_setscheduler/chrt to fail
        if iskthread(pid) and not name.startswith("migration/"):
            rtprio = int(proc[pid]["stat"]["rt_priority"])
            try:
                policy = os.sched_getscheduler(pid)
                affinity = os.sched_getaffinity(pid)
            except OSError as err:
                if err.args[0] == errno.ESRCH:
                    continue
                raise err
            percpu = iskthread(pid) and \
                 proc.is_bound_to_cpu(pid)
            kthreads[name] = sched_tunings(name, pid, policy, rtprio,
                                           affinity, percpu)

    return kthreads

def run_command(cmd, policy, rtprio, cpu_list, background):
    newpid = os.fork()
    if newpid == 0:
        cmd_list = shlex.split(cmd)
        pid = os.getpid()
        if rtprio or policy:
            try:
                thread_set_priority(pid, policy, rtprio)
            except (SystemError, OSError) as err:
                print(f"tuna: {err}")
                sys.exit(2)
        if cpu_list:
            try:
                os.sched_setaffinity(pid, cpu_list)
            except (SystemError, OSError) as err:
                print(f"tuna: {err}")
                sys.exit(2)

        try:
            os.execvp(cmd_list[0], cmd_list)
        except (SystemError, OSError) as err:
            print(f"tuna: {err}")
            sys.exit(2)
    else:
        if not background:
            os.waitpid(newpid, 0)

def generate_rtgroups(filename, kthreads, nr_cpus):
    f = open(filename, "w")
    f.write('''# Generated by tuna
#
# Use it with rtctl:
#
# rtctl --file %s reset
#
# Please use 'man rtctl' for more operations
#
# Associate processes into named groups with default priority and
# scheduling policy.
#
# Format is: <groupname>:<sched>:<prio>:<regex>
#
# groupname must start at beginning of line.
# sched must be one of: 'f' (fifo)
#                       'b' (batch)
#                       'r' (round-robin)
#                       'o' (other)
#                       '*' (leave alone)
# regex is an awk regex
#
# The regex is matched against process names as printed by "ps -eo cmd".

''' % filename)
    f.write(r"kthreads:*:1:*:\[.*\]$" + "\n\n")

    per_cpu_kthreads = []
    names = list(kthreads.keys())
    names.sort()
    for name in names:
        kt = kthreads[name]
        try:
            idx = name.index("/")
            common = name[:idx]
            if common in per_cpu_kthreads:
                continue
            per_cpu_kthreads.append(common)
            name = common
            if common[:5] == "sirq-":
                common = "(sirq|softirq)" + common[4:]
            elif common[:8] == "softirq-":
                common = "(sirq|softirq)" + common[7:]
                name = "s" + name[4:]
            regex = common + r"\/.*"
        except:
            idx = 0
            regex = name
        if kt.percpu or idx != 0 or name == "posix_cpu_timer":
            # Don't mess with workqueues, etc
            # posix_cpu_timer is too long and doesn't
            # have PF_THREAD_BOUND in its per process
            # flags...
            mask = "*"
        else:
            mask = ",".join([hex(a) for a in \
                     procfs.hexbitmask(kt.affinity, nr_cpus)])
        f.write(r"%s:%c:%d:%s:\[%s\]$" % (name, \
                           tuna_sched.sched_str(kt.policy)[6].lower(), \
                            kt.rtprio, mask, regex) + "\n")
    f.close()


def nohz_full_list():
    return [int(cpu) for cpu in procfs.cmdline().options["nohz_full"].split(",")]
