#!/usr/bin/env python3.7
"""Garo Entity downloader"""
#
# -*-mode: python; coding: utf-8 -*-
#
# Copyright 2020-2021 Garo AB

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

windows = os.name == "nt"

def usage():
    "Exit with usage instructions"
    sys.exit(f"""Usage: {sys.argv[0]} <image-file>|--dev [options]
options:
    --dev                          Use image from development tree
    --force=<serial port>          Force fastboot using specified USB port. E.g. 'COM7' on Windows or 'ttyUSB0' in Linux.
    --force-password=<password>    U-boot password to use
    --mem256                       Limit memory to 256MB.
    --loadinterface                Load Interface, same as --model=li
    --compact                      Compact, same as --model=compact
    --model=<model>                Set device tree model
                                   For Spin3, use --model=pro3
    --get-serial                   Get serial number
    --set-serial=<hex-serial>      WARNING: Set serial number (MACs) permanently
                                   <hex-serial> should be MAC address ending with zero
    --yes                          Confirm operations such as --set-serial
    --override-serial=<hex-serial> Temporarily set serial number
    --ucmd <args> ...              Run u-boot ucmd command
    --blank                        Erase uboot - return to blank board - then exit
    --reset                        Reset system
    --download                     Download image to eMMC
    --usb-match=<usb>              USB match, repeat for multiple matches

example:
    entity-downloader garo-entity-0.1.0.img --force=COM14 --force-password=changeme
    """)


# 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
    print("Running", args)
    return run_command(args, *popenargs, **kwargs)


def run_ucmd(uuu, 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(uuu):
    "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

    # Need to restore or uboot will fail with various errors on next operation
    sudo_command([*uuu, "FB:", "ucmd", "setenv", "stdout", "serial"])

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


def set_serial(uuu, sernum, permanent, yes):
    "Set serials and MACs"
    sernum = int(sernum, 16)
    macwords = get_macwords(uuu)
    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("Warning: Cannot derive MAC for if1 - skipping ", file=sys.stderr)

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

    # U-boot in firmware < 2.0 does not set these variables from fuses, so set them
    # manually. FIXME in the future.
    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):
    "Flush serial port, then read line"
    import serial # pylint: disable=C0415
    sio.flush()
    try:
        line = sio.readline()
    except serial.serialutil.SerialException as exc:
        print("warning:", exc)
        return ""
    line = line.rstrip()
    if line != "":
        print(line, file=sys.stderr)
    return line


def setup_serial(serialport):
    "Setup serial port, return sio object"
    import serial # pylint: disable=C0415
    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: # pylint: disable=W0703
        sys.exit("Error setting up serial port: ", exc)
    return sio

def force_fastboot(uuu, sio, ubootfile, password):
    "Enter fastboot mode"
    if sio is None:
        print("Serial port not open, aborting.", file=sys.stderr)
        return False

    linetime = time.time() # Timestamp when we got a line
    print("Waiting for u-boot...", file=sys.stderr)
    while True:
        line = flush_n_read(sio)
        line = line.strip(chr(8)).strip() # Remove backspace and whitespace
        if line != "":
            linetime = time.time()
        if line.startswith("Hit any key"):
            print("Canceling autoboot...", file=sys.stderr)
            sio.write("\r")
            flush_n_read(sio)
            print("Starting fastboot...", file=sys.stderr)
            sio.write("fastboot usb 0\r")
            flush_n_read(sio)
            return True
        if line.startswith("Autoboot in ") or line in ("5", "4", "3", "2", "1", "0"):
            print("Canceling autoboot...", file=sys.stderr)
            sio.write(password + "\r")
            flush_n_read(sio)
            print("Starting fastboot...", file=sys.stderr)
            sio.write("fastboot usb 0\r")
            flush_n_read(sio)
            return True
        if "Login incorrect" in line:
            sys.exit(1)
        elif time.time() - linetime > 2:
            # Check if fastboot is already running
            if check_fb_sdp(uuu, ubootfile):
                print("fastboot is running", file=sys.stderr)
                return True
            serial_login(sio)
            print("Waiting for u-boot...", file=sys.stderr)
            linetime += 60 # Retry in one minute


