#!/usr/bin/env python
# vim: tabstop=4 softtabstop=4 shiftwidth=4 textwidth=80 smarttab expandtab
#
# Copyright (C) 2016 Sangoma Technologies
# All Rights Reserved.
# Moises Silva <msilva@sangoma.com>
#

import sys
import os
import argparse
import logging
import shlex
from plumbum import local, ProcessExecutionError
from plumbum.cmd import pidof, sudo


RMS_ROOT = os.getenv('RMS_INSTALL_PATH')


def sizeof_fmt(num, suffix='B'):
    """
    http://stackoverflow.com/questions/1094841/reusable-library-to-get-human-readable-version-of-file-size
    """
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)


def memory_stat(memdir, stat):
    """ Get cgroup memory stats """
    with open('{}/memory.{}'.format(memdir, stat)) as f:
        return int(f.read())


def do_restart(cmdstr):
    """ Perform restart using the given cmd string """
    cmd = shlex.split(cmdstr)
    try:
        sudo(cmd)
    except ProcessExecutionError:
        logging.exception('Failed to restart rms-client with command %s', cmd)


def run():
    """ Run RMS watchdog checking """
    if not RMS_ROOT:
        sys.stderr.write('No RMS_INSTALL_PATH found, '
                         'did you forget to source rms-activate?')
        sys.exit(1)

    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose', default=False,
                        action='store_true',
                        help="Enable verbose output")
    args = parser.parse_args()
    loglevel = logging.DEBUG if args.verbose else logging.INFO
    logging.basicConfig(level=loglevel,
                        format='[%(asctime)s] [%(levelname)s] %(message)s')

    systemddir = '/usr/lib/systemd/system'
    if os.path.isdir(systemddir):
        memdir = '/sys/fs/cgroup/memory/system.slice/system-sangoma_rms.slice'
        restart_cmd = 'systemctl restart rms-client'
    else:
        memdir = '/cgroup/memory/sangoma_rms'
        restart_cmd = 'service rms-client restart'

    # Check if the process is running at all
    try:
        pidof('rms-client')
        logging.debug('Checking rms-client status from %s', memdir)
    except ProcessExecutionError:
        logging.debug('client is not running')
        sys.exit(0)

    if not os.path.isdir(memdir):
        logging.debug('No sangoma_rms cgroup running')
        sys.exit(0)

    memusage = 'memsw.usage_in_bytes'
    if not os.path.exists('{}/memory.{}'.format(memdir, memusage)):
        memusage = 'usage_in_bytes'

    bytes_in_use = memory_stat(memdir, memusage)
    soft_limit = memory_stat(memdir, 'soft_limit_in_bytes')

    if bytes_in_use >= soft_limit:
        logging.warning('Restarting service due to memory use of %s exceeding '
                        'the soft limit of %s', sizeof_fmt(bytes_in_use),
                        sizeof_fmt(soft_limit))
        do_restart(restart_cmd)
    else:
        logging.debug('bytes_in_use=%s, soft_limit=%s',
                      sizeof_fmt(bytes_in_use),
                      sizeof_fmt(soft_limit))

    sys.exit(0)