#!/usr/bin/python3
# Routine to convert .tzx files to .sna files
# NOTE: these files will not run in an emulator as .tzx files do not contain register information,
# so the header block of the .sna file is left blank. However, the .tzx files can be used in tools
# such as unQuill to extract data from them

import sys
# Represents ZX Spectrum RAM (48k only)
class Memory:
    def __init__(self):
        self.mem = [0]*49152
        self.base = 16384

    def write_byte(self,addr,val):
        self.mem[addr-self.base] = val

    def write_file(self,name):
        print("Writing output file {}".format(name))
        try:
            f = open(name, "wb")
            f.write(bytearray([0]*27)) # Fake header
            f.write(bytearray(self.mem))
        except IOError:
            print("Error opening file {}".format(name))


class TzxFile:
    def __init__(self):
        self.block_num = 1
        self.block_start = 0
        self.block_size = 0
        self.load_addr = None;
        self.mem = Memory()

    def write(self,name):
        self.mem.write_file(name)

    def read_byte(self,offset):
        return self.contents[self.block_start+offset]

    def read_word(self,offset):
        return self.contents[self.block_start+offset] + 256 * self.contents[self.block_start+offset+1]

    def read_dword(self,offset):
        return self.read_word(offset) + 65536 * self.read_word(offset + 2)

    def load(self, filename):
        try:
            f = open(filename, "rb")
            self.contents = f.read(-1)
        except IOError:
            print("Error opening file")
            exit(0)

    def check_header(self):
        preamble_bytes = self.contents[0:7]
        preamble = "".join(map(chr, preamble_bytes))
        if preamble != "ZXTape!":
            print("Unidentified header in TZX file")
            print(preamble)
            exit(0)
        major = self.read_byte(8)
        minor = self.read_byte(9)
        print("Found .TZX file; version: {}.{}".format(major, minor))
        self.block_start += 10

    # Process the next block in the file.
    # Returns true if we haven't reached the end of the file yet
    def process_block(self):
        block_type = self.read_byte(0)
        print("\nOffset: {}".format(self.block_start))
        print("Block type: {}".format(block_type))
        self.block_start += 1
        if block_type == 16:
            self.process_std()
        elif block_type == 17:
            self.process_turbo()
        elif block_type == 18:
            self.process_puretone()
        elif block_type == 19:
            self.process_sequence()
        elif block_type == 20:
            self.process_puredata()
        elif block_type == 21:
            self.process_samples()
        elif block_type == 24:
            self.process_csw()
        elif block_type == 25:
            self.process_general()
        elif block_type == 32:
            self.process_pause()
        elif block_type == 33:
            self.process_grpstart()
        elif block_type == 34:
            self.process_grpend()
        elif block_type == 35:
            self.process_jump()
        elif block_type == 36:
            self.process_loopstart()
        elif block_type == 37:
            self.process_loopend()
        elif block_type == 38:
            self.process_call()
        elif block_type == 39:
            self.process_return()
        elif block_type == 40:
            self.process_select()
        elif block_type == 42:
            self.process_stoptape()
        elif block_type == 43:
            self.process_setsignal()
        elif block_type == 48:
            self.process_text()
        elif block_type == 49:
            self.process_msg()
        elif block_type == 50:
            self.process_archive()
        elif block_type == 51:
            self.process_hw()
        elif block_type == 53:
            self.process_custom()
        elif block_type == 90:
            self.process_glue()
        else:
            self.process_unknown()

        self.block_num += 1
        return self.block_start < len(self.contents)

    def process_std(self):
        print("Standard Speed data block")
        data_len = self.read_word(2)
        print("Data length: {}".format(data_len))
        self.block_start += 4
        self.handle_tap_data(data_len)
        self.block_start += data_len

    def process_turbo(self):
        print("Turbo Speed data block")
        data_len = self.read_byte(0x0f) + 256*self.read_byte(0x10) + 65536*self.read_byte(0x11)
        print("Data length: {}".format(data_len))
        self.block_start += 0x12 
        self.handle_tap_data(data_len)
        self.block_start += data_len

    def process_puredata(self):
        print("Pure data block")
        data_len = self.read_byte(0x07) + 256*self.read_byte(0x08) + 65536*self.read_byte(0x09)
        print("Data length: {}".format(data_len))
        self.block_start += 0x0a 
        self.handle_tap_data(data_len)
        self.block_start += data_len

    def process_archive(self):
        print("Archive block")
        items = self.read_byte(2)
        print("{} items in archive".format(items))
        self.block_start += 3
        for i in range(0,items):
            self.process_item(i)

    def process_item(self,idx):
        item_type = self.read_byte(0)
        item_len = self.read_byte(1)
        item = self.contents[self.block_start + 2:self.block_start + 2 + item_len]
        print("Item {}: {}".format(idx,"".join(map(chr, item))))
        self.block_start += 2 + item_len

    def process_puretone(self):
        print("Puretone block - ignoring")
        self.block_start += 4

    def process_sequence(self):
        print("Sequence block - ignoring")
        num_pulses = self.read_byte(0)
        self.block_start += num_pulses * 2 + 1

    def process_samples(self):
        print("Samples block - ignoring")
        num_samples = self.read_byte(5)
        self.block_start += num_samples + 8

    def process_csw(self):
        print("CSW block - ignoring")
        block_len = self.read_dword(0)
        self.block_start += block_len + 4

    def process_general(self):
        print("General block - ignoring")
        block_len = self.read_dword(0)
        self.block_start += block_len + 4

    def process_pause(self):
        print("Pause block - ignoring")
        self.block_start += 2

    def process_grpstart(self):
        print("Group start - ignoring")
        length = self.read_byte(0)
        self.block_start += length + 1

    def process_grpend(self):
        print("Group end - ignoring")

    def process_jump(self):
        print("Jump block - ignoring")
        self.block_start += 2

    def process_loopstart(self):
        print("Loop start - ignoring")
        length = self.read_byte(0)
        self.block_start += length + 1

    def process_loopend(self):
        print("Loop end - ignoring")

    def process_call(self):
        print("Call block - ignoring")
        calls = self.read_word(0)
        self.block_start += 2*calls + 2

    def process_return(self):
        print("Return block - ignoring")

    def process_select(self):
        print("Select block - ignoring")
        length = self.read_word(0)
        self.block_start += length + 2

    def process_stoptape(self):
        print("Stop tape block - ignoring")
        self.block_start += 4

    def process_setsignal(self):
        print("Set signal block - ignoring")
        self.block_start += 5

    def process_text(self):
        print("Text block")
        length = self.read_byte(0)
        text = self.contents[self.block_start + 1:self.block_start + 1 + length]
        print("{}".format("".join(map(chr, text))))
        self.block_start += length + 1

    def process_msg(self):
        print("Message block")
        length = self.read_byte(1)
        text = self.contents[self.block_start + 2:self.block_start + 2 + length]
        print("{}".format("".join(map(chr, text))))
        self.block_start += length + 2

    def process_hw(self):
        print("Hardware block - ignoring")
        num = self.read_byte(0)
        self.block_start += 3*num + 1

    def process_custom(self):
        print("Custom block - ignoring")
        length = self.read_dword(10)
        self.block_start += length + 14

    def process_glue(self):
        print("Glue block - ignoring")
        self.block_start += 9

    def process_unknown(self):
        print("Unknown block - ignoring")
        block_len = self.read_dword(0)
        self.block_start += block_len + 4

    def handle_tap_data(self, data_length):
        block_type = self.read_byte(0)
        if (block_type == 0):
            print("TAP header block")
            self.read_header()
        elif (block_type == 255):
            print("TAP data block")
            self.read_data(data_length)
        else:
            print("Unknown TAP block marker {}".format(block_type))
            exit(0)

    def read_header(self):
        datatype = self.read_byte(1)
        name = ""
        for i in range(2,12):
            name += chr(self.read_byte(i))
        self.length = self.read_word(12)
        self.start_address = self.read_word(14)
        if datatype == 3:
            self.load_addr = self.start_address
        else:
            self.load_addr = None
        print("Block {}; ".format(self.block_num), end='')
        print("Datatype: {}; ".format(datatype), end='')
        print("Name: {}; ".format(name), end='')
        print("Start address: {}; ".format(self.start_address), end='')
        print("Length: {}".format(self.length))

    def read_data(self,data_len):
        print("Block {}; ".format(self.block_num), end='')
        data_size = data_len - 2
        print("Data Size: {}".format(data_size))
        if self.load_addr != None:
                print("Loading into memory at: {}".format(self.load_addr))
                for i in range(0,data_size):
                    if self.load_addr + i > 65535:
                        break
                    self.mem.write_byte(self.load_addr+i, self.read_byte(i+1))
        else:
            print("No load address specified - ignoring block")
                    
def start():
        tzx = TzxFile()
        inFile = sys.argv[1]
        outFile = sys.argv[2]
        if (len(sys.argv) > 3):
            maxBlock = int(sys.argv[3])
        else:
            maxBlock = 255
        tzx.load(inFile)
        tzx.check_header()

        while tzx.process_block() and tzx.block_num <= maxBlock:
            pass
        
        # Fake a Quill system messages base address
        tzx.mem.write_byte(0x5c7b,0x33)
        tzx.mem.write_byte(0x5c7c,0x69)
        tzx.write(outFile)

if __name__ == "__main__":
	if (len(sys.argv) < 3):
	        print("Usage: tzx2sna.py <input filename> <output filename> [max TZX block to process]")
	else:
	    start()
