#!/usr/bin/env python3
"""psi-monitor"""

import os
from time import sleep
from ctypes import CDLL
from sys import stdout, exit
from argparse import ArgumentParser
from signal import signal, SIGTERM, SIGINT, SIGQUIT, SIGHUP


def form(num):
    """
    """
    s = str(num).split('.')
    return '{}.{:0<2}'.format(s[0], s[1])


def signal_handler(signum, frame):
    """
    """
    def signal_handler_inner(signum, frame):
        pass
    for i in sig_list:
        signal(i, signal_handler_inner)

    log('')

    lpd = len(peaks_dict)
    if lpd != 15:
        exit()

    log('Peak values:  avg10  avg60 avg300')

    log('-----------  ------ ------ ------')

    log('some cpu     {:>6} {:>6} {:>6}'.format(
        form(peaks_dict['c_some_avg10']),
        form(peaks_dict['c_some_avg60']),
        form(peaks_dict['c_some_avg300']),
    ))

    log('-----------  ------ ------ ------')

    log('some memory  {:>6} {:>6} {:>6}'.format(
        form(peaks_dict['m_some_avg10']),
        form(peaks_dict['m_some_avg60']),
        form(peaks_dict['m_some_avg300']),
    ))

    log('full memory  {:>6} {:>6} {:>6}'.format(
        form(peaks_dict['m_full_avg10']),
        form(peaks_dict['m_full_avg60']),
        form(peaks_dict['m_full_avg300']),
    ))

    log('-----------  ------ ------ ------')

    log('some io      {:>6} {:>6} {:>6}'.format(
        form(peaks_dict['i_some_avg10']),
        form(peaks_dict['i_some_avg60']),
        form(peaks_dict['i_some_avg300']),
    ))

    log('full io      {:>6} {:>6} {:>6}'.format(
        form(peaks_dict['i_full_avg10']),
        form(peaks_dict['i_full_avg60']),
        form(peaks_dict['i_full_avg300']),
    ))

    exit()


def cgroup2_root():
    """
    """
    with open(mounts) as f:
        for line in f:
            if cgroup2_separator in line:
                return line.partition(cgroup2_separator)[0].partition(' ')[2]


def mlockall():
    """
    """

    MCL_CURRENT = 1
    MCL_FUTURE = 2
    MCL_ONFAULT = 4

    CDLL('libc.so.6').mlockall(MCL_CURRENT | MCL_FUTURE | MCL_ONFAULT)


def psi_file_mem_to_metrics(psi_path):
    """
    """

    with open(psi_path) as f:
        psi_list = f.readlines()
    some_list, full_list = psi_list[0].split(' '), psi_list[1].split(' ')
    some_avg10 = some_list[1].split('=')[1]
    some_avg60 = some_list[2].split('=')[1]
    some_avg300 = some_list[3].split('=')[1]
    full_avg10 = full_list[1].split('=')[1]
    full_avg60 = full_list[2].split('=')[1]
    full_avg300 = full_list[3].split('=')[1]
    return (some_avg10, some_avg60, some_avg300,
            full_avg10, full_avg60, full_avg300)


def psi_file_cpu_to_metrics(psi_path):
    """
    """

    with open(psi_path) as f:
        psi_list = f.readlines()
    some_list = psi_list[0].split(' ')
    some_avg10 = some_list[1].split('=')[1]
    some_avg60 = some_list[2].split('=')[1]
    some_avg300 = some_list[3].split('=')[1]
    return (some_avg10, some_avg60, some_avg300)


def log(*msg):
    """
    """
    print(*msg)
    if separate_log:
        logging.info(*msg)


parser = ArgumentParser()

parser.add_argument(
    '-t',
    '--target',
    help="""target (cgroup2 or SYTSTEM_WIDE)""",
    default='SYSTEM_WIDE',
    type=str
)


parser.add_argument(
    '-p',
    '--period',
    help="""period in sec""",
    default=2,
    type=float
)


parser.add_argument(
    '-l',
    '--log',
    help="""path to log file""",
    default=None,
    type=str
)


args = parser.parse_args()

target = args.target

period = args.period

log_file = args.log


if log_file is None:
    separate_log = False
else:
    separate_log = True
    import logging

sig_list = [SIGTERM, SIGINT, SIGQUIT, SIGHUP]


for i in sig_list:
    signal(i, signal_handler)


if separate_log:
    logging.basicConfig(
        filename=log_file,
        level=logging.INFO,
        format="%(asctime)s: %(message)s")

log('Starting psi-monitor')
log('target: {}'.format(target))
log('period: {}'.format(period))
if log_file is not None:
    log('log file: {}'.format(log_file))


if not os.path.exists('/proc/pressure'):
    log('ERROR: /proc/pressure does not exist, PSI is not supported, exit')
    exit()


if target == 'SYSTEM_WIDE':

    cpu_file = "/proc/pressure/cpu"
    memory_file = "/proc/pressure/memory"
    io_file = "/proc/pressure/io"

else:

    mounts = '/proc/mounts'
    cgroup2_separator = ' cgroup2 rw,'
    cgroup2_mountpoint = cgroup2_root()

    if cgroup2_mountpoint is None:

        log('ERROR: unified cgroup hierarchy is not mounted, exit')
        exit()

    else:

        log('cgroup2 mountpoint: {}'.format(cgroup2_mountpoint))

    cpu_file = cgroup2_mountpoint + target + "/cpu.pressure"
    memory_file = cgroup2_mountpoint + target + "/memory.pressure"
    io_file = cgroup2_mountpoint + target + "/io.pressure"


mlockall()


log('----------------------------------------------------------------------'
    '--------------------------------------------')
log(' some cpu pressure   || some memory pressure | full memory pressure ||'
    '  some io pressure    |  full io pressure')
