#!/usr/bin/python3
#
# Copyright (C) 2020-2023 NVIDIA Corporation. ALL RIGHTS RESERVED.
# Copyright 2013-2014.  Cumulus Networks, Inc.
#
# cl-resource-query --
#    tool to report hardware route table occupancy
#
try:
    import argparse
    import json
    import sys
    from subprocess import Popen, PIPE
    import os
except ImportError as e:
    raise ImportError(str(e) + "- required module not found")

stat_path = '/cumulus/switchd/run/'

# Used to dump json output, and for cumulative count
stats_dict = {}

def get_cumulative_usage_count(func_arg1, func_arg2):
    count = 0
    for key in func_arg1:
        try:
            count = count + func_arg2[key]["count"]
        except Exception as e:
            print(e)
            return 0
    return count


def get_entry_count_from_software_tables(func_arg1, func_arg2):
    """docstring for get_entry_count_from_software_tables"""
    switchd_sw_tables = "/cumulus/switchd/run/software-tables/root"
    arg2_table_id = 0
    arg2_count = 0

    for l in Popen(['/bin/cat', switchd_sw_tables], shell=False, stdout=PIPE).stdout:
        if (l.find(func_arg1.encode()) != -1):
            arg2_table_id = str(int(l.split(" ".encode())[1]))

    if (arg2_table_id == 0):
        return 0

    arg2_table = "/cumulus/switchd/run/software-tables/" + \
        arg2_table_id
    try:
        stream = open(arg2_table, "r")
    except Exception as e:
        print(e)
        return 0
    arg2_parser = stream.readlines()
    for line in arg2_parser:
        if func_arg2 in line:
            arg2_count+=1
    return arg2_count

class stat:

    def __init__(self, name, key, count_file, max_file, alloc_file=None, rule_class_list = None, att = None, count_function = None, func_arg1 = None, func_arg2 = None):
        self.name = name
        self.key = key
        self.count_file = count_file
        self.max_file = max_file
        self.alloc_file = alloc_file
        self.rule_class_list = rule_class_list
        self.att = att
        self.count_function = count_function
        self.func_arg1 = func_arg1
        self.func_arg2 = func_arg2



