#!/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(1)

###########
# 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="am62pxx-power", description="AM62pxx Power controler"
    )

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

    parser.add_argument("action", choices=["on", "off"], help="Power status")

    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")
    grp.add_argument("--xds", type=int, required=True, help="Port for xds110")

    # Supported boot modes for [am62pxx-sk]:
    #   mmc - 0243
    #   uart - 003B
    #   emmc - 004B
    #   ospi - 0273
    #   qspi - 0273
    #   usbmsc - 0153
    #   spldfu - 0053
    grp.add_argument("--sysboot", type=str, default="003B", help="Boot media")

    return parser


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

    print("Power-on carrier board")
    laacli_usb(options.power, "on")
    print("Power-on the UART port")
    laacli_usb(options.serial, "on")
    print("Power-on the XDS port")
    laacli_usb(options.xds, "on")

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

    with serial.Serial("/dev/ttyACM1", baudrate=115200) as ser:
        ser.write(b"\r\n")
        expect(ser, "=>\n")
        expect(ser, "=>\n")
        ser.write(b"auto set dut am62pxx-sk\r\n")
        expect(ser, "=>\n")
        expect(ser, "=>\n")
        ser.write(f"auto sysboot {options.sysboot}\r\n".encode("utf-8"))
        expect(ser, "=>\n")
        expect(ser, "=>\n")
        if options.action == "on":
            ser.write(b"auto por\r\n")
            expect(ser, "=>\n")
            expect(ser, "=>\n")
            ser.write(b"auto power on\r\n")
            expect(ser, "=>\n")
            expect(ser, "=>\n")
        else:
            ser.write(b"auto power off\r\n")
            expect(ser, "=>\n")
            expect(ser, "=>\n")

    if options.action == "off":
        print("Power-off carrier board")
        laacli_usb(options.power, "off")
        print("Power-off the UART port")
        laacli_usb(options.serial, "off")
        print("Power-off the XDS port")
        laacli_usb(options.xds, "off")



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