#! /usr/bin/python3
import sys
from getopt import gnu_getopt, GetoptError
from dcache_nagios_plugins.dcacheinfo import load_pool
from dcache_nagios_plugins.utils import \
    int_of_sizestr, float_of_timestr, float_of_ratiostr, \
    show_size, show_time, show_percent

report_code = 0
report_msgs = []

def report(code, msg):
    global report_code
    report_code = max(code, report_code)
    if not msg in report_msgs:
        report_msgs.append(msg)

le = lambda x, y: x <= y
ge = lambda x, y: x >= y

def check_limit(what, x, op, threshold, show = show_size):
    xW, xC = threshold
    if not op(x, xC):
        report(2, '%s = %s (critical)' % (what, show(x)))
    elif not op(x, xW):
        report(1, '%s = %s (warning)' % (what, show(x)))
    else:
        report(0, '%s = %s' % (what, show(x)))

def missing_metric(what):
    report(2, 'Missing metric %s.' % what)

def check_heartbeat(pool, threshold):
    check_limit('heartbeat', pool.last_heartbeat, le, threshold, show = show_time)

def check_precious(pool, threshold):
    if pool.space_precious is None:
        missing_metric('precious')
    else:
        precious = pool.space_precious
        check_limit('precious', precious, le, threshold)

def check_nonprecious(pool, threshold):
    if pool.space_total is None:
        missing_metric('total')
    elif pool.space_precious is None:
        missing_metric('precious')
    else:
        nonprecious = pool.space_total - pool.space_precious
        check_limit('nonprecious', nonprecious, ge, threshold)

def check_precious_over_total(pool, threshold):
    if pool.space_precious is None:
        missing_metric('precious')
    elif pool.space_total is None:
        missing_metric('total')
    else:
        precious = float(pool.space_precious) / pool.space_total
        check_limit('precious/total', precious, le, threshold, show = show_percent)

def check_nonprecious_over_total(pool, threshold):
    if pool.space_total is None:
        missing_metric('total')
    elif pool.space_precious is None:
        missing_metric('precious')
    else:
        nonprecious = float(pool.space_total - pool.space_precious)/pool.space_total
        check_limit('nonprecious/total', nonprecious, ge, threshold, show = show_percent)

def threshold_arg(conv, arg):
    try:
        threshold = list(map(conv, arg.split(',')))
    except Exception as err:
        raise GetoptError('Invalid argument: %s' % err)
    if len(threshold) == 1:
        return threshold + threshold
    elif len(threshold) == 2:
        return threshold
    else:
        raise GetoptError('Thresholds options accepts pairs of one or two values.')

def print_time(hdr, amount):
    sys.stdout.write('%s: %s\n' % (hdr, show_time(amount)))

def print_space(hdr, amount):
    if amount is None:
        sys.stdout.write('%s: unknown\n' % hdr)
    else:
        sys.stdout.write('%s: %s\n' % (hdr, show_size(amount)))

info_url = None
certkey = None
cert = None
pool_name = None
heartbeat_threshold = None
precious_threshold = None
nonprecious_threshold = None
precious_over_total_threshold = None
nonprecious_over_total_threshold = None
try:
    opts, args = gnu_getopt(sys.argv[1:], 'P:U:',
        ['heartbeat=',
         'precious=', 'precious-over-total=',
         'nonprecious=', 'nonprecious-over-total=',
         'certkey=', 'cert='])
    for opt, arg in opts:
        if opt == '-P':
            pool_name = arg
        elif opt == '-U':
            info_url = arg
        elif opt == '--certkey':
            certkey = arg
        elif opt == '--cert':
            cert = arg
        elif opt == '--heartbeat':
            heartbeat_threshold = threshold_arg(float_of_timestr, arg)
        elif opt == '--precious':
            precious_threshold = threshold_arg(int_of_sizestr, arg)
        elif opt == '--nonprecious':
            nonprecious_threshold = threshold_arg(int_of_sizestr, arg)
        elif opt == '--precious-over-total':
            precious_over_total_threshold = \
                threshold_arg(float_of_ratiostr, arg)
        elif opt == '--nonprecious-over-total':
            nonprecious_over_total_threshold = \
                threshold_arg(float_of_ratiostr, arg)
        else:
            assert False
    if info_url is None:
        raise GetoptError('The -U option is required.')
    if pool_name is None:
        raise GetoptError('The -P option is required.')
except GetoptError as err:
    sys.stderr.write('%s\n' % err)
    sys.exit(64)

try:
    pool = load_pool(info_url + '/pools/' + pool_name,
                     certkey = certkey, cert = cert)
except IOError as exn:
    sys.stderr.write('Failed to fetch pool info: %s' % exn)
    sys.exit(3)

if pool is None:
    sys.stdout.write('Pool not found.')
    sys.exit(2)
else:
    if pool.space_total:
        report_msgs.append('total = %s' % show_size(pool.space_total))
    if heartbeat_threshold:
        check_heartbeat(pool, heartbeat_threshold)
    if precious_threshold:
        check_precious(pool, precious_threshold)
    if nonprecious_threshold:
        check_nonprecious(pool, nonprecious_threshold)
    if precious_over_total_threshold:
        check_precious_over_total(pool, precious_over_total_threshold)
    if nonprecious_over_total_threshold:
        check_nonprecious_over_total(pool, nonprecious_over_total_threshold)
    sys.stdout.write(', '.join(report_msgs) + '\n')
    print_time('Heartbeat', pool.last_heartbeat)
    print_space('Total space', pool.space_total)
    print_space('Precious space', pool.space_precious)
    print_space('Used space', pool.space_used)
    print_space('Free space', pool.space_free)
    sys.exit(report_code)
