#!/usr/bin/env python3
import sys
import os
import time

infiniband_path="/sys/class/infiniband"
port_counter_list = ['port_rcv_data',
                     'port_rcv_packets',
                     'port_xmit_data',
                     'port_xmit_packets']
HCA_list = []
HCA_num  = 0

All_counter=[]
path_list=os.listdir(infiniband_path)
path_list.sort()
def init_ib_port_counter():
    #for HCA in os.listdir(infiniband_path):
    for HCA in path_list:
       # print HCA
       # if HCA == 'mlx5_2' or HCA == 'mlx5_3':
       #    continue
        HCA_dict = dict(name=HCA, port_list=[])
         
        HCA_path=os.path.join(infiniband_path, HCA, 'ports')
        for port in os.listdir(HCA_path):
            port_path = os.path.join(HCA_path, port)
            #print port_path
            port_dict = dict(name=port, counter_list=[])
            #print port_dict
            for counter in port_counter_list:
                counter_path = os.path.join(port_path, 'counters', counter)
                #print counter_path
                fh = open(counter_path)
                counter_value = int(fh.read())
                counter_dict = dict(name=counter, path=counter_path, fh=fh, bandwidth=0.0, value=counter_value, pre_value=counter_value)
                port_dict['counter_list'].append(counter_dict)
                #print item_dict
            #print port_dict
            HCA_dict['port_list'].append(port_dict)
        HCA_list.append(HCA_dict)
    # counter of all port
    for counter in port_counter_list:
        counter_dict = dict(name=counter, bandwidth=0.0)
        All_counter.append(counter_dict)

def get_ib_port_bandwidth(sleep_sec):
    #reset All counter
    for counter in All_counter:
        counter['bandwidth'] = 0.0
    for HCA in HCA_list:
        #print HCA['name']
        for port in HCA['port_list']:
            #print port['name']
            counter_id = 0
            for counter in port['counter_list']:
                #print counter['name']
                counter['fh'].seek(0)
                counter['value'] = int(counter['fh'].read())
                counter['bandwidth'] = (counter['value'] - counter['pre_value'])/sleep_sec
                counter['pre_value'] = counter['value']
                if counter['name'] == 'port_rcv_data' or counter['name'] == 'port_xmit_data' :
                    counter['bandwidth'] = counter['bandwidth']/1024.0/1024.0*4*8
                #print counter['name'], All_counter[counter_id]['name']
                All_counter[counter_id]['bandwidth'] += counter['bandwidth']
                counter_id += 1
                

def print_ib_port_bandwidth():
    title_str = '{:^10}{:^10}{:^15}{:^15}{:^15}{:^15}'.format('HCA', 'port', 'recv_data/Mbps', 'recv_pkts', 'xmit_data/Mbps', 'xmit_pkts')
    title_str += '{:^22}{:^22}'.format('rcv_pkt_avg_size/Byte', 'xmit_pkt_avg_size/Byte')
    print (title_str)
    sep_line = '{:-^80}'.format('-')
    sep_line += '{:-^44}'.format('-')
    print (sep_line)
    #print "%5s %10s %10s %10s %10s %10s" % ('HCA', 'port', 'recv_data', 'recv_pkts', 'xmit_data', 'xmit_pkts')
    for HCA in HCA_list:
        for port in HCA['port_list']:
            data_str = '{:^10}{:^10}'.format(HCA['name'], port['name'])
            for counter in port['counter_list']:
                data_str += '{:^15.2f}'.format(counter['bandwidth'])
            if port['counter_list'][1]['bandwidth'] != 0 :
                rcv_pkt_avg_size = port['counter_list'][0]['bandwidth']*1024*1024/8/port['counter_list'][1]['bandwidth']
            else:
                rcv_pkt_avg_size = 0.0
            if port['counter_list'][3]['bandwidth'] != 0 :
                xmit_pkt_avg_size = port['counter_list'][2]['bandwidth']*1024*1024/8/port['counter_list'][3]['bandwidth']
            else:
                xmit_pkt_avg_size = 0.0       
            data_str += '{:^22.2f}'.format(rcv_pkt_avg_size)
            data_str += '{:^22.2f}'.format(xmit_pkt_avg_size)
            print (data_str)
    print (sep_line)
    data_str = '{:^10}{:^10}'.format('All',' ')
    for counter in All_counter:
        data_str += '{:^15.2f}'.format(counter['bandwidth'])
    if All_counter[1]['bandwidth'] != 0.0 :
        rcv_pkt_avg_size = All_counter[0]['bandwidth']*1024*1024/8/All_counter[1]['bandwidth']
    else:
        rcv_pkt_avg_size = 0.0
    if All_counter[3]['bandwidth'] != 0.0 :
        xmit_pkt_avg_size = All_counter[2]['bandwidth']*1024*1024/8/All_counter[3]['bandwidth']
    else:
        xmit_pkt_avg_size = 0.0
    data_str += '{:^22.2f}'.format(rcv_pkt_avg_size)
    data_str += '{:^22.2f}'.format(xmit_pkt_avg_size)
    print (data_str)
    print ('\n')


if __name__ == '__main__' :
    if len(sys.argv) > 1 :               
        sleep_sec = float(sys.argv[1])
    else:
        sleep_sec = 1.0
    init_ib_port_counter()
    while(True):
        time.sleep(sleep_sec)
        get_ib_port_bandwidth(sleep_sec)
        print_ib_port_bandwidth()