log('---------------------||----------------------|----------------------||'
    '----------------------|---------------------')
log(' avg10  avg60 avg300 ||  avg10  avg60 avg300 |  avg10  avg60 avg300 ||'
    '  avg10  avg60 avg300 |  avg10  avg60 avg300')
log('------ ------ ------ || ------ ------ ------ | ------ ------ ------ ||'
    ' ------ ------ ------ | ------ ------ ------')


peaks_dict = dict()


while True:

    if not os.path.exists(cpu_file):
        log('ERROR: cpu pressure file does not exist:    {}'.format(cpu_file))
        sleep(period)
        continue

    if not os.path.exists(memory_file):
        log('ERROR: memory pressure file does not exist: {}'.format(
            memory_file))
        sleep(period)
        continue

    if not os.path.exists(io_file):
        log('ERROR: io pressure file does not exist:     {}'.format(cpu_file))
        sleep(period)
        continue

    (c_some_avg10, c_some_avg60, c_some_avg300
     ) = psi_file_cpu_to_metrics(cpu_file)

    (m_some_avg10, m_some_avg60, m_some_avg300,
     m_full_avg10, m_full_avg60, m_full_avg300
     ) = psi_file_mem_to_metrics(memory_file)

    (i_some_avg10, i_some_avg60, i_some_avg300,
     i_full_avg10, i_full_avg60, i_full_avg300
     ) = psi_file_mem_to_metrics(io_file)

    log('{:>6} {:>6} {:>6} || {:>6} {:>6} {:>6} | {:>6} {:>6} {:>6} || {:>6}'
        ' {:>6} {:>6} | {:>6} {:>6} {:>6}'.format(

            c_some_avg10, c_some_avg60, c_some_avg300,

            m_some_avg10, m_some_avg60, m_some_avg300,
            m_full_avg10, m_full_avg60, m_full_avg300,

            i_some_avg10, i_some_avg60, i_some_avg300,
            i_full_avg10, i_full_avg60, i_full_avg300

        ))

    c_some_avg10 = float(c_some_avg10)
    if ('c_some_avg10' not in peaks_dict or
            peaks_dict['c_some_avg10'] < c_some_avg10):
        peaks_dict['c_some_avg10'] = c_some_avg10

    c_some_avg60 = float(c_some_avg60)
    if ('c_some_avg60' not in peaks_dict or
            peaks_dict['c_some_avg60'] < c_some_avg60):
        peaks_dict['c_some_avg60'] = c_some_avg60

    c_some_avg300 = float(c_some_avg300)
    if ('c_some_avg300' not in peaks_dict or
            peaks_dict['c_some_avg300'] < c_some_avg300):
        peaks_dict['c_some_avg300'] = c_some_avg300

    #######################################################################

    m_some_avg10 = float(m_some_avg10)
    if ('m_some_avg10' not in peaks_dict or
            peaks_dict['m_some_avg10'] < m_some_avg10):
        peaks_dict['m_some_avg10'] = m_some_avg10

    m_some_avg60 = float(m_some_avg60)
    if ('m_some_avg60' not in peaks_dict or
            peaks_dict['m_some_avg60'] < m_some_avg60):
        peaks_dict['m_some_avg60'] = m_some_avg60

    m_some_avg300 = float(m_some_avg300)
    if ('m_some_avg300' not in peaks_dict or
            peaks_dict['m_some_avg300'] < m_some_avg300):
        peaks_dict['m_some_avg300'] = m_some_avg300

    m_full_avg10 = float(m_full_avg10)
    if ('m_full_avg10' not in peaks_dict or
            peaks_dict['m_full_avg10'] < m_full_avg10):
        peaks_dict['m_full_avg10'] = m_full_avg10

    m_full_avg60 = float(m_full_avg60)
    if ('m_full_avg60' not in peaks_dict or
            peaks_dict['m_full_avg60'] < m_full_avg60):
        peaks_dict['m_full_avg60'] = m_full_avg60

    m_full_avg300 = float(m_full_avg300)
    if ('m_full_avg300' not in peaks_dict or
            peaks_dict['m_full_avg300'] < m_full_avg300):
        peaks_dict['m_full_avg300'] = m_full_avg300

    #######################################################################

    i_some_avg10 = float(i_some_avg10)
    if ('i_some_avg10' not in peaks_dict or
            peaks_dict['i_some_avg10'] < i_some_avg10):
        peaks_dict['i_some_avg10'] = i_some_avg10

    i_some_avg60 = float(i_some_avg60)
    if ('i_some_avg60' not in peaks_dict or
            peaks_dict['i_some_avg60'] < i_some_avg60):
        peaks_dict['i_some_avg60'] = i_some_avg60

    i_some_avg300 = float(i_some_avg300)
    if ('i_some_avg300' not in peaks_dict or
            peaks_dict['i_some_avg300'] < i_some_avg300):
        peaks_dict['i_some_avg300'] = i_some_avg300

    i_full_avg10 = float(i_full_avg10)
    if ('i_full_avg10' not in peaks_dict or
            peaks_dict['i_full_avg10'] < i_full_avg10):
        peaks_dict['i_full_avg10'] = i_full_avg10

    i_full_avg60 = float(i_full_avg60)
    if ('i_full_avg60' not in peaks_dict or
            peaks_dict['i_full_avg60'] < i_full_avg60):
        peaks_dict['i_full_avg60'] = i_full_avg60

    i_full_avg300 = float(i_full_avg300)
    if ('i_full_avg300' not in peaks_dict or
            peaks_dict['i_full_avg300'] < i_full_avg300):
        peaks_dict['i_full_avg300'] = i_full_avg300

    stdout.flush()
    sleep(period)