#!/usr/bin/env python3
"""Remote-update the SMOD FPGA by streaming a quickboot _update.mcs to the imx6.

This is a Python reimplementation of what ``tools/Xfpga_ru.exe`` (a .NET GUI, no CLI)
does. The device side is ``my_exciter/imx6/src/atsc_app/task/tcp_task.c`` ::
``fpga_upgrade_task``, listening on TCP 1234. It parses the MCS text we send and
pokes the bytes into the FPGA's SpiFlashProgrammer, which erases and writes the
update half of the SPI flash by itself.

This has nothing to do with the web page's "Remote Update" box -- that one runs
check_package.sh, which only accepts *.tar.gz imx6 firmware bundles and deletes
anything else it finds. See references/protocol.md.

Wire format (all values from tcp_task.c / inc/tcp_task.h)::

    request   AA | cmd | datalen:u16le | payload | sum | BB
    response  AA | ack | 01 | 00       | ack     | sum | BB     (7 bytes)
    sum       8-bit sum of every byte from the head up to (not incl.) sum

Usage:
    python fpga_ru.py --mcs .../bitstream/atsc_exciter_update.mcs
    python fpga_ru.py --mcs ... --host 192.168.1.143 --yes
    python fpga_ru.py --mcs ... --dry-run          # parse + frame, touch no hardware
"""

import argparse
import os
import re
import socket
import sys
import time

# ------------------------------------------------------------------ protocol constants

HEAD, TAIL = 0xAA, 0xBB

UPGRADE_START, POLL_ERASER, W_DATA_BLOCK, W_DATA_BLOCK_LAST = 0x01, 0x02, 0x03, 0x04

ACK = {
    0x01: "START_OK",
    0x02: "ERASER_ING",
    0x03: "ERASER_END",
    0x04: "READY",
    0x05: "NOT_READY",
    0x06: "ABORT_SUCCESS",
    0x07: "ABORT_ERROR",
}
START_OK, ERASER_ING, ERASER_END = 0x01, 0x02, 0x03
READY, NOT_READY, ABORT_SUCCESS, ABORT_ERROR = 0x04, 0x05, 0x06, 0x07

# rxdata[] on the device is MAX_DATA_LEN = 544 bytes and holds the *whole* frame,
# so the payload can never exceed 544 - 6 bytes of framing. Going over this silently
# truncates the frame and the tail check fails.
MAX_FRAME = 544
FRAME_OVERHEAD = 6
MAX_PAYLOAD = MAX_FRAME - FRAME_OVERHEAD  # 538

# recv_upgrade_data_parse() rejects any line outside this range with -2 -> ABORT_ERROR.
LINE_MIN, LINE_MAX = 11, 50

UPDATE_START_ADDR = 0x800000
ELA_RE = re.compile(rb"^:02000004([0-9A-Fa-f]{4})")
EOF_RECORD = b":00000001FF"


# ---------------------------------------------------------------------------- framing

def build(cmd, payload=b""):
    body = bytes([HEAD, cmd]) + len(payload).to_bytes(2, "little") + payload
    return body + bytes([sum(body) & 0xFF, TAIL])


def parse_ack(resp):
    if len(resp) != 7 or resp[0] != HEAD or resp[6] != TAIL:
        raise ProtocolError(f"malformed response: {resp.hex(' ')}")
    return resp[4]


class ProtocolError(RuntimeError):
    pass


class Device:
    def __init__(self, host, port, timeout):
        self.sock = socket.create_connection((host, port), timeout=timeout)
        # Every frame is a request/response round trip; Nagle would just add latency.
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

    def xfer(self, cmd, payload=b"", timeout=None):
        if timeout is not None:
            self.sock.settimeout(timeout)
        self.sock.sendall(build(cmd, payload))
        try:
            resp = self.sock.recv(16)
        except socket.timeout:
            raise ProtocolError(
                f"no response to cmd 0x{cmd:02X} within {self.sock.gettimeout()}s.\n"
                "  The app returns an empty reply (and we block) when sys_ready==0 or the\n"
                "  frame was malformed. Check the app is up: ssh root@<host> 'ps | grep atsc'"
            )
        if not resp:
            raise ProtocolError("device closed the connection")
        return parse_ack(resp)

    def close(self):
        try:
            self.sock.close()
        except OSError:
            pass


# ------------------------------------------------------------------------ mcs handling

