#!/usr/bin/python3

import argparse
import subprocess
import sys
import time

import serial

############
# Constant #
############
__version__ = "v0.2"
PREFIX = "| "

###########
# LAACLI #
###########
def laacli_usb(port: int, command: str) -> None:
    if port == 0:
        assert command in ["start", "reset"]
    else:
        assert command in ["on", "off"]

    print(f"=> laacli usb {port} {command}")
    subprocess.check_output(["laacli", "usb", str(port), command])
    time.sleep(3)

###########
# Helpers #
###########
def readline(ser):
    return ser.readline().decode("utf-8").replace("\r", "")


def expect(ser, pattern, failure=None):
    while (line := readline(ser)) != pattern:
        if line == pattern:
            break
        if failure is not None and failure in line:
            raise Exception("Not booting to recovery")
        sys.stdout.write(PREFIX + line)
    sys.stdout.write(PREFIX + line)
    assert line == pattern


##############
# CLI parser #
##############
def setup_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="rzg2l-flash", description="RZ/G2L Flashing tool"
    )

    parser.add_argument(
        "--version", action="version", version=f"%(prog)s, {__version__}"
    )

    grp = parser.add_argument_group("Binaries to flash")
    grp.add_argument(
        "--bl2", type=argparse.FileType("rb"), required=True, help="BL2 binary"
    )
    grp.add_argument(
        "--fip", type=argparse.FileType("rb"), required=True, help="FIP image"
    )
    grp.add_argument(
        "--flash",
        type=argparse.FileType("rb"),
        required=True,
        help="Flash writer binary",
    )

    grp = parser.add_argument_group("USB port numbers")
    grp.add_argument("--power", type=int, required=True, help="Port for power")
    grp.add_argument("--serial", type=int, required=True, help="Port for serial")
    return parser


##############
# Entrypoint #
##############
def main():
    options = setup_parser().parse_args()

    print("Power-off the board", flush=True)
    laacli_usb(options.serial, "off")
    laacli_usb(options.power, "off")
    time.sleep(7)
    laacli_usb(options.serial, "on")

    print("Waiting for the serial", flush=True)
    while True:
        try:
            with serial.Serial("/dev/ttyUSB0", baudrate=115200):
                break
        except Exception:
            print("  retrying", flush=True)
            time.sleep(1)

    with serial.Serial("/dev/ttyUSB0", baudrate=115200) as ser:
        laacli_usb(options.power, "on")
        expect(ser, "please send !\n", failure="U-Boot")

        print(f"Sending Flash driver: {options.flash.name}", flush=True)
        ser.write(options.flash.read())

        expect(ser, " Product Code : RZ/G2L\n")
        ser.write(b"\r\n")
        ser.write(b"SUP\r\n")
        expect(ser, "Please change to 921.6Kbps baud rate setting of the terminal.\n")

    with serial.Serial("/dev/ttyUSB0", baudrate=921600) as ser:
        ser.write(b"\r\n")
        ser.write(b"XLS2\r\n")
        expect(ser, "===== Please Input Program Top Address ============\n")
        ser.write(b"11E00\r\n")
        expect(ser, "===== Please Input Qspi Save Address ===\n")
        ser.write(b"00000\r\n")
        expect(ser, "please send ! ('.' & CR stop load)\n")
        print(f"Sending BL2: {options.bl2.name}", flush=True)
        ser.write(options.bl2.read())
        ser.write(b"y\r\n")
        expect(ser, "===========================================================\n")

        ser.write(b"XLS2\r\n")
        expect(ser, "===== Please Input Program Top Address ============\n")
        ser.write(b"00000\r\n")
        expect(ser, "===== Please Input Qspi Save Address ===\n")
        ser.write(b"20000\r\n")
        expect(ser, "please send ! ('.' & CR stop load)\n")
        print(f"Sending FIP: {options.fip.name}", flush=True)
        ser.write(options.fip.read())
        ser.write(b"y\r\n")
        expect(ser, "===========================================================\n")
        ser.write(b"\r\n")

    print("Power-off the board", flush=True)
    laacli_usb(options.serial, "off")
    laacli_usb(options.power, "off")

    return 0


if __name__ == "__main__":
    sys.exit(main())