#!/usr/bin/env python3

import struct
from typing import NamedTuple, Union
from enum import IntEnum
import sys

WORD_BE = struct.Struct(">I")
WORD_LE = struct.Struct("<I")

WORD = WORD_BE
WORD_SWAPPED = False

def _pack_word(word: int) -> bytes:
    return WORD.pack((word,))
def _unpack_word(word: bytes) -> int:
    return WORD.unpack(word)[0]

def swap_endian() -> None:
    global WORD_SWAPPED
    global WORD
    if WORD_SWAPPED:
        raise Exception("Swapping enddianness twice oops")
    WORD_SWAPPED = True
    WORD = WORD_LE

def _hex_pad(word: int) -> str:
    return "{0:#0{1}x}".format(word, 10)


DUMMY_PAD       = 0xFFFFFFFF
BUS_WIDTH_DET_1 = 0x000000BB
BUS_WIDTH_DET_1_EXCEPT_WRONG_WAY = 0xBB000000
BUS_WIDTH_DET_2 = 0x11220044
SYNC_WORD       = 0xAA995566

PACKET_TYPE_1   = 1
PACKET_TYPE_2   = 2

CFG_FRAME_WORD_LENGTH = 101

class MiscWord(IntEnum):
    DUMMY_PAD = 1
    BUS_WIDTH_DET_1 = 2
    BUS_WIDTH_DET_2 = 3
    SYNC_WORD = 4

class PacketOp(IntEnum):
    NOP = 0
    READ = 1
    WRITE = 2
    RESERVED = 3

class Register(IntEnum):
    CRC    = 0b00000
    FAR    = 0b00001
    FDRI   = 0b00010
    FDRO   = 0b00011
    CMD    = 0b00100
    CTL0   = 0b00101
    MASK   = 0b00110
    STAT   = 0b00111
    LOUT   = 0b01000
    COR0   = 0b01001
    MFWR   = 0b01010
    CBC    = 0b01011
    IDCODE = 0b01100
    AXSS   = 0b01101
    COR1   = 0b01110
    WBSTAR = 0b10000
    TIMER  = 0b10001
    BOOTSTS= 0b10110
    CTL1   = 0b11000
    BSPI   = 0b11111

    UNDOC_1= 0b10011 # Documentation? what's that lol ~ Xilinx, probably

class FarType(IntEnum):
    FABRIC  = 0
    BRAM    = 1
    CFG_CLB = 2
    UNDOC_1 = 7 # And again

class FarValue(NamedTuple):
    blocktype: FarType
    bottom_top: int
    row_addr: int
    col_addr: int
    minor_addr: int

class CmdValue(IntEnum):
    NULL        = 0b00000
    WCFG        = 0b00001
    MFW         = 0b00010
    DGHIGH_LFRM = 0b00011
    RCFG        = 0b00100
    START       = 0b00101
    AGHIGH      = 0b01000
    RCRC        = 0b00111
    SWITCH      = 0b01001
    GRESTORE    = 0b01010
    SHUTDOWN    = 0b01011
    DESYNC      = 0b01101

class Type1(NamedTuple):
    op: PacketOp
    address: Register
    word_count: int

class Type2(NamedTuple):
    op: PacketOp
    word_count: int


def _skip_xil_header(data: bytes) -> bytes:
    # I really don't know what in the nondeterministic fuck is going on with this
    # there's a size field at the end that might be important
    # see http://www.pldtool.com/pdf/fmt_xilinxbit.pdf
    for i in range(len(data) - 4):
        word = _unpack_word(data[i:i+4])
        if word == DUMMY_PAD:
            if i >= 4:
                # we found a header
                print(".......... xilinx header")
                data_len = _unpack_word(data[i-4:i])
                # safety check the size field
                assert data_len == len(data) - i
            else:
                # no header
                assert i == 0
            # return start of bitstream data
            return i
    raise Exception("Didn't find pad word!")

