#
# Copyright (c) 2013-2024 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#  This software is available to you under a choice of one of two
#  licenses.  You may choose to be licensed under the terms of the GNU
#  General Public License (GPL) Version 2, available from the file
#  COPYING in the main directory of this source tree, or the
#  OpenIB.org BSD license below:
#
#      Redistribution and use in source and binary forms, with or
#      without modification, are permitted provided that the following
#      conditions are met:
#
#       - Redistributions of source code must retain the above
#         copyright notice, this list of conditions and the following
#         disclaimer.
#
#       - Redistributions in binary form must reproduce the above
#         copyright notice, this list of conditions and the following
#         disclaimer in the documentation and/or other materials
#         provided with the distribution.
#
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
#  MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
#  BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
#  ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
#  CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
#  SOFTWARE.
from __future__ import print_function
import sys
import os
import subprocess
import re
from collections import namedtuple
import argparse
import shutil

MFT_DEVS_DIR = "/dev/mst/"
RETIMER_DEVS_DIR = "/dev/mst/retimer"
DEV_ID_CRSPACE_ADDRESS = 0xf0014

CMDS = ["add",
        "del",
        "status"]

# Append new retimers to this dictionary
RETIMERS_DIC = {
    "ArcusE": {
        "devid": 0x282,
        "i2c_secondary_address": 0x48,
        "address_width": 4,
        "data_len": 4,
        "devid_address": 0xf0014
    },
    "Arcus2": {
        "devid": 0x288,
        "i2c_secondary_address": 0x48,
        "address_width": 4,
        "data_len": 4,
        "devid_address": 0xf0014
    },
    # "ArcusE": {
    #     "devid": 110,
    #     "i2c_secondary_address": 0x48,
    #     "address_width": 4,
    #     "data_len": 4,
    #     "devid_address": 0xf0014
    # }
}


class Command(object):
    def __init__(self, cmd_str):
        self.cmd_str = cmd_str

    def execute(self):
        p = subprocess.Popen(self.cmd_str,
                             stdin=subprocess.PIPE,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE,
                             universal_newlines=False,
                             shell=True)
        stdout, stderr = p.communicate()
        stat = p.wait()
        return stat, stdout, stderr


class I2cError(Exception):
    pass


def is_remote_device(dev):
    return ":" in dev or "#" in dev


