#!/usr/bin/env python
# --
#                 - Mellanox Confidential and Proprietary -
#
# Copyright (C) Jan 2013, Mellanox Technologies Ltd.  ALL RIGHTS RESERVED.
#
# Except as specifically permitted herein, no portion of the information,
# including but not limited to object code and source code, may be reproduced,
# modified, distributed, republished or otherwise exploited in any form or by
# any means for any purpose without the prior written permission of Mellanox
# Technologies Ltd. Use of software subject to the terms and conditions
# detailed in the file "LICENSE.txt".
# --

from __future__ import print_function
import re
import struct

DEVICE_CONNECTIB = 0
DEVICE_SWITCHIB = 1
DEVICE_CONNECTX4 = 2
DEVICE_SPECTRUM = 3
DEVICE_CONNECTX4LX = 4
DEVICE_SWITCHIB2 = 5
DEVICE_CONNECTX5 = 6
DEVICE_QUANTUM = 7
DEVICE_SPECTRUM2 = 8
DEVICE_BLUEFIELD = 9
DEVICE_CONNECTX6 = 10
DEVICE_CONNECTX6DX = 11
DEVICE_BLUEFIELD2 = 12
DEVICE_CONNECTX6LX = 13
DEVICE_SPECTRUM3 = 14
DEVICE_CONNECTX7 = 15
DEVICE_QUANTUM2 = 16
DEVICE_BLUEFIELD3 = 17
DEVICE_SPECTRUM4 = 18

