#!/usr/bin/env python3
"""Next Gen EVSE downloader"""
#
# -*-mode: python; coding: utf-8 -*-
#
# Copyright 2020-2021 Garo AB

import os
import sys
import getopt
import tarfile
import tempfile
import ngcommon
import subprocess
import serial
import io
import time

windows = os.name == "nt"

def usage():
    "Exit with usage instructions"
    sys.exit(f"""Usage: {sys.argv[0]} [options] <image-file>|--dev 
options:
    --dev                          Use image from development tree
    --force <serial port>          Force and use specified USB port. E.g. 'COM7' on Windows or 'ttyUSB0' in Linux.
    --mem256                       Limit memory to 256MB.
    --loadinterface                Load Interface, force use of imx6ull-garo-nextgen-li.dtb.
    --get-serial                   Get serial number, then exit
    --set-serial <hex-serial>      WARNING: Set serial number (MACs) permanently, then exit
                                   <hex-serial> should be MAC address ending with zero
    --yes                          Confirm operations such as --set-serial
    --override-serial <hex-serial> Temporarily set serial number, then exit
    --ucmd <args> ...              Run u-boot ucmd command, then exit

example:
    nextgen-downloader --force=COM14 garo-nextgen-0.1.0.img
    """)


# Using universal_newlines for Python 3.6 compatibility
def run_command(*popenargs, check=True, universal_newlines=True, **kwargs):
    "Run command with check and as text"
    proc = subprocess.run(*popenargs, check=check, universal_newlines=universal_newlines, **kwargs)
    return proc


def sudo_command(*popenargs, **kwargs):
    "Run command using sudo"
    args = popenargs[0]
    popenargs = popenargs[1:]
    if not windows:
        args = ["sudo"] + args
    return run_command(args, *popenargs, **kwargs)


def run_ucmd(args):
    "Run uuu ucmd command"
    sudo_command(["uuu", "FB:", "ucmd", "setenv", "stdout", "serial,fastboot"])
    cmd = ["uuu", "-v", "FB:", "ucmd"]
    cmd += args
    sudo_command(cmd, check=False)


def get_macwords():
    "Retrieve OTP words for MAC addresses"
    sudo_command(["uuu", "FB:", "ucmd", "setenv", "stdout", "serial,fastboot"])
    cmd = sudo_command(["uuu", "-v", "FB:", "ucmd", "fuse", "read", "4", "2", "3"],
                       stdout=subprocess.PIPE)
    for line in cmd.stdout.splitlines():
        if line.startswith("Word 0x"):
            dummy1, dummy2, ocotp_mac0, ocotp_mac1, ocotp_mac = line.split()
            break

    ocotp_mac0 = int(ocotp_mac0, 16)
    ocotp_mac1 = int(ocotp_mac1, 16)
    ocotp_mac = int(ocotp_mac, 16)
    return ocotp_mac0, ocotp_mac1, ocotp_mac


def set_macwords(macwords, permanent, yes):
    "Write OTP words for MAC addresses"
    ocotp_mac0, ocotp_mac1, ocotp_mac = macwords
    ocotp_mac0 = f"{ocotp_mac0:08x}"
    ocotp_mac1 = f"{ocotp_mac1:08x}"
    ocotp_mac = f"{ocotp_mac:08x}"

    args = ["uuu", "FB:", "ucmd", "fuse"]
    if permanent:
        args += ["prog"]
        if yes:
            args += ["-y"]
    else:
        args += ["override"]
    args += ["4", "2", ocotp_mac0, ocotp_mac1, ocotp_mac]
    sudo_command(args)


def macwords_to_mac(macwords, iface):
    "Convert OTP words to MAC adress, for specified interface (0 or 1)"
    # The 3 fuse banks together contains the bits for 2 MACs:
    # mac0/mac1 contains bits for interface 0
    # mac1/mac contains bits for interface 1
    # Note that u-boot sometimes sets CONFIG_ETHPRIME=eth1, which
    # means that the interfaces are reversed in u-boot and Linux
    ocotp_mac0, ocotp_mac1, ocotp_mac = macwords
    if iface == 0:
        return (ocotp_mac1 & 0xffff) << 32 | ocotp_mac0
    if iface == 1:
        return (ocotp_mac << 16) | ocotp_mac1 >> 16
    return None