def serial_login(sio):
    "Login to device over serial port"
    print("Logging in to device...", file=sys.stderr)
    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...", file=sys.stderr)
    sio.write("\r")
    flush_n_read(sio)
    sio.write("sudo /sbin/reboot\r")
    flush_n_read(sio)


def get_img_files(firstarg):
    "Arrange uboot and image files"
    imgfile = ""
    ubootfile = ""
    if firstarg == "--dev":
        imgfile = "tmp/deploy/images/imx6ull14x14evk/evse-image-imx6ull14x14evk.wic.bz2"
        ubootfile = "tmp/deploy/images/imx6ull14x14evk/u-boot.imx"
        if not os.path.exists("tmp/deploy/images/imx6ull14x14evk/dcd-u-boot.imx"):
            print("Warning: Using unsigned u-boot")
        tdir = None
    else:
        img = tarfile.open(name=firstarg, 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 in ("u-boot.imx-emmc", "u-boot.imx"):
                ubootfile = os.path.join(tdir.name, name)
        img.extractall(tdir.name)
    # Make sure files actually exists
    for fname in (imgfile, ubootfile):
        if not os.path.exists(fname):
            sys.exit(f"{fname} does not exist")
    return tdir, imgfile, ubootfile


def check_fb_sdp(uuu, ubootfile):
    "Check if we can communicate with FB or SDP protocol"
    # First check if fastboot is already active
    ret = sudo_command([*uuu, "-t", "2", "FB:", "ucmd"],
                       check=False, stdout=subprocess.DEVNULL)
    if ret.returncode != 0:
        # Try SDP protocol with blank board, must use DCD version
        head, tail = os.path.split(ubootfile)
        dcd_ubootfile = os.path.join(head, "dcd-" + tail)
        if os.path.exists(dcd_ubootfile):
            ubootfile = dcd_ubootfile
        else:
            print("Warning: DCD version of u-boot not available")
        ret = sudo_command([*uuu, "-t", "2", "SDP:", "boot", "-f", ubootfile],
                           check=False, stdout=subprocess.DEVNULL)
        time.sleep(2)
    return ret.returncode == 0


def arrange_fastboot(uuu, ubootfile, force_port, force_password):
    "Arrange fastboot communication"
    if force_port is not None:
        print("Trying to enter fastboot mode...", file=sys.stderr)
        sio = setup_serial(force_port)
        if not force_fastboot(uuu, sio, ubootfile, force_password):
            print("Failed to enter fastboot mode, aborting.", file=sys.stderr)
            sys.exit(0)
    elif not check_fb_sdp(uuu, ubootfile):
        print("Attach blank board", file=sys.stderr)
        print(" or ", file=sys.stderr)
        print("Activate Fastboot on EVSE using debug console:", file=sys.stderr)
        print("* Reboot", file=sys.stderr)
        print("* Type password", file=sys.stderr)
        print("* Type: fastboot usb 0", file=sys.stderr)
        while not check_fb_sdp(uuu, ubootfile):
            time.sleep(1)


def main():
    "Main function"
    sys.stdout.reconfigure(line_buffering=True)
    sys.stderr.reconfigure(line_buffering=True)
    # An image or --dev is mandatory, since we may need to load
    # uboot
    try:
        firstarg = sys.argv[1]
        if firstarg != "--help":
            del sys.argv[1]
    except IndexError:
        usage()
    try:
        opts, args = getopt.getopt(sys.argv[1:], "", ["help", "download", "force=", "mem256", "loadinterface", "compact", "ucmd", "yes",
                                                      "force-password=", "reset",
                                                      "blank", "get-serial", "set-serial=", "override-serial=", "usb-match=", "model="])
    except getopt.GetoptError as err:
        print(err, file=sys.stderr)
        usage()
    do_download = False
    do_mem256 = False
    force_port = None # If not None, serial port for --force
    force_password = " "
    model = None
    yes = False
    set_serial_a = None
    do_get_serial = False
    override_serial_a = None
    do_ucmd = False
    do_blank = False
    do_reset = False

    if windows:
        uuu = ["uuu.exe"]
    else:
        uuu = ["./uuu"]

    for o, a in opts:
        if o in ("--help",):
            usage()
        elif o in ("--download",):
            do_download = True
        elif o in ("--force",):
            if not windows:
                force_port = os.path.join("/dev", a)
            else:
                force_port = a
        elif o in ("--force-password",):
            force_password = a
        elif o in ("--mem256",):
            do_mem256 = True
        elif o in ("--loadinterface",):
            model = "li"
        elif o in ("--compact",):
            model = "compact"
        elif o in ("--model",):
            model = a
        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 ("--reset",):
            do_reset = True
        elif o in ("--override-serial",):
            override_serial_a = a
        elif o in ("--ucmd",):
            do_ucmd = True
            ucmd_args = args
        elif o in ("--blank",):
            do_blank = True
        elif o in ("--usb-match",):
            uuu.append("-m")
            uuu.append(a)
        else:
            assert False, f"unhandled option {o}"

    if do_download and model is None:
        print("Error: --download requires --model", file=sys.stderr)
        usage()

    dummy_tdir, imgfile, ubootfile = get_img_files(firstarg)

    if os.path.basename(imgfile).startswith("evse-image-imx6"):
        arrange_fastboot(uuu, ubootfile, force_port, force_password)

    if do_ucmd:
        run_ucmd(uuu, ucmd_args)

    if do_get_serial:
        get_serial(uuu)

    if set_serial_a is not None:
        set_serial(uuu, set_serial_a, True, yes)

    if override_serial_a is not None:
        set_serial(uuu, override_serial_a, False, yes)

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

    if do_blank:
        with tempfile.NamedTemporaryFile(delete=False) as ntf:
            ntf.write(bytes(1024))
            ntf.close()
            sudo_command([*uuu, "-b", "emmc", ntf.name])
            os.remove(ntf.name)
        sys.exit(0)

    if do_download:
        download(uuu, imgfile, ubootfile, model)

    if do_reset:
        reset(uuu)

def download(uuu, imgfile, ubootfile, model):
    "Download image and uboot"
    if os.path.basename(imgfile).startswith("evse-image-imx6"):
        pass
    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...", file=sys.stderr)
    sudo_command([*uuu, "-b", "emmc_all", ubootfile, imgfile])

    print("Setting MMC bootbus...", file=sys.stderr)
    sudo_command([*uuu, "FB:", "ucmd", "mmc", "bootbus", "1", "1", "0", "0"])

    if model == "pro2":
        # Legacy
        dtb = "imx6ull-garo-nextgen.dtb"
    else:
        dtb = f"imx6ull-garo-nextgen-{model}.dtb"
    print(f"Setting device tree model to {dtb}...", file=sys.stderr)

    # We are not clearing the environment because U-boot in firmware < 2.0 does not set
    # ethaddr from fuses. FIXME in the future.
    #sudo_command([*uuu, "FB:", "ucmd", "env", "default", "-a"])
    sudo_command([*uuu, "FB:", "ucmd", "setenv", "fdt_file", dtb])
    sudo_command([*uuu, "FB:", "ucmd", "saveenv"])

    print("Restarting device...", file=sys.stderr)


def reset(uuu):
    "Reset system"
    # 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", "gpio", "set", "GPIO3_4"])
    time.sleep(1)
    sudo_command([*uuu, "FB:", "acmd", "gpio", "clear", "GPIO3_4"])

if __name__ == "__main__":
    main()
