#!/usr/bin/env python

import sys
import glob
import json
import argparse

sys.path.append('thrift/gen-py')

from pltfm_pm_rpc.pltfm_pm_rpc import Client
from pltfm_pm_rpc.ttypes import InvalidPltfmPmOperation
from pltfm_pm_rpc.ttypes import pltfm_pm_port_speed_t, pltfm_pm_fec_type_t
from pltfm_pm_rpc.ttypes import pltfm_pm_board_type_t
from thrift import Thrift
from thrift.transport import TSocket
from thrift.transport import TTransport
from thrift.protocol import TBinaryProtocol
from thrift.protocol import TMultiplexedProtocol


class PortMgr:
    """ Port Manager """

    def __init__(self, serverIp, serverPort=9095, portsDef={}):
        self.serverIp = serverIp
        self.serverPort = serverPort
        self.portsDef = portsDef
        self.broadType = None
        self.transport = TSocket.TSocket(serverIp, serverPort)
        self.transport = TTransport.TBufferedTransport(self.transport)
        self.protocol = TBinaryProtocol.TBinaryProtocol(self.transport)
        self.protocol = TMultiplexedProtocol.TMultiplexedProtocol(self.protocol, "pltfm_pm_rpc")
        self.client = Client(self.protocol)

    def __enter__(self):
        self.transport.open()
        bd_id = self.getBroadType()
        self.broadType = pltfm_pm_board_type_t._VALUES_TO_NAMES[bd_id]
        self.log("Connected, broad type: %s" % (self.broadType))
        return self

    def __exit__(self, type, value, traceback):
        self.transport.close()
        self.log("Disconnected")

    def log(self, msg):
        print "[%s] %s" % (self.serverIp, msg)

    def portCleanup(self, deviceId=0):
        self.log("Cleaning up all ports")
        # clean up is not working.....
        # return self.client.pltfm_pm_switchd_port_cleanup(deviceId)
        # try disable all ports and delete all ports
        for portDef in self.portsDef[self.broadType]:
            for dp in portDef:
                try:
                    if self._getPortStatus(dp, deviceId) != -1:
                        self.client.pltfm_pm_port_dis(deviceId, dp)
                        self.client.pltfm_pm_port_del(deviceId, dp)
                except Exception as e:
                    raise e

    def getPortStatus(self, portNum, channelNum=0, deviceId=0):
        dp = self.portsDef[self.broadType][portNum][channelNum]
        return self._getPortStatus(dp, deviceId)

    def _getPortStatus(self, dp, deviceId=0):
        try:
            result = self.client.pltfm_pm_port_oper_status_get(deviceId, dp)
            return result
        except InvalidPltfmPmOperation as e:
            return -1

    def addPort(self, portNum, channelNum=0, deviceId=0,
                speed=pltfm_pm_port_speed_t.BF_SPEED_100G,
                fecType=pltfm_pm_fec_type_t.BF_FEC_TYP_NONE):
        self.log("Adding port %d/%d with speed %s fec %s"
                 % (portNum, channelNum, speed, fecType))
        dp = self.portsDef[self.broadType][portNum][channelNum]
        try:
            result = self.client.pltfm_pm_port_add(deviceId, dp, speed, fecType)
            self.log("Port %d/%d added" % (portNum, channelNum))
            return result
        except InvalidPltfmPmOperation as e:
            self.log("Can't add port %d/%d, code: %d" % (portNum, channelNum, e.code))

    def enablePort(self, portNum, channelNum=0, deviceId=0):
        self.log("Enabling port %d/%d" % (portNum, channelNum))
        dp = self.portsDef[self.broadType][portNum][channelNum]
        try:
            result = self.client.pltfm_pm_port_enable(deviceId, dp)
            self.log("Port %d/%d enabled" % (portNum, channelNum))
            return result
        except Exception as e:
            self.log("Can't enable port %d/%d, code: %d" % (portNum, channelNum, e.code))

    def getBroadType(self):
        return self.client.pltfm_pm_board_type_get();


def main(args):
    with open(args.ports_file) as portsFile:
        portsDef = json.load(portsFile)
    with PortMgr(args.ip, args.port, portsDef) as portMgr:
        if args.do_cleanup:
            portMgr.portCleanup()
        with open(args.port_config) as portConfig:
            for line in portConfig:
                if len(line) == 0:
                    continue
                line = line.strip()
                info = line.split(' ')
                portChannel = info[0]
                speed = pltfm_pm_port_speed_t._NAMES_TO_VALUES[info[1]]
                fecType = pltfm_pm_fec_type_t._NAMES_TO_VALUES[info[2]]
                portNum = int(portChannel.split('/')[0])
                channelNum = int(portChannel.split('/')[1])
                portMgr.addPort(portNum, channelNum, speed=speed, fecType=fecType)
                portMgr.enablePort(portNum, channelNum)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Tofino port manager tool')
    parser.add_argument('ip', help='IP of BF-Platfrom')
    parser.add_argument('port_config', help='Port configuration')
    parser.add_argument('--port', type=int, help='Port of BF-Platfrom', default=9095)
    parser.add_argument('--ports-file', help='A json file which defines ports',
                        default='ports.json')
    parser.add_argument('--do-cleanup', help='Performs port cleanup',
                        type=bool, default=False)

    args = parser.parse_args()
    main(args)