#!/usr/bin/env python3
import serial
import os
import sys
import time
import argparse

def progressbar(it, prefix="", size=60, file=sys.stdout):
    count = len(it)
    def show(j):
        x = int(size*j/count)
        file.write("%s[%s%s] %i/%i\r" % (prefix, "#"*x, "."*(size-x), j, count))
        file.flush()
    show(0)
    for i, item in enumerate(it):
        yield item
        show(i+1)
    file.write("\n")
    file.flush()

def rom_upload(filename):
    # ser = serial.Serial("COM3", timeout=1, write_timeout=1)
    ser = serial.Serial("/dev/ttyUSB0", timeout=1, write_timeout=1, baudrate=115200)
    if ser.is_open:
        # Clear out any existing input
        ser.write(b'\n')
        time.sleep(0.002)

        # Send the upload command
        ser.write(b'#')
        time.sleep(0.002)
        ser.write(b'u')
        time.sleep(0.002)
        ser.write(b'\n')
        time.sleep(0.002)

        size = os.path.getsize(filename)
        ser.write([size & 0xFF])
        time.sleep(0.002)
        ser.write([(size >> 8) & 0xFF])
        time.sleep(0.002)

        with open(filename, "rb") as f:
            for i in progressbar(range(size), "Upload: ", 40):
                byte = f.read(1)
                ser.write(byte)
                time.sleep(0.002)

        ser.close()

    else:
        print("Failed to open serial port")

def serial_upload(filename):
    # ser = serial.Serial("COM3", timeout=1, write_timeout=1)
    ser = serial.Serial("/dev/ttyUSB0", timeout=1, write_timeout=1, baudrate=115200)
    if ser.is_open:
        size = os.path.getsize(filename)

        with open(filename, "rb") as f:
            for i in progressbar(range(size), "Upload: ", 40):
                byte = f.read(1)
                ser.write(byte)

                if byte == b'#':
                    time.sleep(0.002)
                    ser.write(b'#')
                    time.sleep(0.002)
                    ser.write(b'\n')

                time.sleep(0.002)

        ser.close()

    else:
        print("Failed to open serial port")

def main():
    parser = argparse.ArgumentParser(description="Upload binaries to the z80 computer.")
    parser.add_argument("filename", help="file to upload")
    parser.add_argument("--rom", dest="rom", action="store_const", const=True, default=False, help="Upload binary to rom")
    parser.add_argument("--serial", dest="serial", action="store_const", const=True, default=False, help="Upload binary over serial")

    args = parser.parse_args()

    if (args.rom):
        rom_upload(args.filename)
    elif (args.serial):
        serial_upload(args.filename)
    else:
        print("You needs to specify a target")



if "__main__":
    main()