class FwTraceParser(object):
    DEVICES_WITH_IRISC_ID = (DEVICE_SWITCHIB, DEVICE_SWITCHIB2)
    DEVICES_WITHOUT_IRISC_ID = (
        DEVICE_CONNECTIB, DEVICE_CONNECTX4, DEVICE_CONNECTX5,
        DEVICE_CONNECTX6, DEVICE_CONNECTX6DX, DEVICE_SPECTRUM,
        DEVICE_CONNECTX4LX, DEVICE_QUANTUM,
        DEVICE_BLUEFIELD, DEVICE_SPECTRUM2,
        DEVICE_BLUEFIELD2, DEVICE_CONNECTX6LX,
        DEVICE_SPECTRUM3, DEVICE_SPECTRUM4, DEVICE_CONNECTX7,
        DEVICE_QUANTUM2, DEVICE_BLUEFIELD3)

    MISSING_STRING = \
        "*** MISSING STRING POINTER, EVENT DATA SEQUENSE 0 IS MISSING ***"

    #############################
    def __init__(self, deviceType, fwStrDBContents, iriscNameMap,):
        self.deviceType = deviceType
        if deviceType in self.DEVICES_WITHOUT_IRISC_ID:
            self.eventRE = re.compile(
                br"(\S+)\s+ITRACE(\d*)\s+Msn:\s*(\S+)\s+Dsn:\s*(\S+)\s+"
                b"Data:\s*(\S+)\s*")
        elif deviceType in self.DEVICES_WITH_IRISC_ID:
            self.eventRE = re.compile(
                br"(\S+)\s+ITRACE(\d*)\s+IriscId:\s*(\S+)\s+Msn:\s*(\S+)\s+"
                b"Dsn:\s*(\S+)\s+Data:\s*(\S+)\s*")
        else:
            raise RuntimeError(
                "-E- in Failed to initialize the FwParser, "
                "unknown device type=%d" % (deviceType))

        self.iriscMap = iriscNameMap
        # list of tuples (start addr, end addr, data contents)
        self.dataSections = []
        # key is irisc, value is dictionary:
        # {'ts': ts, 'msn' : msn, 'data' : {dsn : data}}
        self.lastMsg = {}
        self.specifierRE = re.compile(
            br"%([-+ #0])?(\d+)?(\.\d*)?(l|ll)?([duxs])")

        # init data sections db
        self.readDataSectionTlvs(fwStrDBContents)

    #############################
    def readDataSectionTlvs(self, fwStrDBContents):
        data = fwStrDBContents

        # read tlv by tlv
        while len(data):
            tlvType, = struct.unpack(">I", data[0:4])
            # print "type = 0x%x" % tlvType
            if tlvType == 1:  # data section
                dataSize, = struct.unpack(">I", data[4:8])
                sectionInfo = data[8: 24]
                sectionDataSize = dataSize - 16
                sectionData = data[24: dataSize + 8]
                data = data[8 + dataSize:]  # prepare data for next TLV
                # print "size = %d" % dataSize
                # print "meta data = %s" % repr(sectionInfo)
                # print "data section raw = %s" % repr(sectionData)
                # print "data section size = %d" % len(sectionData)

                startVirtualAddr, = struct.unpack(">I", sectionInfo[0:4])
                self.dataSections.append(
                    (startVirtualAddr, startVirtualAddr + sectionDataSize,
                     sectionData))

    #############################
    def getFmtString(self, ptr):
        for section in self.dataSections:
            if section[0] <= ptr <= section[1]:
                fmtData = section[2][ptr - section[0]:]
                fmtSize = fmtData.find(b'\0')
                if fmtSize + ptr >= section[1]:
                    # string crosses data section address range
                    return None
                fmt, = struct.unpack("%ds" % fmtSize, fmtData[:fmtSize])
                return fmt.replace(b"%llx", b"%08x%08x")

        return None

    #############################
    def printDataSection(self):
        for idx, section in enumerate(self.dataSections):
            try:
                print("=============== Section: %d ===================" % idx)
                for addr in range(section[0], section[1], 8):
                    print("0x%08x) " % addr,)
                    for i in range(8):
                        byte, = struct.unpack(
                            "B", section[2][addr + i - section[0]])
                        print("%02x " % byte,)

                    for i in range(8):
                        byte, = struct.unpack(
                            "B", section[2][addr + i - section[0]])
                        print("%2s " % chr(byte),)
                    print("")
            except:
                pass

    #############################
    def calcArgLen(self, fmt):
        return len(self.specifierRE.findall(fmt))

    #############################
    def pushLine(self, line):
        line = line.strip()
        m = self.eventRE.search(line)
        if m:
            # print "ts = %s, irisc = %s, msn = %s, dsn = %s, \
            # data = %s" % m.groups()
            if self.deviceType in self.DEVICES_WITHOUT_IRISC_ID:
                ts, irisc, msn, dsn, data = m.groups()
            elif self.deviceType in self.DEVICES_WITH_IRISC_ID:
                ts, irisc, iriscId, msn, dsn, data = m.groups()
                irisc = str(eval(irisc) * 2 + eval(iriscId))

            if irisc.strip() == "":
                irisc = 0xffff
            else:
                irisc = eval(irisc)
            msn = eval(msn)
            dsn = eval(dsn)
            data = eval(data)
            # print "ts = %s, irisc = 0x%x, msn = %d, dsn = 0x%x, \
            # data = 0x%x" % (ts, irisc, msn, dsn, data)
            self.pushEvent(line, ts, irisc, msn, dsn, data)

        else:
            # check if we get to the end of mlxtrace output,
            # then we should flush all buffered events
            if line.find(b"Parsed all events") != -1:
                for irisc in list(self.lastMsg.keys()):
                    self.flushIriscMsg(irisc)

            # Print as is, this is not fwtrace line
            # (maybe hw trace event or some informative message)
            print(line.decode(encoding='UTF-8',errors='ignore'))

    #############################
    def pushEvent(self, line, ts, irisc, msn, dsn, data):
        lastMsg = self.lastMsg.get(irisc)
        if not lastMsg:
            lastMsg = {'ts': ts, 'msn': msn, 'data': {dsn: data}}
            self.lastMsg[irisc] = lastMsg
        else:
            lastMsg = self.lastMsg[irisc]
            if lastMsg['msn'] != msn:
                # print "last msg msn = 0x%x, msg=0x%x, current = 0x%x" % \
                # (lastMsg['msn'], lastMsg['data'].get(0), msn)
                self.flushIriscMsg(irisc)
                lastMsg = {'ts': ts, 'msn': msn, 'data': {dsn: data}}
                self.lastMsg[irisc] = lastMsg
            else:
                lastMsg['data'][dsn] = data

        # Check if we have all required argument, then we can flush
        strPtr = self.lastMsg[irisc]['data'].get(0)
        if strPtr:
            fmtStr = self.getFmtString(strPtr)
            if fmtStr is None:
                print("*** Can't find string with pointer: 0x%x" % strPtr)
            else:
                if self.calcArgLen(fmtStr) < len(self.lastMsg[irisc]['data']):
                    self.flushIriscMsg(irisc)

    #############################
    def flushIriscMsg(self, irisc):
        lastMsg = self.lastMsg[irisc]

        if lastMsg is None:
            return

        strPtr = lastMsg['data'].get(0)
        if strPtr is None:
            print(self.MISSING_STRING)
            del self.lastMsg[irisc]
            return

        ts = lastMsg['ts']
        fmtStr = self.getFmtString(strPtr)
        if fmtStr is None:
            print("*** Can't find string with pointer: 0x%x" % strPtr)
            del self.lastMsg[irisc]
            return
        fmtStr = fmtStr.strip()

        argLen = self.calcArgLen(fmtStr)
        args = [None] * argLen

        if argLen:
            for dsn, data in lastMsg['data'].items():
                if dsn != 0:
                    args[dsn - 1] = data

        for i in range(len(args)):
            if args[i] is None:
                args[i] = "<!!!MISSING-ARGUMENT!!!>"
                fmtStr = self.replaceNthSpecifier(fmtStr, b"%s", i + 1)

        print("%-16s" % (ts.decode(encoding='UTF-8',errors='ignore')), end="")
        if irisc == 0xffff:
            print("",)
        elif irisc in self.iriscMap:
            print(" %-8s" % self.iriscMap[irisc].upper(), end="")
        else:
            print(" %-8s" % ("IRISC%-2d" % irisc), end="")

        if len(args):
            try:
                print(fmtStr.decode(encoding='UTF-8',errors='ignore') % tuple(args))
            except Exception as exp:
                raise Exception(
                    "-E- in StrPtr=0x%x, fmt(%s): %s" % (strPtr, fmtStr, exp))

        else:
            print(fmtStr.decode(encoding='UTF-8',errors='ignore'))
        del self.lastMsg[irisc]

    #############################
    def flushAll(self):
        iriscs = self.lastMsg.keys()
        for irisc in iriscs:
            self.flushIriscMsg(irisc)

    #############################
    def replaceNthSpecifier(self, fmt, rep, n):

        def replaceNthWith(n, replacement):

            def replace(match, c=[0]):
                c[0] += 1
                if c[0] == n:
                    return replacement
                else:
                    return match.group(0)

            return replace

        return self.specifierRE.sub(replaceNthWith(n, rep), fmt)


