#!/usr/bin/python
#
# Copyright 2017.  Cumulus Networks, Inc.
#
# portmap
#
# Utility to dump the ASIC portmap.
#
# See usage (portmap -h) for instructions

import argparse
import json
import os
import sys

import cumulus.platforms
import cumulus.porttab

class ParseError(RuntimeError):
    pass

class BcmAsicPortMapEntry:
    def __init__(self, linux_intf, sdk_intf, unit, is_fabric, logical_port):
        self.linux_intf = linux_intf
        self.sdk_intf = sdk_intf
        self.unit = unit
        self.is_fabric = is_fabric
        self.logical_port = logical_port

    def get_json_obj(self):
        pInfo = {}
        pInfo['linux_intf'] = self.linux_intf
        pInfo['sdk_intf'] = self.sdk_intf
        pInfo['unit'] = self.unit
        pInfo['is_fabric'] = self.is_fabric
        pInfo['logical_port'] = self.logical_port
        return pInfo

class BcmAsic:
    def __init__(self, json_out):
        self.json = json_out
        self.pt = cumulus.porttab.porttab()
        self.bcmPm = []
        # if more than one unit is present we need include unit id in the key
        self.units = []
        for pe in self.pt.porttab:
            # replace the sdk_intf in the porttab with the one from 'bcm portmap'
            if pe.logical_port != -1:
                pe.bcm =  self.pt.logical2dport(pe.logical_port)

            pe.unit = int(pe.unit)
            if pe.unit not in self.units:
                self.units.append(pe.unit)
            bcmPe = BcmAsicPortMapEntry(pe.linux_intf, pe.bcm,
                            pe.unit, pe.is_fabric, pe.logical_port)
            self.bcmPm.append(bcmPe)

    def __str__(self):
        if self.json:
            jout = []
            for bcmPe in self.bcmPm:
                pInfo = bcmPe.get_json_obj()
                jout.append(pInfo)
            return json.dumps(jout, indent=4) + '\n'
        else:
            # for now i am keeping the same format as porttab; more fields can
            # be added later
            portmap = ''
            portmap += '# linux_intf  sdk_intf  unit  is_fabric  logical_port\n'
            for bcmPe in self.bcmPm:
                 portmap += ('%-12s  %8s  %4s  %9s  %12u\n' %
                            (bcmPe.linux_intf, bcmPe.sdk_intf, bcmPe.unit,
                            bcmPe.is_fabric, bcmPe.logical_port))
            return portmap
    
    def dump_by_logical_port(self, lp):
        for bcmPe in self.bcmPm:
            if bcmPe.logical_port == lp:
                if self.json:
                    portmap = json.dumps(bcmPe.get_json_obj(), indent=4)
                else:
                    portmap = bcmPe.linux_intf
                sys.stdout.write(portmap + '\n')
                return 0
        return 1
        

    def dump_by_linux_intf(self, linux_intf):
        for bcmPe in self.bcmPm:
            if bcmPe.linux_intf == linux_intf:
                if self.json:
                    portmap = json.dumps(bcmPe.get_json_obj(), indent=4)
                else:
                    # if more than 1 unit use unit:xe format
                    portmap = ('%d:%s' % (bcmPe.unit, bcmPe.sdk_intf)) if (len(self.units) > 1) else bcmPe.sdk_intf
                sys.stdout.write(portmap + '\n')
                return 0
        return 1

    def dump_by_sdk_intf(self, unit, sdk_intf):
        for bcmPe in self.bcmPm:
            if bcmPe.unit == unit and bcmPe.sdk_intf == sdk_intf:
                if self.json:
                    portmap = json.dumps(bcmPe.get_json_obj(), indent=4)
                else:
                    portmap = bcmPe.linux_intf
                sys.stdout.write(portmap + '\n')
                return 0
        return 1

    def dump_intf(self, intf):
        #if intf is int treat it as logical port
        if intf.isdigit():
            return self.dump_by_logical_port(int(intf))

        #if intf has xe, ge, ce treat it as dport name
        if 'xe' in intf or 'ge' in intf or 'ce' in intf:
            unit = 0
            if ':' in intf:
                intfInfo = intf.split(':', 1)
                try:
                    unit = int(intfInfo[0])
                    intf = intfInfo[1]
                except ValueError:
                    pass
            return self.dump_by_sdk_intf(unit, intf)

        #otherwise linux_intf
        return self.dump_by_linux_intf(intf)


if __name__ == '__main__':
    if os.getuid():
        sys.stderr.write("need to run as root\n")
        sys.exit(1)

    parser = argparse.ArgumentParser(
        description='Provide the linux to ASIC portmap information')
    parser.add_argument('-v', '--verbose', default=False, action='store_true',
                        help='Verbose output')
    parser.add_argument('-j', '--json', default=False, action='store_true',
                        help='JSON output')
    parser.add_argument('intf', nargs='?', default='',
                        help='intf_name. e.g. swp1 or xe0 or 0:xe0 (unit:dport) or 50 (logical_port).')


    try:
        args = parser.parse_args()
    except ParseError, e:
        parser.error(str(e))

    platform_object = cumulus.platforms.probe()
    # create the ASIC object
    asic = None
    if platform_object.switch:
        chip = platform_object.switch.chip
        # currently only used fro Broadcom chips; may extend to Mellanox in the
        # future
        if chip.sw_base == 'bcm':
            asic = BcmAsic(args.json)

    if not asic:
        if json:
            print json.dumps({}, indent=4)
        else:
            sys.stdout.write('Unsupported switchd chip\n')
        sys.exit(1)

    rv = 0
    if args.intf:
        rv = asic.dump_intf(args.intf) 
        sys.exit(rv)
        
    sys.stdout.write('%s' % asic)
    sys.exit(0)
