#!/usr/bin/env python3
'''
UltraScalePlus bitstream analyzer tool.

This script reads a UltraScalePlus bitstream and prints out some useful information.
It can also create a frames file with the configuration data words.
The bitstream is analyzed word by word and interpreted according to
the UG570 Configuration User Guide.

The tool can be used to derive the initialization, startup and finalization
sequence as well as the configuration data. The latter is written to a frames
file which can be used by the bitstream tools such as frames2bit to generate
a valid bitstream.
'''

import argparse
from io import StringIO

conf_regs = {
    0: "CRC",
    1: "FAR",
    2: "FDRI",
    3: "FDRO",
    4: "CMD",
    5: "CTL0",
    6: "MASK",
    7: "STAT",
    8: "LOUT",
    9: "COR0",
    10: "MFWR",
    11: "CBC",
    12: "IDCODE",
    13: "AXSS",
    14: "COR1",
    16: "WBSTAR",
    17: "TIMER",
    22: "BOOTSTS",
    24: "CTL1",
    31: "BSPI"
}

cmd_reg_codes = {
    0: "NULL",
    1: "WCFG",
    2: "MFW",
    3: "LFRM",
    4: "RCFG",
    5: "START",
    7: "RCRC",
    8: "AGHIGH",
    9: "SWITCH",
    10: "GRESTORE",
    11: "SHUTDOWN",
    13: "DESYNC",
    15: "IPROG",
    16: "CRCC",
    17: "LTIMER",
    18: "BSPI_READ",
    19: "FALL_EDGE"
}

opcodes = ("NOP", "READ", "WRITE", "UNKNOWN")

class Bitstream:
    def __init__(self, file_name, verbose=False):
        self.frame_data = []
        self.fdri_write_len = 0
        self.fdri_in_progress = False
        self.words = []
        with open(file_name, "rb") as f:
            word = f.read(4)
            while word:
                self.words.append(int.from_bytes(word, byteorder = 'big'))
                word = f.read(4)
        pos, self.header = self.get_header()
        self.body = self.words[pos + 1:]
        self.parse_bitstream(verbose)

    def get_header(self):
        '''Return position and content of header'''
        pos = self.words.index(0xaa995566)
        return pos, self.words[:pos + 1]

    def parse_bitstream(self, verbose):
        payload_len = 0
        for word in self.body:
            if payload_len > 0:
                payload_len = self.parse_reg(
                    reg_addr, word, payload_len, verbose)
                continue
            else:
                packet_header = self.parse_packet_header(word)
                opcode = packet_header["opcode"]
                reg_addr = packet_header["reg_addr"]
                words = packet_header["word_count"]
                type = packet_header["type"]
                if verbose:
                    if not opcode:
                        print("\n\tNOP")
                    else:
                        print(
                        "\n\tConfiguration Register Word: ", hex(word),
                        'Type: {}, Op: {}, Addr: {} ({}), Words: {}'.format(
                            type, opcodes[opcode], conf_regs[reg_addr] if reg_addr in conf_regs else "UNKNOWN", reg_addr, words))
                if opcode and reg_addr in conf_regs:
                    payload_len = words
                    if conf_regs[reg_addr] == "FDRI" and type == 1:
                        self.fdri_in_progress = True
                        self.fdri_write_len = payload_len
                    continue

    def parse_packet_header(self, word):
        type = (word >> 29) & 0x7
        opcode = (word >> 27) & 0x3
        reg_addr = (word >> 13) & 0x1F
        if type == 1:
            word_count = word & 0x7FF
        elif type == 2:
            word_count = word & 0x7FFFFFF
        else:
            word_count = 0
        return {
            "type": type,
            "opcode": opcode,
            "reg_addr": reg_addr,
            "word_count": word_count
        }

    def parse_command(self, word, verbose):
        if verbose:
            print("\tCommand: {} ({})".format(cmd_reg_codes[word], word))

    #TODO Add COR0 options parsing
    def parse_cor0(self, word, verbose):
        if verbose:
            print("\tCOR0 options: {:08X}".format(word))

    def parse_reg(self, reg_addr, word, payload_len, verbose):
        assert reg_addr in conf_regs
        reg = conf_regs[reg_addr]
        if reg == "CMD":
            self.parse_command(word, verbose)
        elif reg == "COR0":
            self.parse_cor0(word, verbose)
        elif reg == "FDRI":
            # We are in progress of a FDRI operation
            # Keep adding data words 
            if self.fdri_in_progress:
                if verbose:
                    print("\t{:2d}. 0x{:08X}".format(self.fdri_write_len - payload_len, word))
                self.frame_data.append(word)
                if payload_len == 1:
                    self.fdri_in_progress = False
                    return 0
            else:
                #FIXME add Type 2 FDRI writes
                assert False
                #self.curr_fdri_write_len = word
                #self.fdri_in_progress = True
                # Check if 0 words actually means read something
                #payload_len = self.curr_fdri_write_len
                #if verbose:
                #    print("\t{}: {}\n".format(reg, self.curr_fdri_write_len))
                #return payload_len
        else:
            if verbose:
                print("\tRegister: {} Value: 0x{:08X}".format(reg, word))
        payload_len -= 1
        return payload_len

    def write_frames_txt(self, file_name):
        '''Write frame data in a more readable format'''
        frame_stream = StringIO()
        for i in range(len(self.frame_data)):
            if i % 93 == 0:
                frame_stream.write("\nFrame {:4}\n".format(i // 93))
        with open(file_name, "w") as f:
            print(frame_stream.getvalue(), file=f)

    #FIXME Add storing of frame address
    def write_frames(self, file_name):
        '''Write configuration data to frames file'''
        frame_stream = StringIO()
        for i in range(len(self.frame_data)):
            if i % 93 == 0:
                frame_stream.write("0x{:08x} ".format(i // 93))
            frame_stream.write("0x{:04x}".format(self.frame_data[i]))
            if i % 93 == 92:
                frame_stream.write("\n")
            elif i < len(self.frame_data) - 1:
                frame_stream.write(",")
        with open(file_name, "w") as f:
            print(frame_stream.getvalue(), file=f)


def main(args):
    verbose = not args.silent
    bitstream = Bitstream(args.bitstream, verbose)
    print("Frame data length: ", len(bitstream.frame_data))
    if args.frames_out:
        bitstream.write_frames(args.frames_out)
        #bitstream.write_frames_txt(args.frames_out + ".txt")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--bitstream', help='Input bitstream')
    parser.add_argument('--frames_out', help='Output frames file')
    parser.add_argument(
        '--silent', help="Don't print analysis details", action='store_true')
    args = parser.parse_args()
    main(args)