#!/usr/bin/env python3

import os
import signal
import operator
from time import sleep

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

k = 0.0042
oom_score_min = 10
t_min = 0.5

rate_mem = 3
rate_swap = 1
rate_zram = 0.5

# уровни доступной памяти для срабатывания киллеров
mem_term_level = 200 # MiB
mem_kill_level = 100 # MiB

swap_term_level = 200 # MiB
swap_kill_level = 100 # MiB

zram_term_level = 3 * 1024 # MiB
zram_kill_level = 4 * 1024 # MiB


mtl = mem_term_level * 1024
mkl = mem_kill_level * 1024
stl = swap_term_level * 1024
skl = swap_kill_level * 1024
ztl = zram_term_level * 1024
zkl = zram_kill_level * 1024


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

# перевод дроби в проценты
def percent(num):
    a = str(round(num * 100, 1)).split('.')
    a0 = a[0].rjust(3, ' ')
    a1 = a[1]
    return '{}.{}'.format(a0, a1)


# B -> MiB, KiB -> GiB
def human(num):
    a = str(round(num / 1048576, 3))
    a0 = a.split('.')[0].rjust(4, ' ')
    a1 = a.split('.')[1]
    if len(a1) == 1:
        a1 += '00'
    if len(a1) == 2:
        a1 += '0'
    return '{}.{}'.format(a0, a1)


# возвращает disksize и mem_used_total по zram id
def zram_stat(zram_id):

    try:
        with open('/sys/block/' + zram_id + '/disksize') as file:
            disksize = file.readlines()
    except FileNotFoundError:
        return '0', '0'

    if disksize == ['0\n']:
        return '0', '0'

    try:

        with open('/sys/block/' + zram_id + '/mm_stat') as file:
            mm_stat = file.readlines()[0][:-1].split(' ')

        mm_stat_list = []

        # улучшить, сократить цикл
        for i in mm_stat:
            if i != '':
                mm_stat_list.append(i)

        mem_used_total = mm_stat_list[2]

    except FileNotFoundError:

        with open('/sys/block/' + zram_id + '/mem_used_total') as file:
            mem_used_total = file.readlines()[0][:-1]

    return disksize[0][:-1], mem_used_total # BYTES, str


# имя через пид
def pid_to_name(pid):
    try:
        with open('/proc/' + pid + '/status') as f:
            for line in f:
                return line[:-1].split('\t')[1]
    except FileNotFoundError:
        return '<unknown1>'
    except ProcessLookupError:
        return '<unknown2>'


# поиск пид жертвы
def find_victim(signal):

    oom_list = []

    for i in os.listdir('/proc'):

        if i.isdigit() is not True:
            continue

        try:
            with open('/proc/' + i + '/oom_score') as file:
                oom_score = int(file.readlines()[0][:-1])
        except FileNotFoundError:
            oom_score = 0

        oom_list.append((i, oom_score))

    pid_tuple = sorted(oom_list, key=operator.itemgetter(1), reverse=True)[0]
    oom_score = pid_tuple[1]

    if oom_score >= oom_score_min:

        pid = pid_tuple[0]

        name = pid_to_name(pid)

        print('\nSend signal {} to process {}, Pid {}, oom_score {}'.format(signal, name, pid, oom_score))

        try:
            os.kill(int(pid), signal)
            print('Success\n')
        except ProcessLookupError:
            print('No such process\n')
        except PermissionError:
            print('Operation not permitted\n')

    else:

        print('\noom_score {} < oom_score_min {}\n'.format(oom_score, oom_score_min))


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


# START

# повышаем приоритет

try:
    os.nice(-20)
    print('nice = -20')
except PermissionError:
    pass


# снизить oom_adj


# ищем позиции

with open('/proc/meminfo') as file:
    mem_list = file.readlines()

mem_list_names = []
for s in mem_list:
    mem_list_names.append(s.split(':')[0])

if mem_list_names[2] != 'MemAvailable':
    print('Your Linux kernel is too old (3.14+ requied), bye!')
    exit()

swap_total_index = mem_list_names.index('SwapTotal')
swap_free_index = swap_total_index + 1

mem_total = int(mem_list[0].split(':')[1].split(' ')[-2])


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




# рабочий цикл
while True:

    # находим mem_available, swap_total, swap_free
    with open('/proc/meminfo') as f:
        for n, line in enumerate(f):
            if n == 2:
                mem_available = int(line.split(':')[1].split(' ')[-2])
                continue
            if n == swap_total_index:
                swap_total = int(line.split(':')[1].split(' ')[-2])
                continue
            if n == swap_free_index:
                swap_free = int(line.split(':')[1].split(' ')[-2])
                break

    # тут находим фулл зрам
    disksize_sum = 0
    mem_used_total_sum = 0

    for dev in os.listdir('/sys/block'):

        if dev.startswith('zram'):

            stat = zram_stat(dev)

            disksize_sum += int(stat[0])
            mem_used_total_sum += int(stat[1])

    full_zram = (disksize_sum * k + mem_used_total_sum) / 1024.0


    # если не печатать периоды, то можно это вынести в конец
    t_mem = mem_available / 1024.0 / 1024.0 / rate_mem

    t_swap = swap_free / 1024.0 / 1024.0 / rate_swap

    # fullzram может превысить 09, будет отриц значение
    # memtotal * 0.9 - это фактически макс память для зрам
    t_zram = (mem_total * 0.8 - full_zram)  / 1024.0 / 1024.0 / rate_zram
    if t_zram <= 0:
        t_zram = 0.01

    t1 = t_mem + t_swap
    t2 = t_mem + t_zram

    if t1 <= t2:
        t = t1
    else:
        t = t2


    print(
        'MA: {} G, SA: {} G, ZF: {} G, Periods: {} {}'.format(
            human(mem_available), 
            human(swap_free), 
            human(full_zram), 
            round(t1, 1), 
            round(t2, 1)
            )
        )

    if mem_available < mkl and swap_free < skl:
        print('\nmem_available < mkl and swap_free < skl')
        find_victim(signal.SIGKILL)
        sleep(t_min)
        continue

    if full_zram > zkl:
        print('\nfull_zram > zkl')
        find_victim(signal.SIGKILL)
        sleep(t_min)
        continue

    if mem_available < mtl and swap_free < stl:
        print('\nmem_available < mtl and swap_free < stl')
        find_victim(signal.SIGTERM)
        sleep(t_min)

    if full_zram > ztl:
        print('\nzram_part > zram_tl')
        find_victim(signal.SIGTERM)
        sleep(t_min)


    sleep(t)