#!/usr/bin/env python3

import json
import os
import re
import requests
import subprocess


JUGGLER_URL = 'http://juggler-push.search.yandex.net/events'
JUGGLER_TIMEOUT = 5

PORTO_NAME = '/'.join(os.getenv('PORTO_NAME').split('/')[:2])

DMESG_OOM_MESSAGE_RE = re.compile(r'^\[(\d+)\.\d+\] Memory cgroup out of memory: (.*)$')
DMESG_CGROUP_RE = re.compile(r'\[\d+\.\d+\] Memory cgroup stats for \/porto%(.*?)%')
OOM_CRITICAL_TIME = int(os.getenv('OOM_CRITICAL_TIME', 1800))


def send_check_result_to_juggler(status, description):
    if status not in ('OK', 'CRIT', 'WARN'):
        raise ValueError('Bad status. Expected OK, CRIT or WARN.')
    data = {
        'source': '{}.{}.{}.{}'.format(os.getenv('QLOUD_PROJECT'), os.getenv('QLOUD_APPLICATION'),
                                       os.getenv('QLOUD_ENVIRONMENT'), os.getenv('QLOUD_COMPONENT')),
        'events': [
            {
                'description': description,
                'host': os.getenv('QLOUD_HOSTNAME'),
                'instance': '',
                'service': 'dmesg-oom-check',
                'status': status,
                'tags': [
                    'oom',
                ],
            },
        ],
    }
    requests.post(JUGGLER_URL, data=json.dumps(data), timeout=JUGGLER_TIMEOUT)


def get_uptime():
    with open('/proc/uptime', 'r') as f:
        return int(f.read().split('.')[0])


def get_out_of_memory_error():
    uptime = get_uptime()
    proc = subprocess.Popen(['dmesg'], stdout=subprocess.PIPE)
    last_oom_cgroup = None
    while True:
        line = proc.stdout.readline().decode()
        if not line:
            break
        porto_match = DMESG_CGROUP_RE.match(line)
        if porto_match:
            last_oom_cgroup = porto_match.group(1)
        oom_message_match = DMESG_OOM_MESSAGE_RE.match(line)
        if oom_message_match:
            oom_time = int(oom_message_match.group(1))
            oom_error_message = oom_message_match.group(2)
            if uptime - oom_time < OOM_CRITICAL_TIME and last_oom_cgroup == PORTO_NAME:
                return oom_error_message


if __name__ == '__main__':
    if os.getenv('SKIP_OOM_CHECK'):
        send_check_result_to_juggler('OK', 'OK')
        exit()
    error_message = get_out_of_memory_error()
    if error_message:
        send_check_result_to_juggler('CRIT', error_message)
    else:
        send_check_result_to_juggler('OK', 'OK')
