#!/usr/bin/env python3

import os
from time import sleep
from ctypes import CDLL
from argparse import ArgumentParser


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


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

    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:
            pass
        else:
            pass
    else:
        pass


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

    with open(psi_path) as f:
        psi_list = f.readlines()
    # print(psi_list)
    some_list, full_list = psi_list[0].split(' '), psi_list[1].split(' ')
    # print(some_list, full_list)
    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()
    # print(psi_list)
    some_list = psi_list[0].split(' ')
    # print(some_list, full_list)
    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


def sleeeep():
    try:
        sleep(period)
    except KeyboardInterrupt:
        log('Exit')
        exit()


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

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

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


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('Starting psi-monitor, target: {}, period: {}'.format(target, period))
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('------ ------ ------ || ------ ------ ------ | ------ ------ ------ ||'
    ' ------ ------ ------ | ------ ------ ------')


while True:

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

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

    if not os.path.exists(io_file):
        log('ERROR: io pressure file does not exist:     {}'.format(cpu_file))
        sleeeep()
        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('{} {} {} || {} {} {} | {} {} {} || {} {} {} | {} {} {}'.format(

        c_some_avg10.rjust(6),
        c_some_avg60.rjust(6),
        c_some_avg300.rjust(6),

        m_some_avg10.rjust(6),
        m_some_avg60.rjust(6),
        m_some_avg300.rjust(6),
        m_full_avg10.rjust(6),
        m_full_avg60.rjust(6),
        m_full_avg300.rjust(6),

        i_some_avg10.rjust(6),
        i_some_avg60.rjust(6),
        i_some_avg300.rjust(6),
        i_full_avg10.rjust(6),
        i_full_avg60.rjust(6),
        i_full_avg300.rjust(6)

    ))

    sleeeep()