#####################################################
if __name__ == "__main__":
    iriscs = {0: 'ir0', 1: 'ir1', 2: 'ir2', 3: 'ir3', 4: 'ir4'}
    import sys
    if len(sys.argv) < 3:
        print("-E- missing fw strings db file path and string pointer")
        sys.exit(1)

    data = open(sys.argv[1], "r+b").read()
    fwParser = FwTraceParser(DEVICE_CONNECTIB, data, iriscs)
    ptr = eval(sys.argv[2])
    # fwParser.printDataSection()
    fmt = fwParser.getFmtString(ptr)
    if fmt is None:
        print("Coudn't find string ptr: 0x%x" % ptr)
    else:
        print("ptr 0x%x == > %s" % (ptr, fmt))

    sys.exit(0)

    hash = open(sys.argv[1], "r").read()
    fwParser = FwTraceParser(DEVICE_CONNECTIB, hash, iriscs)
    # for line in open("testing/connectib_dump_example_ooo.txt", "r").\
    # readlines():
    import sys
    import string
    import subprocess
    import os
    import signal

    # for line in sys.stdin.readlines():
    #    t.pushLine(line)
    # t.flushAll()

    cmd = ["./tmp.py"]
    print(" ".join(cmd))
    stdout = None
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
    try:
        while True:
            fwParser.pushLine(proc.stdout.readline())
            # print line,
    except KeyboardInterrupt:
        os.kill(proc.pid, signal.SIGKILL)
        proc.wait()
        print("\nStopping & Flushing... "\
            "(messages below may have missing arguments)")
        fwParser.flushAll()