# --
# Copyright (c) 2021 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.
# --

#!/usr/bin/python
### imports ###
import os
import sys
import subprocess
import argparse


def buildTree(linesList, tabIndex, nodeSons):
    if len(linesList) == 0:
        return ("", nodeSons)
    nodeStr = linesList[0]
    if not isNodeStr(nodeStr):
        return buildTree(linesList[1:], tabIndex, nodeSons)
    nodeTab = getNodeTab(nodeStr)
    if nodeTab == tabIndex:
        return (nodeSons, linesList)
    if (tabIndex - nodeTab) == 1:
        nodeName = getNodeName(nodeStr)
        remineLines, nodeSons = buildTree(linesList[1:], nodeTab, [])
        nodeStruct = {'nodeName': nodeName, 'nodeSons': nodeSons}
        return buildTree(linesList[1:], nodeTab, nodeSons.add(nodeStruct))


def isNewNode(line, queryCommand):
    if queryCommand:
        if ":" in line:
            index = line.find(":")
            charBF = line[index - 1]
            return (charBF == "" or charBF == " ")
    removeLeadingZero = line.strip(" ")
    if removeLeadingZero and removeLeadingZero[0] == "-":
        return True
    return False


def isInfoFlag(line):
    if line[0] != " " and line[0] != "-":
        return True
    return False


def keepOnlyNodes(filesList):
    helpLineOnlyNodes = []
    onNodedescription = False
    queryCommand = False
    node = ""
    for line in filesList:
        if line.strip(" ") == "" or line == '\n':
            continue
        if "Examples" in line:
            helpLineOnlyNodes.append(node)
            return helpLineOnlyNodes
        if "COMMANDS SUMMARY" in line:
            queryCommand = True
            continue
        if isInfoFlag(line):
            queryCommand = False
            onNodedescription = False
            continue
        if isNewNode(line, queryCommand):
            if node != "":
                helpLineOnlyNodes.append(node)
            onNodedescription = True
            node = line
            if ":" not in line:
                node = node + " :"
            continue
        if onNodedescription:
            node = node + " " + line.strip(" ")
    helpLineOnlyNodes.append(node)
    return helpLineOnlyNodes


# -l |--loopback   <loopback>             : Configure Loopback Mode [NO(no loopback)/PH(phy loopback)/EX(external loopback)]

def isShortCutsLine(line):
    if "|" in line:
        lineSplit = line.split("|")[0].strip(" ")
        lineSpaceSplit = lineSplit.split(" ")
        lineSpaceWords = list(filter(lambda x: x != "", lineSpaceSplit))
        return (len(lineSpaceWords) <= 1)
    return False


def getNodeShortCutsAndName(line):
    shortCuts = ""
    name = ""
    line = line.strip(" ")
    if isShortCutsLine(line):
        lineSplit = line.split("|")
        rightFlag = lineSplit[0].strip(" ")
        leftFlag = lineSplit[1].strip(" ").split(" ")[0]
        if len(rightFlag) > len(leftFlag):
            return leftFlag, rightFlag
        else:
            return rightFlag, leftFlag
    name = line.split(" ")[0]
    return shortCuts, name


def getSonsAndUpperNeededFromVals(sonsValsSplit):
    upperNeed = "1"
    sons = ""
    lastVal = ""
    if "..." in sonsValsSplit:
        threePointsIndex = sonsValsSplit.index("...")
        if threePointsIndex < len(sonsValsSplit) - 1:
            numBefore = int(sonsValsSplit[threePointsIndex - 1]) + 1
            numAfter = int(sonsValsSplit[threePointsIndex + 1])
            arr = range(numBefore, numAfter)
            sonsValsSplit[threePointsIndex] = " ".join(map(lambda x: str(x), arr))
    for val in sonsValsSplit:
        sonVal = val.split('(')[0]
        sons = sons + " " + sonVal
        if not sonVal.isupper():
            upperNeed = ""
    return sons, upperNeed


def getSonsValsFromLine(line):
    ret = ""
    lineSplit = line.split("[")[1:]
    for sons in lineSplit:
        ret = ret + sons.split("]")[0]
    return ret


def getNodeTypeSonsUpperNeedAndLastCommandIndex(line):
    lastCommandIndex = "-1"
    sons = ""
    nodeType = "0"
    upperNeed = ""
    if "[" in line:
        lastCommandIndex = "1"
        sonsVals = getSonsValsFromLine(line)
        sonsValsSplit = []
        if "/" in sonsVals:
            nodeType = "3"
            sonsValsSplit = sonsVals.split("/")
        elif "," in sonsVals:
            nodeType = "4"
            sonsValsSplit = sonsVals.split(",")
        else:
            nodeType = "3"
            sonsValsSplit = [sonsVals]
        sons, upperNeed = getSonsAndUpperNeededFromVals(sonsValsSplit)
    return nodeType, sons, upperNeed, lastCommandIndex


