#!/usr/bin/env python3
"""A fake SMOD imx6 that speaks the FPGA remote-update protocol on TCP 1234.

Mirrors ``fpga_upgrade_process()`` in tcp_task.c as literally as possible, including
the parts that bite: the ':'-based re-splitting of our payload, the 11..50 byte line
limit, and the READY-then-ABORT_SUCCESS handshake at the end.

The point is to be able to exercise fpga_ru.py end-to-end without erasing a real
flash. If a change to fpga_ru.py passes against this and against the real board,
the two agree; if it passes here but not there, this file is the thing that's wrong.

    python mock_device.py --port 12340
    python fpga_ru.py --mcs some_update.mcs --host 127.0.0.1 --port 12340 --yes

Fault injection, to prove the client's retry/abort paths actually work:
    --not-ready-every N   answer NOT_READY to every Nth data frame
    --erase-polls N       answer ERASER_ING to the first N POLL_ERASER frames
    --abort-at-frame N    answer ABORT_ERROR on the Nth data frame
"""

import argparse
import socket
import sys

HEAD, TAIL = 0xAA, 0xBB
UPGRADE_START, POLL_ERASER, W_DATA_BLOCK, W_DATA_BLOCK_LAST = 0x01, 0x02, 0x03, 0x04
START_OK, ERASER_ING, ERASER_END = 0x01, 0x02, 0x03
READY, NOT_READY, ABORT_SUCCESS, ABORT_ERROR = 0x04, 0x05, 0x06, 0x07

LINE_MIN, LINE_MAX = 11, 50
MAX_DATA_LEN = 544


def reply(ack):
    tx = bytes([HEAD, ack, 1, 0, ack])
    return tx + bytes([sum(tx) & 0xFF, TAIL])


def split_lines_like_c(payload):
    """Reproduce the strchr(':') walk in W_DATA_BLOCK.

    The device never looks at newlines: it finds the first ':', then each following
    ':' ends the previous line. The final line runs to the NUL terminator, which is
    why a payload that ends mid-line silently mis-measures that line.
    """
    out = []
    start = payload.find(b":")
    if start < 0:
        return None
    while True:
        nxt = payload.find(b":", start + 1)
        out.append(payload[start:] if nxt < 0 else payload[start:nxt])
        if nxt < 0:
            return out
        start = nxt


def parse_line(line):
    """recv_upgrade_data_parse(): -2 bad length, -1 EOF record, 16 data bytes, 0 control."""
    if not (LINE_MIN <= len(line) <= LINE_MAX):
        return -2
    length = int(line[1:3], 16)
    rectype = int(line[7:9], 16)
    if length == 0x10 and rectype == 0x00:
        return 16
    if length == 0x00 and rectype == 0x01:
        return -1
    return 0


class Mock:
    def __init__(self, args):
        self.args = args
        self.bytes_written = 0
        self.data_frames = 0
        self.erase_polls = 0
        self.file_end = False
        self.done_polls = 0

    def on_frame(self, cmd, payload):
        if cmd == UPGRADE_START:
            self.bytes_written = 0
            self.erase_polls = 0
            self.file_end = False
            return START_OK

        if cmd == POLL_ERASER:
            self.erase_polls += 1
            return ERASER_ING if self.erase_polls <= self.args.erase_polls else ERASER_END

        if cmd == W_DATA_BLOCK:
            self.data_frames += 1
            if self.args.abort_at_frame and self.data_frames == self.args.abort_at_frame:
                return ABORT_ERROR
            if self.args.not_ready_every and self.data_frames % self.args.not_ready_every == 0:
                return NOT_READY  # client must re-send this exact frame

            lines = split_lines_like_c(payload)
            if lines is None:
                return ABORT_ERROR
            ack = READY
            for line in lines:
                n = parse_line(line)
                if n == -2:
                    print(f"  [mock] ABORT: bad line ({len(line)} bytes) {line[:20]!r}", file=sys.stderr)
                    return ABORT_ERROR
                if n == -1:
                    self.file_end = True
                    ack = READY  # check_fpga_upgrade_done() is not DONE yet
                elif n > 0:
                    self.bytes_written += n
            return ack

        if cmd == W_DATA_BLOCK_LAST:
            self.done_polls += 1
            if self.done_polls < self.args.done_polls:
                return READY
            return ABORT_SUCCESS if self.file_end else ABORT_ERROR

        return ABORT_ERROR

    def serve_one(self, conn):
        buf = b""
        while True:
            data = conn.recv(MAX_DATA_LEN)
            if not data:
                return
            buf += data
            # The real device does a single recv() per frame; we tolerate splits so a
            # client bug shows up as a protocol error rather than a hang.
            while len(buf) >= 6:
                datalen = int.from_bytes(buf[2:4], "little")
                total = datalen + 6
                if len(buf) < total:
                    break
                frame, buf = buf[:total], buf[total:]
                if frame[0] != HEAD or frame[-1] != TAIL:
                    print(f"  [mock] bad framing {frame[:4].hex(' ')}", file=sys.stderr)
                    return
                if sum(frame[:datalen + 4]) & 0xFF != frame[datalen + 4]:
                    conn.sendall(reply(NOT_READY))
                    continue
                conn.sendall(reply(self.on_frame(frame[1], frame[4:4 + datalen])))


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--port", type=int, default=12340)
    ap.add_argument("--erase-polls", type=int, default=2)
    ap.add_argument("--done-polls", type=int, default=1)
    ap.add_argument("--not-ready-every", type=int, default=0)
    ap.add_argument("--abort-at-frame", type=int, default=0)
    args = ap.parse_args()

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("127.0.0.1", args.port))
    srv.listen(1)
    print(f"[mock] listening on 127.0.0.1:{args.port}", flush=True)

    conn, _ = srv.accept()
    mock = Mock(args)
    try:
        mock.serve_one(conn)
    finally:
        conn.close()
    mib = mock.bytes_written / (1 << 20)
    print(f"[mock] {mock.data_frames} data frames, {mock.bytes_written} bytes ({mib:.3f} MiB), "
          f"file_end={mock.file_end}")
    # A correct run streams exactly the 8 MiB update image.
    sys.exit(0 if mock.file_end else 1)


if __name__ == "__main__":
    main()