class Retimer(object):

    @classmethod
    def discover_external_chips(cls, cables=None):
        devices_to_check = cables if cables else [dev for dev in os.listdir(MFT_DEVS_DIR) if not is_remote_device(dev)]
        devices_list = []

        devices_list = cls._discover_chips_in_cables(devices_to_check)
        # retimers_in_evb = cls._discover_chip_on_EVB(mtusb_devices)
        # devices_list.extend(retimers_in_evb)
        retimers_num = cls._create_files(devices_list)
        return retimers_num

    @staticmethod
    def convert(in_str):
        result = []
        if in_str is not None and len(in_str):
            for token in in_str.split("\n"):
                if len(token):
                    try:
                        result.append((int(token.split(",")[0]), int(token.split(",")[1])))
                    except BaseException:
                        pass
        return result

    @classmethod
    def _retimers_ids(cls):
        ids_list = []
        for dev_info in RETIMERS_DIC.values():
            ids_list.append(dev_info["devid"])
        return ids_list

    @classmethod
    def _create_files(cls, devices_list):
        retimer_devs = 0
        if len(devices_list) > 0:
            # crate folder if not exist
            cls._create_dir(RETIMER_DEVS_DIR)
            for device in devices_list:
                # create new file
                dev_file = os.path.join(RETIMER_DEVS_DIR, device)
                if not os.path.exists(dev_file):
                    open(dev_file, 'w').close()
                    retimer_devs += 1
        return retimer_devs

    @classmethod
    def _discover_chips_in_cables(cls, mst_devices):
        """
        Discover retimers in a module via MCU (FW Gateway)
        The function returns a list with the name of discovered retimer devices
        """
        cable_devices = []
        linkx_devices = []
        for mst_device in mst_devices:
            if 'cable' in mst_device:
                cable_devices.append(mst_device)
            # if 'mtusb' in mst_device and 'cable' not in mst_device:
            #     mtusb_devices.append(MFT_DEVS_DIR + mst_device)

        for cable_device in cable_devices:
            cmd = "mlxcables -d {0} --discover".format(cable_device)
            rc, out, _ = Command(cmd).execute()
            if not rc:
                result = Retimer.convert(out.decode("utf-8"))
                for device_id, num_of_devices in result:
                    if device_id in cls._retimers_ids():
                        if num_of_devices == 1:
                            linkx_device_name = "{0}_rt{1}".format(cable_device, device_id)
                            linkx_devices.append(linkx_device_name)
                        else:  # num_of_devices > 1:
                            for device_num in range(num_of_devices):
                                linkx_device_name = "{0}_rt{1}_{2}".format(cable_device, device_id, device_num)
                                linkx_devices.append(linkx_device_name)
        return linkx_devices

    @classmethod
    def _discover_chip_on_EVB(cls, mtusb_devices):
        retimer_devs = []
        retimer_folder_exist = False
        for mtusb_dev in mtusb_devices:  # for each mtusb device found in mst status
            for dev_info in RETIMERS_DIC.values():  # check if retimer by reading i2c dev id
                ret = cls._i2c_cmd_read(mtusb_dev, dev_info["i2c_secondary_address"], dev_info["address_width"],
                                        DEV_ID_CRSPACE_ADDRESS, dev_info["data_len"])
                if ret != -1 and int(ret, 16) == dev_info["devid"]:
                    retimer_dev_name = "{}_rt{}".format(mtusb_dev.replace("/dev/mst/", ""), dev_info["devid"])
                    # # create retimer dirif not exist
                    # if not retimer_folder_exist:
                    #     cls._create_dir(RETIMER_DEVS_DIR)
                    #     retimer_folder_exist = True
                    # #create new file+link
                    # dev_file = os.path.join(RETIMER_DEVS_DIR, retimer_dev_name)
                    # if not os.path.exists(dev_file):
                    #     os.symlink(mtusb_dev, dev_file)
                    #     retimer_devs += 1
                    # break
                    dev_file = os.path.join(RETIMER_DEVS_DIR, retimer_dev_name)
                    retimer_devs.append(dev_file)
        return retimer_devs

    @classmethod
    def _i2c_cmd_read(cls, mtusb_device, i2c_secondary_address, address_width, address, size):
        cmd = "i2c -a {0} -x {1} {2} read {3:#x} {4:#x}".format(
            address_width, size, mtusb_device, i2c_secondary_address, address)
        rc, out, err = Command(cmd).execute()
        if (rc == 0):
            return out
        else:
            return -1

    @classmethod
    def _create_dir(cls, dir_name):
        if not os.path.isdir(dir_name):
            try:
                os.makedirs(dir_name)
            except OSError:
                print("-E- Creation of the directory %s failed" % dir_name)


if __name__ == '__main__':
    try:
        parser = argparse.ArgumentParser(prog="mst_retimer",
                                         description="Retimers discovery script")
        parser.add_argument("command", choices=CMDS)
        parser.add_argument("cables", nargs='?', default=None)
        args = parser.parse_args()

        if args.command == "del":
            if os.path.exists(RETIMER_DEVS_DIR):
                shutil.rmtree(RETIMER_DEVS_DIR,
                              ignore_errors=False, onerror=None)

        elif args.command == "add":
            cables = None
            if args.cables:
                # convert to set and back to list to get rid of duplicates
                cables = list(set(args.cables.split(',')))
            chipset_cnt = Retimer.discover_external_chips(cables=cables)
            if chipset_cnt:
                print("-I- Added {} retimer device{} ..".format(chipset_cnt, "s" if chipset_cnt > 1 else ""))
            else:
                print("-I- No Retimer devices found")

        elif args.command == "status":
            if os.path.exists(RETIMER_DEVS_DIR):
                print("Retimer devices:")
                print("----------------")
                for file in os.listdir(RETIMER_DEVS_DIR):
                    print(file)

    except Exception as err:
        print(err, file=sys.stderr)
        sys.exit(1)
