#!/usr/bin/env python3
import argparse

import psutil

from checker.markdown import list_to_markdown_table

# Parse the bandwidth limit argument
parser = argparse.ArgumentParser(description='Check network interface bandwidth utilization.')
parser.add_argument('--bandwidth', type=float, required=True, help='Bandwidth speed in Mbps.')
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.')
args = parser.parse_args()

data = []
warn_value = (args.bandwidth * args.warn)
crit_value = (args.bandwidth * args.critical)

# 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():
    # Get the number of bytes sent and received
    bytes_sent = stats.bytes_sent
    bytes_recv = stats.bytes_recv

    # Wait for 1 second
    psutil.cpu_percent(interval=1)

    # Get the number of bytes sent and received after 1 second
    new_stats = psutil.net_io_counters(pernic=True)[interface]
    new_bytes_sent = new_stats.bytes_sent
    new_bytes_recv = new_stats.bytes_recv

    # Calculate the bandwidth utilization in bits per second
    bandwidth_utilization = (8 * (new_bytes_sent - bytes_sent + new_bytes_recv - bytes_recv)) / (1 * 1000 * 1000)
    data.append([interface, bandwidth_utilization, 'none'])

critical = []
warn = []
ok = []
perf_data = []
for i in range(len(data)):
    interface = data[i][0]
    bandwidth_utilization = data[i][1]
    if bandwidth_utilization >= crit_value:
        critical.append(interface)
        data[i][2] = 'critical'
    elif bandwidth_utilization >= warn_value:
        warn.append(interface)
        data[i][2] = 'warning'
    else:
        ok.append(interface)
        data[i][2] = 'ok'
    perf_data.append(f'{interface}={round(bandwidth_utilization, 2)}B;{warn_value};{crit_value}; ')

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'))

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