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

import mojo

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 upload_rom(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 upload_bload(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)

        print("Size: {} ({} pages)".format(size, math.ceil(size/256)))

        input("Press enter to start upload")

        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 upload_sd(filename):
    if os.geteuid() == 0:
        print("Running as root")
        file_in = open(filename, "rb")
        file_out = open("/dev/sdb", "wb")

        file_out.write(file_in.read())
    else:
        print("Reqeusting root permission")
        subprocess.call(["sudo", "python3", *sys.argv])
        sys.exit()


# @todo For now we just wrap the existing mojo.py tool
# At some point we should rewrite it from scratch
def upload_mojo(filename):
    # Serial code
    try:
        ser = serial.Serial("/dev/ttyACM0", 19200, timeout=20)
    except:
        print("No serial port found named /dev/ttyACM0")
        sys.exit(1)

    mojo.install_mojo(ser, filename, True, False, False, True)

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("--bload", dest="bload", action="store_const", const=True, default=False, help="Upload binary to bload")
    parser.add_argument("--mojo", dest="mojo", action="store_const", const=True, default=False, help="Upload binary to mojo")
    parser.add_argument("--sd", dest="sd", action="store_const", const=True, default=False, help="Upload binary to sd card")

    args = parser.parse_args()

    if (args.rom):
        upload_rom(args.filename)
    elif (args.bload):
        upload_bload(args.filename)
    elif (args.mojo):
        upload_mojo(args.filename)
    elif (args.sd):
        upload_sd(args.filename)
    else:
        print("You needs to specify a target")

if __name__ == "__main__":
    main()