#!/usr/bin/env python3

# nohang_notify_low_mem --mem '14% 12%' --name 'stress' --pid '6666'

# need UID=0

# output:
# Low memory: 14% 12%
# Fattest process: 6666, stress

# need to remove this slow and fat parser
from argparse import ArgumentParser

from os import listdir
from subprocess import Popen, TimeoutExpired

wait_time = 2

parser = ArgumentParser()

parser.add_argument(
    '--mem',
    help="""available memory percent (15%, for example)""",
    default=None,
    type=str
)

parser.add_argument(
    '--pid',
    help="""pid""",
    default=None,
    type=str
)

parser.add_argument(
    '--name',
    help="""process name""",
    default=None,
    type=str
)

args = parser.parse_args()

pid = args.pid
name = args.name
mem = args.mem

title = 'Low memory: {}'.format(mem)

body = 'Fattest process: <b>{}</b>, <b>{}</b>'.format(pid, name)

display_env = 'DISPLAY='
dbus_env = 'DBUS_SESSION_BUS_ADDRESS='
user_env = 'USER='


def rline1(path):
    """read 1st line from path."""
    with open(path) as f:
        for line in f:
            return line


def re_pid_environ(pid):
    """
    read environ of 1 process
    returns tuple with USER, DBUS, DISPLAY like follow:
    ('user', 'DISPLAY=:0',
     'DBUS_SESSION_BUS_ADDRESS=unix:path=/run/user/1000/bus')
    returns None if these vars is not in /proc/[pid]/environ
    """
    try:
        env = str(rline1('/proc/' + pid + '/environ'))
        if display_env in env and dbus_env in env and user_env in env:
            env_list = env.split('\x00')

            # iterating over a list of process environment variables
            for i in env_list:
                if i.startswith(user_env):
                    user = i
                    continue

                if i.startswith(display_env):
                    display = i[:10]
                    continue

                if i.startswith(dbus_env):
                    if ',guid=' in i:
                        return None
                    dbus = i
                    continue

                if i.startswith('HOME='):
                    # exclude Display Manager's user
                    if i.startswith('HOME=/var'):
                        return None

            env = user.partition('USER=')[2], display, dbus
            return env

    except FileNotFoundError:
        pass
    except ProcessLookupError:
        pass


def root_notify_env():
    """return set(user, display, dbus)"""
    unsorted_envs_list = []
    # iterates over processes, find processes with suitable env
    for pid in listdir('/proc'):
        if pid[0].isdecimal() is False:
            continue
        one_env = re_pid_environ(pid)
        unsorted_envs_list.append(one_env)
    env = set(unsorted_envs_list)
    env.discard(None)
    return env


b = root_notify_env()

# if somebody logged in with GUI
if len(b) > 0:
    # iterating over logged-in users
    for i in b:
        username, display_env, dbus_env = i[0], i[1], i[2]
        display_tuple = display_env.partition('=')
        dbus_tuple = dbus_env.partition('=')
        display_key, display_value = display_tuple[0], display_tuple[2]
        dbus_key, dbus_value = dbus_tuple[0], dbus_tuple[2]
        try:
            Popen(['sudo', '-u', username,
                'notify-send', '--icon=dialog-warning',
                '{}'.format(title), '{}'.format(body)
                ], env={
                    display_key: display_value, dbus_key: dbus_value
                    }).wait(wait_time)
        except TimeoutExpired:
            print('TimeoutExpired: notify ' + username)
else:
    print('Low memory warnings: nobody logged in with GUI. Nothing to do.')