import argparse
import serial
import struct
import sys
import threading
import time

offset = 0x80000000
tty = "/dev/ttyUSB1"
baud = 115200
chunksize = 1024


def write(s, offset, dat):
    for i in range(0, len(dat), chunksize):
        chunk = dat[i : i + chunksize]
        cmd = struct.pack("<cII", b"c", offset + i, len(chunk))

        print(f"Sending {len(chunk)} bytes @0x{offset + i:04x}")

        s.write(cmd)
        s.write(chunk)

        ack = s.read(2)
        if len(ack) < 2:
            raise RuntimeError(f"timeout waiting for full ack. got {ack}")
        if ack[0] != b"a"[0]:
            raise RuntimeError(f"expected ack, got this instead: {ack}")
        print(f"Ack! len={ack[1]}")


def jump(s, offset):
    cmd = struct.pack("<cI", b"j", offset)
    print(f"Jumping to 0x{offset:04x}")
    s.write(cmd)


def stream_logs(s):
    while True:
        dat = s.read()
        if not dat:
            continue
        sys.stdout.buffer.write(dat.replace(b"\r", b""))
        sys.stdout.buffer.flush()


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("binary")
    parser.add_argument(
        "--monitor",
        action="store_true",
        help="wait for and display program serial output",
    )
    return parser.parse_args()


def main():
    args = parse_args()

    with open(args.binary, "rb") as f:
        dat = f.read()

    s = serial.Serial(tty, baud, timeout=1)
    write(s, offset, dat)
    jump(s, offset)

    if args.monitor:
        t = threading.Thread(target=lambda: stream_logs(s), daemon=True)
        t.start()

        try:
            while True:
                dat = input("") + "\r"
                s.write(dat.encode())
        except KeyboardInterrupt:
            print("Bye.")


if __name__ == "__main__":
    main()