def load_chunks(path, max_payload, allow_any_image):
    """Split the MCS into whole-line payloads.

    Lines must stay whole: the device re-splits our payload on ':' and measures each
    line's length itself, so a line straddling two frames is parsed as two malformed
    lines and aborts the upgrade.
    """
    with open(path, "rb") as fh:
        lines = fh.readlines()
    if not lines:
        raise SystemExit(f"[FAIL] {path} is empty")

    first_ela = None
    for line in lines:
        m = ELA_RE.match(line)
        if m:
            first_ela = int(m.group(1), 16) << 16
            break
    if first_ela != UPDATE_START_ADDR and not allow_any_image:
        got = "none" if first_ela is None else f"0x{first_ela:X}"
        raise SystemExit(
            f"[FAIL] {os.path.basename(path)} starts at {got}, expected 0x{UPDATE_START_ADDR:X}.\n"
            "       That is not a quickboot _update.mcs. Uploading a base or _initial mcs\n"
            "       writes garbage into the update slot; the board then falls back to golden\n"
            "       on the next cold boot. Regenerate with scripts/gen_update_mcs.py.\n"
            "       (Override with --allow-any-image if you really mean it.)"
        )

    for n, line in enumerate(lines, 1):
        if not line.startswith(b":"):
            raise SystemExit(f"[FAIL] line {n} does not start with ':' -- not an Intel-HEX mcs")
        if not (LINE_MIN <= len(line) <= LINE_MAX):
            raise SystemExit(
                f"[FAIL] line {n} is {len(line)} bytes (incl EOL); the device only accepts "
                f"{LINE_MIN}..{LINE_MAX} and aborts otherwise"
            )
    if not lines[-1].startswith(EOF_RECORD):
        raise SystemExit(f"[FAIL] last line is {lines[-1]!r}, expected the EOF record {EOF_RECORD.decode()}")

    chunks, cur = [], []
    size = 0
    for line in lines:
        if size + len(line) > max_payload and cur:
            chunks.append(b"".join(cur))
            cur, size = [], 0
        cur.append(line)
        size += len(line)
    if cur:
        chunks.append(b"".join(cur))
    return chunks, len(lines)


# ------------------------------------------------------------------------------ stages

def do_erase(dev, timeout):
    print("[2/4] erasing the update image region (this is the slow part)")
    ack = dev.xfer(UPGRADE_START, timeout=30)
    if ack != START_OK:
        raise ProtocolError(f"UPGRADE_START -> {ACK.get(ack, hex(ack))}, expected START_OK")

    deadline = time.monotonic() + timeout
    dots = 0
    while time.monotonic() < deadline:
        ack = dev.xfer(POLL_ERASER, timeout=30)
        if ack == ERASER_END:
            print("\n      erase complete")
            return
        if ack != ERASER_ING:
            raise ProtocolError(f"POLL_ERASER -> {ACK.get(ack, hex(ack))}")
        dots += 1
        print(f"\r      erasing... {dots * 0.5:.0f}s", end="", flush=True)
        time.sleep(0.5)
    raise ProtocolError(
        f"flash still erasing after {timeout}s.\n"
        "  POLL_ERASER only ever reports ERASER_ING/ERASER_END -- the app never checks the\n"
        "  ERROR bit, so a programmer that aborted looks identical to one that is still busy.\n"
        "  Read the real status on the device:\n"
        "    ssh root@<host> \"grep 'ing|end' /var/module_log/imx_atsc_app.log | tail -3\"\n"
        "  rstate=0x7 (ERROR|DONE, EraseOK=0) means SpiFlashProgrammer aborted immediately --\n"
        "  the running design cannot reach the flash. See references/protocol.md.")


def do_write(dev, chunks, total_lines, retries):
    print(f"[3/4] streaming {total_lines} mcs lines in {len(chunks)} frames")
    sent_lines = 0
    t0 = time.monotonic()
    for i, chunk in enumerate(chunks):
        for attempt in range(retries + 1):
            ack = dev.xfer(W_DATA_BLOCK, chunk, timeout=60)
            if ack == READY:
                break
            if ack == NOT_READY:
                # Device asks us to re-send this exact block after a short pause.
                time.sleep(0.001)
                continue
            if ack in (ABORT_SUCCESS, ABORT_ERROR):
                return ack
            raise ProtocolError(f"W_DATA_BLOCK -> {ACK.get(ack, hex(ack))}")
        else:
            raise ProtocolError(f"frame {i} still NOT_READY after {retries} retries")

        sent_lines += chunk.count(b"\n")
        if i % 2000 == 0 or i == len(chunks) - 1:
            pct = 100.0 * sent_lines / total_lines
            rate = sent_lines / max(time.monotonic() - t0, 1e-6)
            eta = (total_lines - sent_lines) / max(rate, 1e-6)
            print(f"\r      {pct:5.1f}%  {sent_lines}/{total_lines} lines  ETA {eta/60:.1f} min",
                  end="", flush=True)
    print()
    return None


