#!/usr/bin/env python3

import os


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


def psi_path_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 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):
    """
    """
    path_list = []

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

    return path_list


def psi_path_to_cgroup2(path):
    """
    """
    return path.partition(cgroup2_mountpoint)[2][:-16]


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


psi_path = '/proc/pressure/memory'
mounts = '/proc/mounts'
cgroup2_separator = ' cgroup2 rw,'

cgroup2_mountpoint = cgroup2_root()


if cgroup2_mountpoint is None:
    print('cgroup2 is not mounted')
else:
    print('cgroup2 mountpoint:', cgroup2_mountpoint)


psi_support = os.path.exists(psi_path)

if not psi_support:
    print('PSI is not supported, /proc/pressure/memory does not exist. Exit.')
    exit(1)


if cgroup2_mountpoint is not None:
    y = get_psi_mem_files(cgroup2_mountpoint)
    for path in y:
        pass  # print(psi_path_to_cgroup2(path))

path_list = get_psi_mem_files(cgroup2_mountpoint)

print('      avg10  avg60 avg300         avg10  avg60 avg300  cgroup2')
print('      -----  ----- ------         -----  ----- ------  ---------')

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

print('some {} {} {} | full {} {} {}  {}'.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:
    (some_avg10, some_avg60, some_avg300,
     full_avg10, full_avg60, full_avg300) = psi_path_to_metrics(psi_path)

    print('some {} {} {} | full {} {} {}  {}'.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)))