#!/usr/bin/env python3

import os
from argparse import ArgumentParser


def psi_path_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_path_to_metrics_cpu(psi_path):
    """
    """
    with open(psi_path) as f:
        psi_list = f.readlines()

    some_list = psi_list[0].rstrip().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 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 get_psi_mem_files(cgroup2_path, met):
    """
    """
    path_list = []

    for root, dirs, files in os.walk(cgroup2_path):
        for file in files:
            path = os.path.join(root, file)
            if path.endswith('/{}.pressure'.format(met)):
                path_list.append(path)

    return path_list


def psi_path_to_cgroup2(path):
    """
    """

    if path.endswith('/cpu.pressure'):
        return path.partition(cgroup2_mountpoint)[
            2].partition('/cpu.pressure')[0]

    if path.endswith('/io.pressure'):
        return path.partition(cgroup2_mountpoint)[
            2].partition('/io.pressure')[0]

    if path.endswith('/memory.pressure'):
        return path.partition(cgroup2_mountpoint)[
            2].partition('/memory.pressure')[0]


parser = ArgumentParser()

parser.add_argument(
    '-m',
    '--metrics',
    help="""metrics (memory, io or cpu)""",
    default='memory',
    type=str
)


args = parser.parse_args()


met = args.metrics


if not (met == 'memory' or met == 'io' or met == 'cpu'):
    print('Invalid metrics:', met)
    exit(1)


psi_path = '/proc/pressure/{}'.format(met)
mounts = '/proc/mounts'
cgroup2_separator = ' cgroup2 rw,'

cgroup2_mountpoint = cgroup2_root()


if cgroup2_mountpoint is None:
    print('ERROR: cgroup_v2 hierarchy is not mounted')
    exit(1)


try:
    psi_path_to_metrics('/proc/pressure/memory')
except Exception as e:
    print('ERROR: {}'.format(e))
    print('PSI metrics are not provided by the kernel. Exit.')
    exit(1)


if cgroup2_mountpoint is not None:
    y = get_psi_mem_files(cgroup2_mountpoint, met)


path_list = get_psi_mem_files(cgroup2_mountpoint, met)


head_mem_io = '''PSI metrics: {}
cgroup_v2 mountpoint: {}
---------------------------------------------------------
         some        |          full        |
-------------------- | -------------------- | -----------
 avg10  avg60 avg300 |  avg10  avg60 avg300 | cgroup_v2
------ ------ ------ | ------ ------ ------ | -----------'''.format(
    met, cgroup2_mountpoint)


head_cpu = '''PSI metrics: {}
cgroup_v2 mountpoint: {}
----------------------------------
         some        |
-------------------- | -----------
 avg10  avg60 avg300 | cgroup_v2
------ ------ ------ | -----------'''.format(
    met, cgroup2_mountpoint)


if met == 'cpu':
    print(head_cpu)
else:
    print(head_mem_io)


if met == 'cpu':

    some_avg10, some_avg60, some_avg300 = psi_path_to_metrics_cpu(psi_path)

    print('{} {} {} | {}'.format(
        some_avg10.rjust(6),
        some_avg60.rjust(6),
        some_avg300.rjust(6),
        'SYSTEM_WIDE'))

else:

    (some_avg10, some_avg60, some_avg300, full_avg10, full_avg60, full_avg300
     ) = psi_path_to_metrics(psi_path)

    print('{} {} {} | {} {} {} | {}'.format(
        some_avg10.rjust(6),
        some_avg60.rjust(6),
        some_avg300.rjust(6),
        full_avg10.rjust(6),
        full_avg60.rjust(6),
        full_avg300.rjust(6), 'SYSTEM_WIDE'))


for psi_path in path_list:

    if met == 'cpu':
        some_avg10, some_avg60, some_avg300 = psi_path_to_metrics_cpu(psi_path)

        print('{} {} {} | {}'.format(
            some_avg10.rjust(6),
            some_avg60.rjust(6),
            some_avg300.rjust(6),
            psi_path_to_cgroup2(psi_path)))

    else:
        (some_avg10, some_avg60, some_avg300,
         full_avg10, full_avg60, full_avg300) = psi_path_to_metrics(psi_path)

        print('{} {} {} | {} {} {} | {}'.format(
            some_avg10.rjust(6),
            some_avg60.rjust(6),
            some_avg300.rjust(6),
            full_avg10.rjust(6),
            full_avg60.rjust(6),
            full_avg300.rjust(6), psi_path_to_cgroup2(psi_path)))