def parseLine(line):
    node = {'shortCut': "", 'name': "", 'lastCommandIndex': "", 'nodeType': "", 'sons': "", 'extra': "", 'upperNeed': "", 'description': ""}
    lineSplitByColoumn = line.split(":")
    shortCut, name = getNodeShortCutsAndName(lineSplitByColoumn[0])
    node["shortCut"] = " ".join(shortCut.split())
    node["name"] = " ".join(name.split())
    node["description"] = lineSplitByColoumn[1]
    nodeType, sons, upperNeed, lastCommandIndex = getNodeTypeSonsUpperNeedAndLastCommandIndex(line)
    node["nodeType"] = nodeType
    node["sons"] = sons
    node["upperNeed"] = upperNeed
    node["lastCommandIndex"] = lastCommandIndex
    return node


def getNumOfLeadingZeroFromLine(line):
    count = 0
    for char in line:
        if char == " ":
            count += 1
        elif char == "\t":
            count += 4
        else:
            return count
    return count


def getTabLevelFromLine(line):
    numOfLeadingZeros = getNumOfLeadingZeroFromLine(line)
    return int(numOfLeadingZeros / 4)


def getUserInput():
    my_parser = argparse.ArgumentParser()
    my_parser.add_argument('-n', '--name', action='store', type=str, help="Tool Name", required=True)
    my_parser.add_argument('-o', '--output', action='store', type=int, default=1, choices=[1], help="Output: 1- build structs for Auto-Compeletion(Default)")
    my_parser.add_argument('-c', '--command', action='store', type=str, default="", help="Full Command Line: ")

    args = my_parser.parse_args()
    return vars(args)


def isCommandNode(node):
    return node["name"][0] != "-"


def isNeedUpdateNode(prevNode, prevIndex, currNode, CurrIndex):
    if prevNode is None:
        return True
    if isCommandNode(prevNode):
        if isCommandNode(currNode):
            return prevIndex > CurrIndex
        else:
            return False
    else:
        return CurrIndex > prevIndex


def buildMSTNodesTree(commandLine):
    nodesTree = []
    filesList = subprocess.check_output("sudo mst help", shell=True).decode().split("\n")
    options = ""
    params = ""
    sons = ""
    extra = ""
    lastNodeStruct = {'shortCut': "", 'name': "", 'lastCommandIndex': "1", 'nodeType': "0", 'sons': "", 'extra': "", 'upperNeed': "", 'description': ""}
    mstMainNodeSons = ""
    for line in reversed(filesList):
        lineRemovingLeadingZero = line.strip(" ")
        if len(lineRemovingLeadingZero) == 0:
            continue
        if lineRemovingLeadingZero[0:3] == "mst":
            splitLine = lineRemovingLeadingZero[4:].split(" ")
            nodeName = splitLine.pop(0)
            if nodeName not in mstMainNodeSons.split(" "):
                mstMainNodeSons = mstMainNodeSons + " " + nodeName
            if nodeName not in commandLine.split(" "):
                continue
            if lastNodeStruct["name"] != nodeName:
                if lastNodeStruct["name"] != "" and lastNodeStruct["sons"].strip(" ") != "":
                    nodesTree.append(lastNodeStruct)
                lastNodeStruct = {'shortCut': "", 'name': nodeName, 'lastCommandIndex': "1", 'nodeType': "0", 'sons': "", 'extra': "", 'upperNeed': "", 'description': ""}
            for son in reversed(splitLine):
                if son[0] == "<":
                    continue
                if son[0] == "[":
                    sonName = son.replace("[", "").replace("]", "")
                    if sonName == "OPTIONS":
                        extra = extra + options + sons
                        sons = ""
                    elif sonName == "params":
                        extra = params
                    else:
                        sons = sons + " " + sonName
                else:
                    if sons == "":
                        if extra != "":
                            sons = extra
                            extra = ""
                    if sons != "":
                        nodesTree.append({'shortCut': "", 'name': son, 'lastCommandIndex': "1", 'nodeType': "0", 'sons': sons, 'extra': extra, 'upperNeed': "", 'description': ""})
                    sons = son
                    params = ""
                    options = ""
            lastNodeStruct["sons"] = lastNodeStruct["sons"] + " " + sons
            sons = ""
            params = ""
            options = ""
        elif "params:" in line:
            paramsSplit = lineRemovingLeadingZero.split(" ")
            for param in paramsSplit:
                if param[0] == "[":
                    paramMame = param.replace("[", "").replace("]", "")
                    params = params + " " + paramMame
        elif lineRemovingLeadingZero[0] == "-":
            options = options + " " + lineRemovingLeadingZero.split(" ")[0].replace(":", "")
    nodesTree.append(lastNodeStruct)
    nodesTree.append({'shortCut': "", 'name': "mst", 'lastCommandIndex': "1", 'nodeType': "0", 'sons': mstMainNodeSons, 'extra': "", 'upperNeed': "", 'description': ""})
    return nodesTree


