#!/usr/bin/env python3
import argparse
import re
import sys
import time
import traceback

import psutil

import checker.nagios as nagios
from checker.markdown import list_to_markdown_table

parser = argparse.ArgumentParser(description='Check network interface bandwidth utilization.')
parser.add_argument('--bandwidth', type=float, required=True, help='Bandwidth speed in Mbps. Used to calculate percentage.')
parser.add_argument('--critical', type=int, default=75, help='Critical if percent of bandwidth usage is greater than or equal to this.')
parser.add_argument('--warn', type=int, default=50, help='Warning if percent of bandwidth usage is greater than or equal to this.')
parser.add_argument('--max', type=int, default=None, help='Set the max value the bandwidth can be. Useful for graphs and whatever.')
parser.add_argument('--ignore', nargs='*', default=[], help='Interface names to ignore, separated by a space.')
parser.add_argument('--ignore-re', default=None, help='Regex matching interface names to ignore.')
args = parser.parse_args()

if args.ignore_re:
    ignore_re = re.compile(args.ignore_re)
else:
    ignore_re = None


def get_network_traffic(interface):
    net_io = psutil.net_io_counters(pernic=True)
    if interface in net_io:
        return net_io[interface]
    else:
        raise ValueError(f"Interface '{interface}' not found")


def calculate_network_traffic(interface, interval=1):
    initial_traffic = get_network_traffic(interface)
    start_time = time.perf_counter()

    # Should be more accurate that time.sleep()
    while True:
        current_time = time.perf_counter()
        elapsed_time = current_time - start_time
        if elapsed_time >= interval:
            break

    final_traffic = get_network_traffic(interface)

    sent_bytes = final_traffic.bytes_sent - initial_traffic.bytes_sent
    recv_bytes = final_traffic.bytes_recv - initial_traffic.bytes_recv

    sent_speed = sent_bytes / elapsed_time
    recv_speed = recv_bytes / elapsed_time

    # Convert bytes per second to megabits per second
    sent_speed_mbps = sent_speed * 8 / (1024 * 1024)
    recv_speed_mbps = recv_speed * 8 / (1024 * 1024)

    return sent_speed_mbps, recv_speed_mbps


def main():
    data = []
    warn_value = (args.bandwidth * args.warn / 100)
    crit_value = (args.bandwidth * args.critical / 100)

    # Get network interface statistics
    net_io_counters = psutil.net_io_counters(pernic=True)

    # Calculate bandwidth utilization for each interface
    for interface, stats in net_io_counters.items():
        if interface in args.ignore or (ignore_re and ignore_re.search(interface)):
            continue
        sent_speed, recv_speed = calculate_network_traffic(interface)
        bandwidth_utilization = sent_speed + recv_speed
        data.append([interface, sent_speed, recv_speed, bandwidth_utilization, 'none'])

    exit_code = nagios.OK
    critical = []
    warn = []
    ok = []
    perf_data = []
    for i in range(len(data)):
        interface = data[i][0]
        bandwidth_utilization = data[i][3]
        if bandwidth_utilization >= crit_value:
            critical.append(interface)
            state = 'critical'
            exit_code = nagios.CRITICAL
        elif bandwidth_utilization >= warn_value:
            warn.append(interface)
            state = 'warning'
            if exit_code < nagios.WARNING:
                exit_code = nagios.WARNING
        else:
            ok.append(interface)
            state = 'ok'
        data[i][2] = f'[{state.upper()}]'
        perf_data.append(f'{interface}={round(bandwidth_utilization, 2)}Mbps;{warn_value};{crit_value};{f"0;{args.max};" if args.max else ""} ')

    if len(ok):
        print(f'OK: {", ".join(ok)}')
    if len(warn):
        print(f'WARNING: {", ".join(warn)}')
    if len(critical):
        print(f'CRITICAL: {", ".join(critical)}')

    data = [(x[0], f'{round(x[1], 2)} Mbps', x[2]) for x in data]
    data.insert(0, ('Interface', 'Bandwidth', 'State'))
    print(list_to_markdown_table(data, align='left', seperator='!', borders=False))

    print(f' |{"".join(perf_data)}')
    sys.exit(exit_code)


if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f'UNKNOWN: exception "{e}"')
        print(traceback.format_exc())
        sys.exit(nagios.UNKNOWN)