def mac_to_macwords(macwords, newmac, permanent, iface):
    "Convert MAC to OTP words, for specified interface"
    ocotp_mac0, ocotp_mac1, ocotp_mac = macwords

    oldmac = macwords_to_mac(macwords, iface)
    # We can only change 0 to 1, not the opposite
    if permanent and oldmac & ~newmac:
        print(f"Impossible to un-burn fuse bits required for if{iface} MAC {newmac:012x}", file=sys.stderr)
        sys.exit(1)
    print(f"Replacing if{iface} MAC {oldmac:012x} with {newmac:012x}", file=sys.stderr)

    if iface == 0:
        ocotp_mac0 = newmac & 0xffffffff
        ocotp_mac1 = (ocotp_mac1 & 0xffff0000) | newmac >> 32
    elif iface == 1:
        ocotp_mac1 = (ocotp_mac1 & 0xffff) | (newmac & 0xffff) << 16
        ocotp_mac = newmac >> 16
    else:
        raise Exception(f"Internal error: bad interface {iface}")

    macwords = ocotp_mac0, ocotp_mac1, ocotp_mac
    return macwords

# Note about naming:
# Interface 0 = ENET1 in reference manual
# Interface 1 = ENET2 in reference manual

def get_serial():
    "Retrieve and print serial number"
    # Warnings/infos printed to stderr, serial to stdout
    macwords = get_macwords()
    if0mac = macwords_to_mac(macwords, 0)
    print(f"if0 MAC address: {if0mac:012x}", file=sys.stderr)
    if1mac = macwords_to_mac(macwords, 1)
    print(f"if1 MAC address: {if1mac:012x}", file=sys.stderr)
    if (if0mac >> 4) != (if1mac >> 4):
        print(f"Warning: if0/if1 MACs do not have common prefix", file=sys.stderr)


def set_serial(sernum, permanent, yes):
    "Set serials and MACs"
    sernum = int(sernum, 16)
    macwords = get_macwords()
    oldwords = macwords

    newmac0 = sernum
    macwords = mac_to_macwords(macwords, newmac0, permanent, 0)

    newmac1 = ngcommon.serial_to_mac(sernum, 1)
    if newmac1 is not None:
        macwords = mac_to_macwords(macwords, newmac1, permanent, 1)
    else:
        print(f"Warning: Cannot derive MAC for if1 - skipping ", file=sys.stderr)

    if oldwords != macwords:
        set_macwords(macwords, permanent, yes)

    # When using CONFIG_NET=n these env variables are no longer set by u-boot so set them manually.
    sudo_command(["uuu", "FB:", "ucmd", "setenv", "ethaddr", ngcommon.mac_to_readable(newmac0)])
    if newmac1 is not None:
        sudo_command(["uuu", "FB:", "ucmd", "setenv", "eth1addr", ngcommon.mac_to_readable(newmac1)])

    # Save environment if changes are to be permanent.
    if permanent:
        sudo_command(["uuu", "FB:", "ucmd", "saveenv"])


def flush_n_read(sio):
    sio.flush()
    try:
        line = sio.readline()
    except serial.serialutil.SerialException as e:
        print("warning:", e)
        return ""
    line = line.rstrip()
    if line != "":
        print(line)
    return line

ser = None
sio = None
def setup_serial(serialport):
    global ser, sio
    try:
        ser = serial.Serial(port=serialport,\
                        baudrate=115200,\
                        parity=serial.PARITY_NONE,\
                        stopbits=serial.STOPBITS_ONE,\
                        bytesize=serial.EIGHTBITS,\
                        timeout=0.1)
        sio = io.TextIOWrapper(io.BufferedRWPair(ser, ser))
    except Exception as exc:
        print("Error setting up serial port: ", exc)

def force_usb_boot():
    global ser, sio
    if sio is None:
        print("Serial port not open, aborting.")
        return False
    print("Logging in to device...")
    sio.write("\r")
    while True:
        line = flush_n_read(sio)
        if line == "":
            break
    sio.write("admin\r")
    flush_n_read(sio)
    sio.write("iapc-hw9s-mgn5\r")
    flush_n_read(sio)
    time.sleep(1)
    print("Rebooting device...")
    sio.write("\r")
    flush_n_read(sio)
    sio.write("sudo /sbin/reboot\r")
    flush_n_read(sio)
    print("Waiting for autoboot message...")
    while True:
        try:
            line = flush_n_read(sio)
            if "autoboot" in line or "any key" in line:
                print("Canceling autoboot...")
                sio.write("\r")
                flush_n_read(sio)
                print("Setting to boot from USB...")
                sio.write("fastboot usb 0\r")
                flush_n_read(sio)
                break
            elif "Login incorrect" in line:
                sys.exit(1)
        except OSError as e:
            raise e
    return True


