import os
import subprocess
import configparser
import re
import json
import sys

mrzcpd_root = "/opt/tsg/mrzcpd"
mrzcpd_cfg = ""
dpdk_bind = ""
dpdk_driver = "vfio-pci"
interfaces = []
hwfile_location = '/var/run/mrzcpd/hwfile.json'
hwinfo = {}


def init():
    global mrzcpd_root
    global mrzcpd_cfg
    global dpdk_bind
    global dpdk_driver
    global interfaces
    global hwinfo

    # Get driver mode
    with open("/etc/sysconfig/mrzcpd", 'r') as drive_fp:
        for line in drive_fp.readlines():
            if re.search("DEFAULT_UIO_MODULE=*", line):
                str_driver = line.split('DEFAULT_UIO_MODULE=')[1]
                match = re.search(r'"(.*?)"', str_driver)
                if match:
                    dpdk_driver = match.group(1)
                    if match.group(1) == "vfio_pci".strip():
                        dpdk_driver = "vfio-pci"
            if re.search("MRZCPD_ROOT=*", line):
                mrzcpd_root = line.split('MRZCPD_ROOT=')[1]

    # Get cfg path
    mrzcpd_cfg = mrzcpd_root.replace("\n", "") + "/etc/mrglobal.conf"
    dpdk_bind = mrzcpd_root.replace("\n", "") + "/bin/dpdk-devbind.py"

    # Check driver mode
    if dpdk_driver.strip() not in {"vfio-pci", "uio_pci_generic", "igb_uio"}:
        print("driver mode invalid.", dpdk_driver)
        sys.exit(1)

    # Get dev list from mrglobal
    config = configparser.ConfigParser()
    config.read(mrzcpd_cfg)
    devices = config.get('device', 'device').split(',')
    for d in devices:
        if config.get('device:' + d,'driver') != '2':
            interfaces.append(d)

    # hwfile decode
    with open(hwfile_location, 'r') as json_fp:
        json_info = json.load(json_fp)
        for item in json_info:
            dev_info = {}
            dev_name = item.get('Interface', None)
            dev_info['pci'] = item.get('Slot_str', None)
            dev_info['driver'] = item.get('Driver_str', None)
            hwinfo[dev_name] = dev_info


def module_load():
    # vfio_pci
    if dpdk_driver.strip() == "vfio-pci".strip():
        if os.system("modprobe vfio") or os.system("echo 1 > /sys/module/vfio/parameters/enable_unsafe_noiommu_mode") or os.system("modprobe vfio_pci"):
            print("modprode vfio err")
            sys.exit(1)

    # uio_pci_generic
    if dpdk_driver.strip() == "uio_pci_generic".strip():
        if os.system("modprobe uio_pci_generic"):
            print("modprode uio_pci_generic err")
            sys.exit(1)

    # igb_uio
    if dpdk_driver.strip() == "igb_uio".strip():
        if os.system("modprobe igb_uio"):
            print("modprode igb_uio err")
            sys.exit(1)


def nic_bind():
    for d in interfaces:
        if d in hwinfo and hwinfo.get(d).get('driver') != 'mlx5_core':
            bind_cmd = dpdk_bind + ' --bind=' + \
                dpdk_driver + " " + hwinfo.get(d).get('pci')
            print("bind cmd:", bind_cmd)
            if os.system(bind_cmd):
                print("bind dev err")
                sys.exit(1)


def nic_unbind():
    for d in interfaces:
        if d in hwinfo and hwinfo.get(d).get('driver') != 'mlx5_core':
            dpdk_driver_unbind = hwinfo.get(d).get('driver')
            unbind_cmd = dpdk_bind + ' --bind=' + \
                dpdk_driver_unbind + " " + hwinfo.get(d).get('pci')
            print("unbind cmd:", unbind_cmd)
            if os.system(unbind_cmd):
                print("unbind dev err")
                sys.exit(1)


def main():
    args = sys.argv
    if args[1] == "bind":
        # module load
        module_load()
        # bind dev
        nic_bind()
    elif args[1] == "bind-without-module-load":
        # bind dev
        nic_bind()
    elif args[1] == "unbind":
        # unbind dev
        nic_unbind()


if __name__ == "__main__":
    init()
    main()