def _parse_far_word(word: int) -> FarValue:
    try:
        blk_type = (word >> 23) & 0b111
        bottom_top = (word >> 22) & 0b1
        row_addr = (word >> 17) & 0b11111
        col_addr = (word >> 7) & 0b1111111111
        minor_addr = word & 0b1111111
        return FarValue(FarType(blk_type), bottom_top, row_addr, col_addr, minor_addr)
    except:
        raise Exception("Invalid FAR value: " + _hex_pad(word))

def _parse_cmd_word(word: int) -> CmdValue:
    try:
        return CmdValue(word)
    except:
        raise Exception("Invalid command value: " + _hex_pad(word))

def _parse_next_word(word: int) -> Union[MiscWord, Type1, Type2]:
    try:
        if word == DUMMY_PAD:
            return MiscWord.DUMMY_PAD
        elif word == BUS_WIDTH_DET_1_EXCEPT_WRONG_WAY:
            swap_endian()
            return MiscWord.BUS_WIDTH_DET_1
        elif word == BUS_WIDTH_DET_1:
            return MiscWord.BUS_WIDTH_DET_1
        elif word == BUS_WIDTH_DET_2:
            return MiscWord.BUS_WIDTH_DET_2
        elif word == SYNC_WORD:
            return MiscWord.SYNC_WORD
        else:
            header_bits = (word >> 29) & 0b111
            opcode = (word >> 27) & 0b11
            if header_bits == PACKET_TYPE_1:
                addr = (word >> 13) & 0b11111
                count = word & 0b1111111111
                return Type1(PacketOp(opcode), Register(addr), count)
            elif header_bits == PACKET_TYPE_2:
                count = word & 0b11111111111111111111111111
                return Type2(PacketOp(opcode), count)
        raise Exception()
    except:
        raise Exception("Unknown word: " + _hex_pad(word))

def parse_bitstream(data: bytes) -> None:
    i = _skip_xil_header(data)
    data_counter = 0
    last_far = None
    bram_addrs = set()
    while i < len(data):
        word = _unpack_word(data[i:i+4])
        obj = _parse_next_word(word)
        print(_hex_pad(word), obj)

        i += 4

        if isinstance(obj, Type1) and obj.address == Register.FAR:
            assert obj.word_count == 1
            dat_word = _unpack_word(data[i:i+4])
            last_far = _parse_far_word(dat_word)
            if last_far.blocktype == FarType.BRAM:
                bram_addrs.add(dat_word)
            print(_hex_pad(dat_word), last_far)
            i += 4
        elif isinstance(obj, Type1) and obj.address == Register.CMD:
            assert obj.word_count == 1
            dat_word = _unpack_word(data[i:i+4])
            print(_hex_pad(dat_word), _parse_cmd_word(dat_word))
            i += 4
        elif isinstance(obj, Type1) or isinstance(obj, Type2):
            if obj.word_count > 0:
                if obj.word_count < 4:
                    for j in range(obj.word_count):
                        dat_word = _unpack_word(data[i+j*4:i+j*4+4])
                        print(_hex_pad(dat_word), "data")
                else:
                    l = f"{data_counter}"
                    if last_far is not None and last_far.blocktype == FarType.BRAM:
                        l = f"_bram_{last_far.bottom_top}_{last_far.row_addr}_{last_far.col_addr}_{last_far.minor_addr}"
                    out_file = f"{sys.argv[1]}-frames-{l}.dat"
                    data_counter += 1
                    print("..........", obj.word_count, "words (written to ", out_file, ")")
                    with open(out_file, "wb") as f:
                        f.write(data[i:i+obj.word_count*4])
                i += obj.word_count * 4

    if i != len(data):
        raise Exception("Desync??")

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: ./xilbitstream.py <bitstream-file>")
        sys.exit()
    with open(sys.argv[1], "rb") as f:
        data = f.read()
    parse_bitstream(data)