def buildMFTTollNodesTree(toolName, commandLine):
    nodesTree = []
    nodesTabs = ["", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", "", ""]
    try:
        # force decoding to unicode (utf-8) because python2 decodes to ascii and creates and error
        filesList = subprocess.check_output(toolName + " -h", shell=True).decode('utf-8').split("\n")
    # some tools supply --help only when executed by root. at this point auto-complete can not support them
    except subprocess.CalledProcessError:  # this error exists in both python2 and python3
        sys.exit(0)

    helpLineOnlyNodes = keepOnlyNodes(filesList)
    tabLevel = 0
    tabLevelBestIndex = [0] * 8
    tabLevelBestNode = [None] * 8
    for line in reversed(helpLineOnlyNodes):
        parsedLine = parseLine(line)
        tabLevel = getTabLevelFromLine(line)
        if tabLevel > 2:  # The current line belongs to the previous command.
            continue
        nodesTabs[tabLevel] = nodesTabs[tabLevel] + " " + parsedLine["name"]
        sons = nodesTabs[tabLevel + 1]
        nodesTabs[tabLevel + 1] = ""
        if parsedLine["sons"] == "":
            parsedLine["sons"] = sons
        else:
            parsedLine["extra"] = sons
        flagIndex = -1
        splitCommand = commandLine.split(" ")
        if parsedLine["name"] in splitCommand:
            flagIndex = splitCommand.index(parsedLine["name"])
        elif parsedLine["shortCut"] in splitCommand and parsedLine["shortCut"] != "":
            flagIndex = splitCommand.index(parsedLine["shortCut"])
        if flagIndex > -1:
            if isNeedUpdateNode(tabLevelBestNode[tabLevel], tabLevelBestIndex[tabLevel], parsedLine, flagIndex):
                tabLevelBestNode[tabLevel] = parsedLine
                tabLevelBestIndex[tabLevel] = flagIndex
                for i in range(tabLevel + 1, 8):
                    if tabLevelBestIndex[i] < flagIndex:
                        tabLevelBestNode[i] = None
    for bestTab in tabLevelBestNode:
        if bestTab is not None:
            nodesTree.append(bestTab)

    nodesTree.append({'shortCut': "", 'name': toolName, 'lastCommandIndex': "-1", 'nodeType': "0", 'sons': " ".join(nodesTabs), 'extra': "", 'upperNeed': "", 'description': ""})
    return nodesTree


def buildNodesTree(toolName, commandLine):
    if toolName == "mst":
        return buildMSTNodesTree(commandLine)
    else:
        return buildMFTTollNodesTree(toolName, commandLine)


def buildShortCutsNodeListNodesDeclary(nodesTree):
    shortCuts = ""
    nodesList = ""
    nodesDeclareArray = []
    for node in nodesTree:
        shortCutVal = node["shortCut"]
        nodeNAme = node["name"]
        nodeSons = node["sons"]
        extra = node["extra"]
        if nodeNAme == "--device" or nodeNAme == "-d" or nodeNAme == "-dev" or shortCutVal == "-d":
            nodeSons = "temp"
            node["lastCommandIndex"] = "1"
            node["nodeType"] = "2"
        if shortCutVal != "":
            shortCuts = shortCuts + " [\"" + shortCutVal + "\"]=\"" + nodeNAme + "\""
        if nodeSons != "" or extra != "":
            nodesList = nodesList + " " + nodeNAme
            nodeDeclare = "declare -A " + nodeNAme.replace("-", "") + "=( [\"lastCommandIndex\"]=" + node["lastCommandIndex"] + " [\"nodeType\"]=\"" + node["nodeType"] + "\" [\"sons\"]=\"" + nodeSons + "\" [\"extra\"]=\"" + node["extra"] + "\" [\"upperNeed\"]=" + node["upperNeed"] + " )\n"
            nodesDeclareArray.append(nodeDeclare)
    return shortCuts, nodesList, nodesDeclareArray


def printNodesStructures(nodesTree):
    shortCuts, nodesList, nodesDeclareArray = buildShortCutsNodeListNodesDeclary(nodesTree)

    print("declare -A shortcuts=( " + shortCuts + " )\n")
    print("nodes=\"" + nodesList + "\"\n")
    for declareLine in nodesDeclareArray:
        print(declareLine)


if __name__ == "__main__":
    userInput = getUserInput()
    toolName = userInput["name"]
    commandLine = userInput["command"]

    nodesTree = buildNodesTree(toolName, commandLine)

    printNodesStructures(nodesTree)