stat_list = [stat('Forwarding Profile',  'forwarding_profile',      'datapath_info/forwarding_profile',      None),
             stat('Host 0 entries',  'host_0_entry',      'route_info/host/count_0',      'route_info/host/max_0'),
             stat('Host 1 entries',  'host_1_entry',      'route_info/host/count_1',      'route_info/host/max_1'),
             stat('IPv4 neighbors',  'host_v4_entry',     'route_info/host/count_v4',     None),
             stat('IPv6 neighbors',  'host_v6_entry',     'route_info/host/count_v6',     None),
             stat('Route 0 entries', 'route_0_entry',     'route_info/route/count_0',     'route_info/route/max_0'),
             stat('Route 1 entries', 'route_1_entry',     'route_info/route/count_1',     'route_info/route/max_1'),
             stat('IPv4 Routes',     'route_v4_entry',    'route_info/route/count_v4',    None),
             stat('IPv6 Routes',     'route_v6_entry',    'route_info/route/count_v6',    None),
             stat('Total Routes',    'route_total_entry', 'route_info/route/count_total', 'route_info/route/max_total'),
             stat('ECMP nexthops',   'ecmp_nh_entry',     'route_info/ecmp_nh/count',     'route_info/ecmp_nh/max'),
             stat('ECMP entries',    'ecmp_entry',        'ecmp/count',                   None,  att = 'mlx'),
             stat('MAC entries',     'mac_entry',         'route_info/mac/count',         'route_info/mac/max', att = 'bcm'),
             stat('MAC entries',     'mac_entry', ['route_info/mac/count', 'route_info/route/router_mac_count'], 'route_info/mac/max', att = 'mlx'),
             stat('Total Mcast Routes', 'mroute_total_entry', 'route_info/mroute/hw_count_total', 'route_info/mroute/max_total'),
             stat('MPLS Labels',     'mpls_entry',        'mpls/count',                   'mpls/max'),
             stat('Ingress ACL entries', 'in_acl_entry',  'acl_info/ingress/entries',     'acl_info/ingress/entries_total'),
             stat('Ingress ACL counters','in_acl_counter','acl_info/ingress/counters',    'acl_info/ingress/counters_total', att = 'bcm'),
             stat('Ingress ACL meters',  'in_acl_meter',  'acl_info/ingress/meters',      'acl_info/ingress/meters_total', att = 'bcm'),
             stat('Ingress ACL slices',  'in_acl_slice',  'acl_info/ingress/slices',      'acl_info/ingress/slices_total', att = 'bcm'),
             stat('Egress ACL entries',  'eg_acl_entry',  'acl_info/egress/entries',      'acl_info/egress/entries_total'),
             stat('Egress ACL counters', 'eg_acl_counter','acl_info/egress/counters',     'acl_info/egress/counters_total', att = 'bcm'),
             stat('Egress ACL meters',   'eg_acl_meter',  'acl_info/egress/meters',       'acl_info/egress/meters_total', att = 'bcm'),
             stat('Egress ACL slices',   'eg_acl_slice',  'acl_info/egress/slices',       'acl_info/egress/slices_total', att = 'bcm'),

             stat('Ingress ACL ipv4_mac filter table',   'in_acl_v4mac_filter',  'acl_info/ingress/v4mac_filter/entries_used', 'acl_info/ingress/v4mac_filter/entries_max', 'acl_info/ingress/v4mac_filter/entries_allocated', att = 'bcm'),
             stat('Ingress ACL ipv6 filter table',   'in_acl_v6_filter',  'acl_info/ingress/v6_filter/entries_used', 'acl_info/ingress/v6_filter/entries_max', 'acl_info/ingress/v6_filter/entries_allocated', att = 'bcm'),
             stat('Ingress ACL mirror table',   'in_acl_mirror_filter',  'acl_info/ingress/mirror_filter/entries_used', 'acl_info/ingress/mirror_filter/entries_max', 'acl_info/ingress/mirror_filter/entries_allocated', att = 'bcm'),
             stat('Ingress ACL 8021x filter table',   'in_acl_8021x_filter',  'acl_info/ingress/8021x_filter/entries_used', 'acl_info/ingress/8021x_filter/entries_max', 'acl_info/ingress/8021x_filter/entries_allocated', att = 'bcm'),
             stat('Ingress PBR ipv4_mac filter table',   'in_pbr_v4mac_filter',  'iprule/info/ingress/v4mac_filter/entries_used', 'iprule/info/ingress/v4mac_filter/entries_max', 'iprule/info/ingress/v4mac_filter/entries_allocated', att = 'bcm'),
             stat('Ingress PBR ipv6 filter table',   'in_pbr_v6_filter',  'iprule/info/ingress/v6_filter/entries_used', 'iprule/info/ingress/v6_filter/entries_max', 'iprule/info/ingress/v6_filter/entries_allocated', att = 'bcm'),

             stat('Ingress ACL ipv4_mac mangle table',   'in_acl_v4mac_mangle',  'acl_info/ingress/v4mac_mangle/entries_used', 'acl_info/ingress/v4mac_mangle/entries_max', 'acl_info/ingress/v4mac_mangle/entries_allocated', att = 'bcm'),
             stat('Ingress ACL ipv6 mangle table',   'in_acl_v6_mangle',  'acl_info/ingress/v6_mangle/entries_used', 'acl_info/ingress/v6_mangle/entries_max', 'acl_info/ingress/v6_mangle/entries_allocated', att = 'bcm'),

             stat('Egress ACL ipv4_mac filter table',   'eg_acl_v4mac_filter',  'acl_info/egress/v4mac_filter/entries_used', 'acl_info/egress/v4mac_filter/entries_max', 'acl_info/egress/v4mac_filter/entries_allocated', att = 'bcm'),
             stat('Egress ACL ipv6 filter table',   'eg_acl_v6_filter',  'acl_info/egress/v6_filter/entries_used', 'acl_info/egress/v6_filter/entries_max', 'acl_info/egress/v6_filter/entries_allocated', att = 'bcm'),

             stat('ACL L4 port range checkers',   'acl_l4_port_range_checkers',  'acl_info/l4_port_range_checkers/entries_used', 'acl_info//l4_port_range_checkers/entries_max', att = 'bcm'),

             stat('ACL Regions',   'acl_mlx_regions',  'acl_info/mlx_regions/used', 'acl_info/mlx_regions/max', att = 'mlx'),
             stat('ACL 18B Rules Key',   'acl_mlx_18b_rules_key',  'acl_info/mlx_18b_rules_key/entries_used', 'acl_info/mlx_18b_rules_key/entries_max', att = 'mlx'),
             stat('ACL 36B Rules Key',   'acl_mlx_36b_rules_key',  'acl_info/mlx_32b_rules_key/entries_used', 'acl_info/mlx_32b_rules_key/entries_max', att = 'mlx'),
             stat('ACL 54B Rules Key',   'acl_mlx_54b_rules_key',  'acl_info/mlx_54b_rules_key/entries_used', 'acl_info/mlx_54b_rules_key/entries_max', att = 'mlx'),
             stat('Ingress ACL mac filter table',   'in_acl_mac_filter',  'acl_info/ingress/mac_filter/entries_used', None, None, ['acl_info/ingress/mac_filter_18b/entries_used','acl_info/ingress/mac_filter_36b/entries_used','acl_info/ingress/mac_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv4 filter table',   'in_acl_v4_filter',  'acl_info/ingress/v4_filter/entries_used', None, None, ['acl_info/ingress/v4_filter_18b/entries_used','acl_info/ingress/v4_filter_36b/entries_used','acl_info/ingress/v4_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv6 filter table',   'in_acl_v6_filter',  'acl_info/ingress/v6_filter/entries_used', None, None, ['acl_info/ingress/v6_filter_18b/entries_used','acl_info/ingress/v6_filter_36b/entries_used','acl_info/ingress/v6_filter_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL mac filter table',   'eg_acl_mac_filter',  'acl_info/egress/mac_filter/entries_used', None, None, ['acl_info/egress/mac_filter_18b/entries_used','acl_info/egress/mac_filter_36b/entries_used','acl_info/egress/mac_filter_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL ipv4 filter table',   'eg_acl_v4_filter',  'acl_info/egress/v4_filter/entries_used', None, None, ['acl_info/egress/v4_filter_18b/entries_used','acl_info/egress/v4_filter_36b/entries_used','acl_info/egress/v4_filter_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL ipv6 filter table',   'eg_acl_v6_filter',  'acl_info/egress/v6_filter/entries_used', None, None, ['acl_info/egress/v6_filter_18b/entries_used','acl_info/egress/v6_filter_36b/entries_used','acl_info/egress/v6_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv4 mangle table',   'in_acl_v4_mangle',  'acl_info/ingress/v4_mangle/entries_used', None, None,['acl_info/ingress/v4_mangle_18b/entries_used','acl_info/ingress/v4_mangle_36b/entries_used','acl_info/ingress/v4_mangle_54b/entries_used'], att = 'mlx'),
             stat('Ingress ACL ipv6 mangle table',   'in_acl_v6_mangle',  'acl_info/ingress/v6_mangle/entries_used', None, None, ['acl_info/ingress/v6_mangle_18b/entries_used','acl_info/ingress/v6_mangle_36b/entries_used','acl_info/ingress/v6_mangle_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL ipv4 mangle table',   'eg_acl_v4_mangle',  'acl_info/egress/v4_mangle/entries_used', None, None,['acl_info/egress/v4_mangle_18b/entries_used','acl_info/egress/v4_mangle_36b/entries_used','acl_info/egress/v4_mangle_54b/entries_used'], att = 'mlx'),
             stat('Egress ACL ipv6 mangle table',   'eg_acl_v6_mangle',  'acl_info/egress/v6_mangle/entries_used', None, None, ['acl_info/egress/v6_mangle_18b/entries_used','acl_info/egress/v6_mangle_36b/entries_used','acl_info/egress/v6_mangle_54b/entries_used'], att = 'mlx'),
             stat('Ingress PBR ipv4 filter table',   'in_pbr_v4_filter',  'iprule/info/ingress/v4_filter/entries_used', None, None, ['iprule/info/ingress/v4_filter_18b/entries_used','iprule/info/ingress/v4_filter_36b/entries_used','iprule/info/ingress/v4_filter_54b/entries_used'], att = 'mlx'),
             stat('Ingress PBR ipv6 filter table',   'in_pbr_v6_filter',  'iprule/info/ingress/v6_filter/entries_used', None, None, ['iprule/info/ingress/v6_filter_18b/entries_used','iprule/info/ingress/v6_filter_36b/entries_used','iprule/info/ingress/v6_filter_54b/entries_used'], att = 'mlx'),
             stat('Flow Counters',   'flow',  'stats/counters/flow/used', 'stats/counters/flow/max', att = 'mlx'),
             stat('RIF Basic Counters',   'rif_basic',  'stats/counters/rif_basic/used', 'stats/counters/rif_basic/max', att = 'mlx'),
             stat('RIF Enhanced Counters',   'rif_enh',  'stats/counters/rif_enh/used', 'stats/counters/rif_enh/max', att = 'mlx'),
             stat('Static SNAT entries',          'static_snat_count',    'nat_info/static/snat_count',       'nat_info/static/snat_max'),
             stat('Static DNAT entries',          'static_dnat_count',    'nat_info/static/dnat_count',       'nat_info/static/dnat_max'),
             stat('Dynamic SNAT entries',         'dynamic_snat_count',   'nat_info/dynamic/snat_count',      'nat_info/dynamic/snat_max'),
             stat('Dynamic DNAT entries',         'dynamic_dnat_count',   'nat_info/dynamic/dnat_count',      'nat_info/dynamic/dnat_max'),
             stat('Dynamic Config SNAT entries',  'config_snat_count',    'nat_info/config_rule/snat_count',  'nat_info/config_rule/snat_max'),
             stat('Dynamic Config DNAT entries',  'config_dnat_count',    'nat_info/config_rule/dnat_count',  'nat_info/config_rule/dnat_max'),
             stat('Downstream-VNI FID count',   'dvni_fid_count',  None, None, att = 'mlx', count_function=get_entry_count_from_software_tables, func_arg1="HAL Downstream-VNI Table", func_arg2="hal-dvni-entry"),
             stat('Vport FID count',   'vport_fid_count',  None, None, att = 'mlx', count_function=get_entry_count_from_software_tables, func_arg1="Backend VFID to Bridge,VID Table", func_arg2="bridge_id-"),
             stat('Total FID count',   'total_fid_count',  None, None, att = 'mlx', count_function=get_cumulative_usage_count, func_arg1=["dvni_fid_count","vport_fid_count"], func_arg2=stats_dict),
             ]

class bcolors:
    HEADER  = '\033[95m'
    OKBLUE  = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL    = '\033[91m'
    ENDC    = '\033[0m'

    def disable(self):
        self.HEADER  = ''
        self.OKBLUE  = ''
        self.OKGREEN = ''
        self.WARNING = ''
        self.FAIL    = ''
        self.ENDC    = ''

def get_max_broadcast_domain():
    """max broadcast domain count """
    max_bd = "/cumulus/switchd/run/vlan/max_bridge_domains"

    try:
        stream = open(max_bd, "r")
    except Exception as e:
        print(e)
        return 0
    arg2_parser = stream.readlines()
    for line in arg2_parser:
        return int(line)

def generate_report(key_value_flag, use_json):
    global stats_dict
    MLX_MAX_FID_COUNT = get_max_broadcast_domain()

    read_as_string = False
    str_val = ""
    # fetch the route table mode
    route_mode_file = stat_path + 'route_info/route/mode'
    try:
        with open(route_mode_file, 'r') as f:
            route_mode = int(f.readline())
    except IOError:
        print('route table mode is not available')
        return -1

    # fetch the host table mode
    host_mode_file = stat_path + 'route_info/host/mode'
    try:
        with open(host_mode_file, 'r') as f:
            host_mode = int(f.readline())
    except IOError:
        print('host table mode is not available')
        return -1

    # fetch nat support state
    nat_support_file = stat_path + 'nat_support'
    try:
        with open(nat_support_file, 'r') as f:
            nat_support = (f.readline())
    except IOError:
        print('nat support file is not available')
        return -1

    #fetch system/asic info
    system_info_file = '/cumulus/switchd/config/system_info'
    try:
        with open(system_info_file, 'r') as f:
            system_info = (f.readline())
    except IOError:
        print('system info file is not available')
        return -1

    forwarding_profile = ''
    for stat in stat_list:
        max_value = 0

        if stat.key == 'route_0_entry':
            if route_mode == 1:
                stat.name = 'IPv4/IPv6 route entries'
            elif route_mode == 2:
                stat.name = 'IPv4 route entries'

        if stat.key == 'route_1_entry':
            if route_mode == 1:
                stat.name = 'Long IPv6 route entries'
            elif route_mode == 2:
                stat.name = 'IPv6 route entries'

        if stat.key == 'host_0_entry':
            if host_mode == 1:
                stat.name = 'IPv4/IPv6 host entries'
            elif host_mode == 2:
                stat.name = 'IPv4 host entries'

        if stat.key == 'host_1_entry':
            if host_mode == 1:
                # not used
                continue
            elif host_mode == 2:
                stat.name = 'IPv6 host entries'

        if stat.key == 'forwarding_profile':
            read_as_string = True
        else:
            read_as_string = False

        if stat.key == 'mpls_entry':
           if forwarding_profile != 'mpls':
               continue

        if 'nat' in stat.key:
            if 'TRUE' in nat_support:
                if (stat.key == 'static_snat_count' or \
                   stat.key == 'static_dnat_count') and \
                   'Spectrum' in system_info:
                    continue
                if (stat.key == 'config_snat_count' or \
                   stat.key == 'config_dnat_count') and \
                   'BCM' in system_info:
                    continue
            else:
                continue
        if (stat.att == 'mlx' and \
           'BCM' in system_info):
               continue

        if (stat.att == 'bcm' and \
           'Spectrum' in system_info):
               continue
        count_files = list()
        if type(stat.count_file) is list:
            for file in stat.count_file:
                count_files.append(stat_path + file)
        elif stat.count_file:
            count_file = stat_path + stat.count_file
        max_file = stat_path + stat.max_file if stat.max_file else None
        alloc_file = stat_path + stat.alloc_file if stat.alloc_file else None
        percent_occupancy = 0
        rule_class_list = []
        rule_count_list = []

        if stat.rule_class_list:
            for _ in stat.rule_class_list:
                rule_class_list.append(stat_path + _)

        if ((stat.key == 'in_acl_entry' or stat.key == 'eg_acl_entry') and \
            'Spectrum' in system_info):
            max_file = None
            alloc_file = None

        count = 0
        if (stat.count_function):
            count = stat.count_function(stat.func_arg1, stat.func_arg2)
            stats_dict[stat.key] = {}
            stats_dict[stat.key]["count"] = count

            if stat.key == "total_fid_count":
                max_value = MLX_MAX_FID_COUNT
        elif type(stat.count_file) is list:
            for item in count_files:
                try:
                    with open(item, 'r') as f:
                        count += int(f.readline())
                except IOError:
                    print('%s count is not available' % stat.name)
                    return -1
        else:
            try:
                with open(count_file, 'r') as f:
                    if read_as_string is True:
                        str_val = f.readline().rstrip()
                        if stat.key == 'forwarding_profile':
                            forwarding_profile = str_val

                    else:
                        count = int(f.readline())
            except IOError:
                print('%s count is not available' % stat.name)
                return -1

        if max_file:
            try:
                with open(max_file, 'r') as f:
                    max_value = int(f.readline())
                    #Max value being -1 indicates that this resource
                    #is not available, as per switchd
                    #This check makes sure that non availability of
                    #resource should be displayed as 0
                    if max_value == -1:
                        max_value = 0

            except IOError:
                print('%s max value is not available' % stat.name)
                return -1

        if float(max_value) > 0:
            percent_occupancy = float(count) / float(max_value)

        if alloc_file:
            try:
                with open(alloc_file, 'r') as f:
                    alloc_value = int(f.readline())
            except IOError:
                print('%s alloc value is not available' % stat.name)
                return -1
        else:
            alloc_value = 0

        if len(rule_class_list) != 0:
            for fh in rule_class_list:
                try:
                    with open(fh,'r') as f:
                        t = int(f.readline())
                        rule_count_list.append(t)
                except IOError:
                    print('%s rule_count is not available' % fh)
                    return -1

        if key_value_flag:
            if read_as_string is False:
                print(stat.key + '_count=' + '%d' % count)
            else:
                print(stat.key + '=' + str_val)
            if max_file:
                print(stat.key + '_max=' + '%d' % max_value)
            if alloc_file:
                print(stat.key + '_allocated=' + '%d' % alloc_value)

        elif use_json:
            if stat.rule_class_list:
                if read_as_string is False:
                    stats_dict[stat.key] = {
                        'count': int(count),
                        'max' : 0,
                        'name': stat.name,
                        '18B' : rule_count_list[0],
                        '36B' : rule_count_list[1],
                        '54B' : rule_count_list[2],
                }
                else:
                    stats_dict[stat.key] = {
                        'val':  str_val,
                        'max' : 0,
                        'name': stat.name,
                        '18B' : rule_count_list[0],
                        '36B' : rule_count_list[1],
                        '54B' : rule_count_list[2],
                    }
            else:
                if read_as_string is False:
                    stats_dict[stat.key] = {
                    'count': int(count),
                    'max': int(max_value),
                    'allocated': int(alloc_value),
                    'name': stat.name,
                    'percentage': float(percent_occupancy * 100),
                }
                else:
                    stats_dict[stat.key] = {
                        'val': str_val,
                        'max': int(max_value),
                        'allocated': int(alloc_value),
                        'name': stat.name,
                        'percentage': float(percent_occupancy * 100),
                    }

        else:
            if max_file or max_value:
                if alloc_file:
                        output = '%-34s %6d, %3d%% of maximum value %6d (allocated: %d)' % (stat.name + ':',
                                                                            count,
                                                                            int(percent_occupancy * 100),
                                                                            max_value, alloc_value)
                else:
                        output = '%-34s %6d, %3d%% of maximum value %6d' % (stat.name + ':',
                                                                    count,
                                                                    int(percent_occupancy * 100),
                                                                    max_value)
                if percent_occupancy > 0.90:
                    color = bcolors.FAIL
                elif percent_occupancy > 0.75:
                    color = bcolors.WARNING
                else:
                    color = None
                if color:
                    print(color + output + bcolors.ENDC)
                else:
                    print(output)
            else:
                if len(rule_count_list) != 0:
                    print('%-34s %6d %6s :%2d 36B :%2d 54B :%2d ' % (stat.name + ':', count,'18B',
                                                                  rule_count_list[0],
                                                                  rule_count_list[1],
                                                                  rule_count_list[2],))
                else:
                    if read_as_string is False:
                        print('%-34s %6d' % (stat.name + ':', count))
                    else:
                        print('%-38s %s' % (stat.name + ':', str_val))

    if use_json:
        print(json.dumps(stats_dict, sort_keys=True, indent=4))


def main(argv):
    """ main function """
    descr = 'Cumulus Resource Reporting'

    arg_parser = argparse.ArgumentParser(description=descr)
    # Command line arg parser
    #
    option = arg_parser.add_mutually_exclusive_group(required=False)
    option.add_argument('-k', '--key-value', dest='key_value',
                        action='store_true',
                        default=False,
                        help='Report key=value pairs')
    option.add_argument('-j', '--json',
                        action='store_true',
                        default=False,
                        help='Display JSON output')

    # Parse command line arguments
    cmdline_args = arg_parser.parse_args()

    ret = generate_report(cmdline_args.key_value, cmdline_args.json)

if __name__ == "__main__":
    main(sys.argv[1:])