def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:], "", ["help", "dev", "force=", "mem256", "loadinterface", "ucmd", "yes",
                                                      "get-serial", "set-serial=", "override-serial="])
    except getopt.GetoptError as err:
        print(err)
        usage()
    dev = False
    mem256 = False
    force = False
    usb_port = None
    load_interface = False
    yes = False
    set_serial_a = None
    do_get_serial = False
    override_serial_a = None
    for o, a in opts:
        if o in ("--help",):
            usage()
        elif o in ("--dev",):
            dev = True
        elif o in ("--force",):
            force = True
            if not windows:
                usb_port = os.path.join("/dev", a)
            else:
                usb_port = a

        elif o in ("--mem256",):
            mem256 = True
        elif o in ("--loadinterface",):
            load_interface = True
        elif o in ("--yes",):
            yes = True
        elif o in ("--get-serial",):
            do_get_serial = True
        elif o in ("--set-serial",):
            set_serial_a = a
        elif o in ("--override-serial",):
            override_serial_a = a
        elif o in ("--ucmd",):
            run_ucmd(args)
            sys.exit(0)
        else:
            assert False, f"unhandled option {o}"

    if force:
        print("Forcing usb boot on next startup...")
        setup_serial(usb_port)
        if not force_usb_boot():
            print("Forced usb boot failed, aborting.")
            sys.exit(0)

    if do_get_serial:
        get_serial()
        sys.exit(0)

    if set_serial_a is not None:
        set_serial(set_serial_a, True, yes)
        sys.exit(0)

    if override_serial_a is not None:
        set_serial(override_serial_a, False, yes)
        sys.exit(0)

    if windows:
        uuu = "uuu.exe"
    else:
        uuu = "uuu"


    imgfile = ""
    ubootfile = ""
    if dev:
        imgfile = "tmp/deploy/images/imx6ull14x14evk/evse-image-imx6ull14x14evk.wic.bz2"
        ubootfile = "tmp/deploy/images/imx6ull14x14evk/u-boot.imx-emmc"
    else:
        if len(args) < 1:
            usage()
        img = tarfile.open(name=args[0], errorlevel=2)
        tdir = tempfile.TemporaryDirectory(suffix=".img")
        for name in img.getnames():
            if name.startswith("evse-image-"):
                imgfile = os.path.join(tdir.name, name)
            elif name == "u-boot.imx-emmc":
                ubootfile = os.path.join(tdir.name, name)
        img.extractall(tdir.name)

    if os.path.basename(imgfile).startswith("evse-image-imx6"):
        if not force:
            print("Activate Fastboot on EVSE using debug console:")
            print("* Reboot")
            print("* Press space")
            print("* Type: fastboot usb 0")
    elif os.path.basename(imgfile).startswith("evse-image-"):
        if os.name == "nt":
            rpi = os.path.join(os.getenv("PROGRAMFILES(x86)"), "Raspberry Pi Imager", "rpi-imager.exe")
            proc = subprocess.run([rpi, imgfile], shell=True)
            if proc.returncode:
                print()
                print("For this machine, this tool cannot automatically transfer the image to target flash.")
                print("Please transfer the following file using suitable tool:")
                print()
                print(imgfile)
                print()
                input("Press Enter to exit.")
            sys.exit(proc.returncode)
    else:
        sys.exit("Error: No file system found in image")

    print("Downloading u-boot and filesystem...")
    sudo_command([uuu, "-b", "emmc_all", ubootfile, imgfile])

    print("Setting MMC bootbus...")
    sudo_command([uuu, "FB:", "ucmd", "mmc", "bootbus", "1", "1", "0", "0"])

    # Use gpio to trigger the hardware watchdog to perform a power-cycle reset.
    # The reason being that when using the built-in 'reset' command in u-boot it makes the 'reboot'
    # command in Linux fail and hang until the watchdog triggers and power-cycles it anyway.
    sudo_command([uuu, "FB:", "ucmd", "setenv", "trigger_watchdog", "gpio set GPIO3_4; sleep 1; gpio clear GPIO3_4"])
    sudo_command([uuu, "FB:", "ucmd", "saveenv"])

    if mem256:
        print("Limiting memory to 256 MB...")
        sudo_command([uuu, "FB:", "ucmd", "setenv", "mmcargs", "'setenv bootargs console=\${console},\${baudrate} root=\${mmcroot} mem=256M'"])

    if load_interface:
        print("Using Load Interface device-tree. Setting dtb to imx6ull-garo-nextgen-li.dtb...")
        sudo_command([uuu, "FB:", "ucmd", "setenv", "fdt_file", "imx6ull-garo-nextgen-li.dtb"])

    if mem256 or load_interface:
        sudo_command([uuu, "FB:", "ucmd", "saveenv"])

    print("Restarting device...")
    sudo_command([uuu, "FB:", "acmd", "run", "trigger_watchdog"])


if __name__ == "__main__":
    main()