def do_finish(dev, timeout):
    print("[4/4] waiting for the FPGA to report the write is done")
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            ack = dev.xfer(W_DATA_BLOCK_LAST, timeout=30)
        except (ProtocolError, ConnectionResetError, OSError) as exc:
            # On success the app calls reset_fpga_system() and reboots the imx6, so the
            # socket dying right here is expected -- but we can no longer read the ack.
            raise ProtocolError(f"link dropped while waiting for the final ack: {exc}")
        if ack in (ABORT_SUCCESS, ABORT_ERROR):
            return ack
        if ack != READY:
            raise ProtocolError(f"W_DATA_BLOCK_LAST -> {ACK.get(ack, hex(ack))}")
        time.sleep(0.2)
    raise ProtocolError(f"no final verdict after {timeout}s")


# -------------------------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--mcs", required=True, help="path to the quickboot <tag>_update.mcs")
    ap.add_argument("--host", default="192.168.1.143")
    ap.add_argument("--port", type=int, default=1234)
    ap.add_argument("--yes", action="store_true", help="skip the confirmation prompt")
    ap.add_argument("--dry-run", action="store_true", help="validate and frame the mcs, contact nothing")
    ap.add_argument("--allow-any-image", action="store_true",
                    help="skip the 0x800000 start-address check (you almost never want this)")
    ap.add_argument("--max-payload", type=int, default=MAX_PAYLOAD)
    ap.add_argument("--erase-timeout", type=float, default=600)
    ap.add_argument("--finish-timeout", type=float, default=120)
    ap.add_argument("--retries", type=int, default=200)
    args = ap.parse_args()

    if not 0 < args.max_payload <= MAX_PAYLOAD:
        raise SystemExit(f"[FAIL] --max-payload must be 1..{MAX_PAYLOAD} (device rxdata is {MAX_FRAME} bytes)")

    print(f"[1/4] validating {args.mcs}")
    chunks, total_lines = load_chunks(args.mcs, args.max_payload, args.allow_any_image)
    mb = os.path.getsize(args.mcs) / 1e6
    print(f"      {mb:.1f} MB, {total_lines} lines, {len(chunks)} frames, starts at 0x{UPDATE_START_ADDR:X}")

    if args.dry_run:
        print("[dry-run] mcs is well-formed and would stream cleanly. Nothing was sent.")
        return

    if not args.yes:
        print(f"\nAbout to erase and rewrite the update image on {args.host}.")
        print("The golden image is untouched, so a failure here is recoverable by cold boot.")
        print("Do NOT power the board off while this runs.")
        print("\nMake sure DPD is OFF. The exciter's DPD sample&load writes to the same SPI1 bus\n"
              "we stream the image over, and the FPGA's spi_slave32 will happily fold those words\n"
              "into the flash image. See references/protocol.md.")
        if input("Continue? [y/N] ").strip().lower() not in ("y", "yes"):
            raise SystemExit("aborted")

    dev = Device(args.host, args.port, timeout=30)
    try:
        do_erase(dev, args.erase_timeout)
        verdict = do_write(dev, chunks, total_lines, args.retries)
        if verdict is None:
            verdict = do_finish(dev, args.finish_timeout)
    finally:
        dev.close()

    if verdict == ABORT_SUCCESS:
        print("\n[OK] FPGA reports the update image was written and verified.")
        print("The app has reset the FPGA and is rebooting the imx6.")
        print("Now cold-boot the board (observe the SMOD power-up sequence) to load the new image.")
        return
    raise SystemExit(
        f"\n[FAIL] device returned {ACK.get(verdict, hex(verdict))}.\n"
        "The update slot is now half-written, but the switch word was erased at the start,\n"
        "so the board still cold-boots golden. Nothing is bricked.\n\n"
        "If this was ABORT_ERROR late in the stream, suspect SPI1 contention: the exciter's\n"
        "DPD sample&load shares fpga_spi_fd with this upload, and the FPGA folds any SPI1 word\n"
        "into the image. Turn DPD off and retry. references/protocol.md has the full story.")


if __name__ == "__main__":
    main()