#!/usr/bin/env python3

import os
from ctypes import CDLL
from time import sleep, monotonic
from operator import itemgetter
from sys import stdout, stderr, argv, exit
from re import search
from sre_constants import error as invalid_re
from signal import signal, SIGKILL, SIGTERM, SIGINT, SIGQUIT, SIGHUP, SIGUSR1, SIGSTOP, SIGCONT
import threading
import shlex
from subprocess import Popen, TimeoutExpired


cmd_list = argv[1:]

psi_support = True


exit_d = dict()
exit_d[0] = False


def find_psi_metrics_value():
    """
    """
    psi_path = '/proc/pressure/memory'
    psi_metrics = 'full_avg10'

    if psi_support:

        if psi_metrics == 'some_avg10':
            return float(rline1(psi_path).split(' ')[1].split('=')[1])
        if psi_metrics == 'some_avg60':
            return float(rline1(psi_path).split(' ')[2].split('=')[1])
        if psi_metrics == 'some_avg300':
            return float(rline1(psi_path).split(' ')[3].split('=')[1])

        if psi_metrics == 'full_avg10':
            with open(psi_path) as f:
                psi_list = f.readlines()
                return float(psi_list[1].split(' ')[1].split('=')[1])
        if psi_metrics == 'full_avg60':
            with open(psi_path) as f:
                psi_list = f.readlines()
                return float(psi_list[1].split(' ')[2].split('=')[1])
        if psi_metrics == 'full_avg300':
            with open(psi_path) as f:
                psi_list = f.readlines()
                return float(psi_list[1].split(' ')[3].split('=')[1])

def exe(cmd):
    """ execute cmd in subprocess.Popen()
    """
    cmd_list = shlex.split(cmd)

    cmd_num_dict['cmd_num'] += 1
    cmd_num = cmd_num_dict['cmd_num']

    log('Execute the command ({}) in {}: {}'.format(
        cmd_num,
        threading.current_thread().getName(),
        cmd_list))
    t3 = monotonic()
    with Popen(cmd_list) as proc:
        try:
            proc.wait(timeout=exe_timeout)
            exit_status = proc.poll()
            t4 = monotonic()
            log('Command ({}) execution completed in {} sec; exit status'
                ': {}'.format(cmd_num, round(t4 - t3, 3), exit_status))
        except TimeoutExpired:
            proc.kill()
            t4 = monotonic()
            log('TimeoutExpired for the command ({}) in {} sec'.format(
                cmd_num, round(t4 - t3, 3)))

def pid_to_autogroup(pid):
    """
    """
    try:
        return rline1('/proc/{}/autogroup'.format(pid)).partition(
            '/autogroup-')[2].partition(' nice ')[0]
    except IndexError:
        return None
    except FileNotFoundError:
        return None
    except ProcessLookupError:
        return None
    except AttributeError:
        return None

def rline1(path):
    """Read 1st line from the path."""
    try:
        with open(path) as f:
            for line in f:
                return line.rstrip()
    except UnicodeDecodeError:
        with open(path, 'rb') as f:
            return f.read(999).decode(
                'utf-8', 'ignore').split('\n')[0]  # use partition()!


def exe(cmd_list):
    """ execute cmd in subprocess.Popen()
    """
    th_name = threading.current_thread().getName()
    t3 = monotonic()
    try:
        with Popen(cmd_list) as proc:
            proc.wait()
            exit_status = proc.poll()
            t4 = monotonic()
            print('Command execution completed in {}s; exit status: {}'.format(round(t4 - t3, 3), exit_status))
            exit_d[0] = True

    except Exception as e:
        print('Exception in {}: {}'.format(th_name, e))
        exit_d[0] = True


def start_thread(func, *a, **k):
    """ run function in a new thread
    """
    th = threading.Thread(target=func, args=a, kwargs=k, daemon=True)
    th_name = th.getName()

    print('Starting {} from {}'.format(
        th_name, threading.current_thread().getName()
    ))

    try:

        t1 = monotonic()
        th.start()
        t2 = monotonic()

        print('{} has started in {} ms, {} threads are '
            'currently alive'.format(th_name, round((
                t2 - t1) * 1000, 1), threading.active_count()))

    except RuntimeError:
        print('RuntimeError: cannot start {}'.format(th_name))
        exit_d[0] = True
        return 1


def mlockall():
    """
    """
    MCL_CURRENT = 1
    MCL_FUTURE = 2
    MCL_ONFAULT = 4

    libc = CDLL('libc.so.6', use_errno=True)
    result = libc.mlockall(MCL_CURRENT | MCL_FUTURE | MCL_ONFAULT)

    if result != 0:
        result = libc.mlockall(MCL_CURRENT | MCL_FUTURE)
        if result != 0:
            print('WARNING: cannot lock all memory: [Errno {}]'.format(result))
        else:
            print('All memory locked with MCL_CURRENT | MCL_FUTURE')
    else:
        print('All memory locked with MCL_CURRENT | MCL_FUTURE | MCL_ONFAULT')


# mlockall()


def alive_pid_list():
    """
    """
    pid_list = []

    proc_list = os.listdir('/proc')

    for pid in proc_list:

        if not pid[0].isdecimal():
            continue

        pid_list.append(pid)

    return pid_list



def pid_to_name(pid):
    """
    """
    try:
        with open('/proc/{}/comm'.format(pid), 'rb', buffering=0) as f:
            return f.read().decode('utf-8', 'ignore')[:-1]
    except FileNotFoundError:
        return ''
    except ProcessLookupError:
        return ''


self_pid = os.getpid()

self_au = pid_to_autogroup('self')


pid_list = alive_pid_list()
stop_pid_set = set()
for pid in pid_list:
    au = pid_to_autogroup(pid)
    if au == self_au:
        print(pid, au, pid_to_name(pid))
        stop_pid_set.add(pid)



start_thread(exe, cmd_list)






def cont():
    """
    """
    pid_list = alive_pid_list()

    for pid in pid_list:
        au = pid_to_autogroup(pid)
        name = pid_to_name(pid)

        if au == self_au:
            if pid in stop_pid_set:
                # print('EXCLUDE', pid, au, name)
                continue
            else:
                print('CONT', pid, name)
                try:
                    os.kill(int(pid), SIGCONT)
                except (ProcessLookupError, FileNotFoundError, PermissionError):
                    pass


def stop():
    """
    """
    pid_list = alive_pid_list()

    for pid in pid_list:
        au = pid_to_autogroup(pid)
        name = pid_to_name(pid)

        if au == self_au:
            if pid in stop_pid_set:
                # print('EXCLUDE', pid, au, name)
                continue
            else:
                print('STOP', pid, name)
                try:
                    os.kill(int(pid), SIGSTOP)
                except (ProcessLookupError, FileNotFoundError, PermissionError):
                    pass





def correction():
    """
    """
    stop()
    stop()
    sleep(interval)
    while True:

        if exit_d[0]:
            cont()
            cont()
            exit()

        m = find_psi_metrics_value()
        print(m)

        if m > m_th:
            sleep(interval)
        else:
            cont()
            break


m_th = 10
interval = 1


while True:
    # print('---------------------------------')

    if exit_d[0]:
        cont()
        cont()
        exit()

    m = find_psi_metrics_value()
    print(m)

    if m > m_th:
        correction()

    sleep(interval)