"""Interactive UART command shell (sync domain, 24 MHz).

A tiny line-oriented console for bring-up over the FT2232H channel-B UART
(115200 8N1).  Prints a ``rebbarb> `` prompt, echoes typed characters (with
backspace editing), and on Enter parses one command:

    help                     list commands
    udp unicast <ip> [msg]   send a UDP datagram to <ip>:DST_PORT (W5100 ARPs)
    udp broadcast [msg]      send a UDP datagram to 255.255.255.255:DST_PORT

`msg` is optional; when omitted a built-in default payload is sent.  The shell
drives the W5100 master's runtime UDP-send interface (`udp_send_req`,
`udp_dst_ip`, and the `udp_pl_*` payload stream) — it does not touch the MACRAW
BBA path.

Grammar notes (it's a debug shell, not bash):
    * commands and keywords are case-insensitive;
    * tokens are separated by exactly one space;
    * the destination is parsed as a dotted IPv4 literal;
    * everything after the address token (unicast) / keyword (broadcast) is the
      raw payload, preserving case and embedded spaces, to end of line.

On the PC:
    udp unicast   →  nc -ul 6464      (or: sudo tcpdump -n udp port 6464)
    udp broadcast →  nc -ul 6464      (limited broadcast reaches the segment)
"""

from amaranth import *
from amaranth.lib.cdc import FFSynchronizer
from amaranth.lib.memory import Memory

__all__ = ["UARTShell"]

# ── Printable message ROM ──────────────────────────────────────────────────
_MSGS = {
    "BANNER": b"\r\nre-bba-rb ethernet test\r\n",
    "PROMPT": b"rebbarb> ",
    "HELP":   (b"commands:\r\n"
               b"  help                   show this\r\n"
               b"  udp unicast <ip> [msg] send UDP to <ip>\r\n"
               b"  udp broadcast [msg]    send UDP broadcast\r\n"),
    "SENT":   b"sent\r\n",
    "ERR":    b"? (try 'help')\r\n",
    "BADIP":  b"bad ip\r\n",
    "BUSY":   b"busy\r\n",
    "CRLF":   b"\r\n",
    "BKSP":   b"\b \b",
    "RXPFX":  b"rx ",
    "TIMEOUT": b"timeout\r\n",
}
_ORDER = list(_MSGS)
_ROM   = b"".join(_MSGS[k] for k in _ORDER)
_OFF   = {}
_acc = 0
for _k in _ORDER:
    _OFF[_k] = (_acc, _acc + len(_MSGS[_k]))
    _acc += len(_MSGS[_k])

_DEFAULT_PAYLOAD = b"rebbarb-udp-test"

# return codes for the shared print routine
_RET_PROMPT = 0   # after printing, (re)issue the prompt then read input
_RET_INPUT  = 1   # after printing, go straight back to reading input
_RET_PARSE  = 2   # after printing (the CR echo), parse the line
_RET_RXWAIT = 3   # after printing, wait for a UDP reply (with timeout)
_RET_RXBODY = 4   # after printing "rx ", stream the reply payload


class UARTShell(Elaboratable):
    def __init__(self, clk_freq=24_000_000, baud_rate=115_200,
                 default_payload=_DEFAULT_PAYLOAD, lbuf_len=64,
                 reply_timeout_cycles=24_000_000, poll_gap_cycles=24_000):
        self._div      = round(clk_freq / baud_rate)
        self._defpl    = list(default_payload)
        self._lbuf_n   = lbuf_len
        # After a send, wait this many sync cycles for a reply before giving up
        # (default ~1 s at 24 MHz) so the console never hangs; poll the socket
        # RX every `poll_gap_cycles` (~1 ms) in between.
        self._to_cyc   = reply_timeout_cycles
        self._gap_cyc  = poll_gap_cycles

        # UART pins
        self.uart_tx = Signal(init=1)
        self.uart_rx = Signal(init=1)

        # W5100 UDP-send interface (drive the W5100ParallelMaster)
        self.udp_send_req  = Signal()
        self.udp_dst_ip    = Signal(32)
        self.udp_pl_data   = Signal(8)
        self.udp_pl_valid  = Signal()
        self.udp_pl_last   = Signal()
        self.udp_pl_ready  = Signal()
        self.udp_test_busy = Signal()

        # W5100 UDP-receive interface (poll for a reply after each send)
        self.udp_rx_req    = Signal()
        self.udp_rx_busy   = Signal()
        self.udp_rx_none   = Signal()
        self.udp_rx_data   = Signal(8)
        self.udp_rx_valid  = Signal()
        self.udp_rx_eof    = Signal()
        self.udp_rx_ready  = Signal()

    def elaborate(self, platform):
        m = Module()
        div = self._div
        N   = self._lbuf_n

        # Message ROM lives in a block RAM (SB_RAM40_4K) instead of LUTs — the
        # ~230-byte string table was the single biggest LUT consumer (a wide
        # addressed mux).  PRINT already streams one byte per UART-byte-time, so
        # the block RAM's 1-cycle read latency is hidden (one prime cycle on
        # entry).  The tiny default payload stays in logic (cheap, and keeps the
        # payload path combinational — no streaming-latency hazard).
        m.submodules.rom_mem = rom_mem = Memory(
            shape=unsigned(8), depth=len(_ROM), init=list(_ROM))
        rom_rd  = rom_mem.read_port()          # synchronous (block RAM)
        defrom  = Array([Const(b, 8) for b in self._defpl])
        DEF_LEN = len(self._defpl)

        # ── UART TX (8N1) ─────────────────────────────────────────────────
        tx_cnt   = Signal(range(div))
        tx_bits  = Signal(range(11))
        tx_shreg = Signal(10, init=0b1111111111)
        tx_busy  = Signal()
        tx_load  = Signal()
        tx_byte  = Signal(8)
        m.d.comb += [tx_busy.eq(tx_bits != 0), self.uart_tx.eq(tx_shreg[0])]
        with m.If(tx_load & ~tx_busy):
            m.d.sync += [tx_shreg.eq(Cat(Const(0, 1), tx_byte, Const(1, 1))),
                         tx_bits.eq(10), tx_cnt.eq(div - 1)]
        with m.Elif(tx_busy):
            with m.If(tx_cnt == 0):
                m.d.sync += [tx_shreg.eq(Cat(tx_shreg[1:], Const(1, 1))),
                             tx_bits.eq(tx_bits - 1), tx_cnt.eq(div - 1)]
            with m.Else():
                m.d.sync += tx_cnt.eq(tx_cnt - 1)

        # ── UART RX (8N1, mid-bit sampling) ───────────────────────────────
        rx_sync = Signal(init=1)
        rx_prev = Signal(init=1)
        rx_cnt  = Signal(range(div))
        rx_bits = Signal(range(9))
        rx_shr  = Signal(8)
        rx_byte = Signal(8)
        rx_valid = Signal()
        m.submodules += FFSynchronizer(self.uart_rx, rx_sync, init=1)
        m.d.sync += rx_valid.eq(0)
        with m.FSM(name="rx"):
            with m.State("IDLE"):
                m.d.sync += rx_prev.eq(rx_sync)
                with m.If(rx_prev & ~rx_sync):
                    m.d.sync += rx_cnt.eq(div // 2 - 1)
                    m.next = "START"
            with m.State("START"):
                with m.If(rx_cnt == 0):
                    with m.If(~rx_sync):
                        m.d.sync += [rx_cnt.eq(div - 1), rx_bits.eq(7)]
                        m.next = "DATA"
                    with m.Else():
                        m.next = "IDLE"
                with m.Else():
                    m.d.sync += rx_cnt.eq(rx_cnt - 1)
            with m.State("DATA"):
                with m.If(rx_cnt == 0):
                    m.d.sync += [rx_shr.eq(Cat(rx_shr[1:], rx_sync)),
                                 rx_cnt.eq(div - 1)]
                    with m.If(rx_bits == 0):
                        m.next = "STOP"
                    with m.Else():
                        m.d.sync += rx_bits.eq(rx_bits - 1)
                with m.Else():
                    m.d.sync += rx_cnt.eq(rx_cnt - 1)
            with m.State("STOP"):
                with m.If(rx_cnt == 0):
                    with m.If(rx_sync):
                        m.d.sync += [rx_byte.eq(rx_shr), rx_valid.eq(1)]
                    m.next = "IDLE"
                with m.Else():
                    m.d.sync += rx_cnt.eq(rx_cnt - 1)

        # 1-deep RX holding register: decouples byte capture from the shell's
        # echo/print activity so a character arriving while the shell is busy
        # echoing (or printing a short response) isn't dropped.  This keeps a
        # back-to-back paste of a command line intact — the incoming and echo
        # rates match, so at most one byte is ever in flight during an echo.
        # (A byte arriving during a long print, e.g. 'help' output, can still
        # be lost, but you don't type into a response.)
        rx_pending = Signal()
        rx_hold    = Signal(8)
        with m.If(rx_valid):
            m.d.sync += [rx_hold.eq(rx_byte), rx_pending.eq(1)]

        # ── Line buffer (block RAM) ───────────────────────────────────────
        # Stored in an SB_RAM40 rather than flip-flops; the parser reads it
        # SEQUENTIALLY (one byte per cycle) through a single registered read
        # port.  This is what lets the 64-byte buffer AND its wide read muxes
        # leave the LC fabric (they were the dominant cost).  `la` drives the
        # read address; `lbuf_rd.data` is valid the cycle after `la` settles.
        m.submodules.lbuf_mem = lbuf_mem = Memory(
            shape=unsigned(8), depth=N, init=[])
        lbuf_wr = lbuf_mem.write_port()
        lbuf_rd = lbuf_mem.read_port()
        llen = Signal(range(N + 1))
        la   = Signal(range(N))
        m.d.comb += lbuf_rd.addr.eq(la)
        m.d.comb += la.eq(0)          # default; read states override per-cycle
        # write port: addr/data always presented; a store pulses `en` in INPUT.
        m.d.comb += [lbuf_wr.addr.eq(llen), lbuf_wr.data.eq(rx_hold),
                     lbuf_wr.en.eq(0)]

        def low(c):
            return Mux((c >= ord('A')) & (c <= ord('Z')), c | 0x20, c)

        # Command templates, matched sequentially in the SCAN state.  Padded to
        # a common length so a runtime index is always in range (guarded by the
        # per-template length so padding bytes are never actually compared).
        _TH, _TU, _TB, _TA = "help", "udp unicast ", "udp broadcast", "udp bcast"
        _TMAX = 14
        def _tmpl(s):
            return Array([Const(ord(c), 8) for c in s.ljust(_TMAX, "\x00")])
        THarr, TUarr, TBarr, TAarr = _tmpl(_TH), _tmpl(_TU), _tmpl(_TB), _tmpl(_TA)

        # ── Shared print routine ──────────────────────────────────────────
        pr_ptr = Signal(range(len(_ROM) + 1))
        pr_end = Signal(range(len(_ROM) + 1))
        pr_ret = Signal(3)
        to_ctr = Signal(range(self._to_cyc + 1))    # reply timeout countdown
        gap    = Signal(range(self._gap_cyc + 1))    # inter-poll gap countdown

        # ROM read port is addressed by the print pointer (1-cycle latency).
        m.d.comb += rom_rd.addr.eq(pr_ptr)

        def start_print(msg_key, ret, nxt="PRINT_PRIME"):
            s, e = _OFF[msg_key]
            m.d.sync += [pr_ptr.eq(s), pr_end.eq(e), pr_ret.eq(ret)]
            m.next = nxt

        # ── Parse / send scratch ──────────────────────────────────────────
        ip_b   = Array([Signal(8, name=f"ip{i}") for i in range(4)])
        ip_pos = Signal(range(N + 1))
        octet  = Signal(9)
        noct   = Signal(3)
        kwend  = Signal(range(N + 1))
        use_def = Signal()
        req_r  = Signal()
        ech    = Signal(8)   # character currently being echoed
        # sequential command scan
        si     = Signal(range(_TMAX + 2))          # scan index
        mh, muni, mbc, mbca = (Signal(init=1), Signal(init=1),
                               Signal(init=1), Signal(init=1))
        dh, dbc, dbca = Signal(), Signal(), Signal()   # delimiter-ok captures
        # payload streaming (prefetch: `la` runs one byte ahead of pl_byte)
        pl_base = Signal(range(N + 1))
        pl_len  = Signal(range(N + 1))
        j       = Signal(range(N + 1))
        pl_byte = Signal(8)
        m.d.comb += [self.udp_send_req.eq(req_r),
                     self.udp_dst_ip.eq(Cat(ip_b[3], ip_b[2], ip_b[1], ip_b[0])),
                     self.udp_pl_data.eq(pl_byte)]

        tx_load_i = Signal()
        tx_byte_i = Signal(8)
        m.d.comb += [tx_load.eq(tx_load_i), tx_byte.eq(tx_byte_i)]

        # ── Shell FSM ─────────────────────────────────────────────────────
        with m.FSM(name="shell"):
            with m.State("BOOT"):
                start_print("BANNER", _RET_PROMPT)

            # One prime cycle so the block RAM read (rom_rd.data = rom[pr_ptr])
            # is valid before the first byte is loaded.
            with m.State("PRINT_PRIME"):
                m.next = "PRINT"

            # Generic ROM printer → dispatch on pr_ret.  rom_rd.data holds
            # rom[pr_ptr] (address is comb-driven; stable across the tx-busy
            # wait, so it is valid whenever ~tx_busy lets a byte load).
            with m.State("PRINT"):
                with m.If(~tx_busy):
                    with m.If(pr_ptr == pr_end):
                        with m.Switch(pr_ret):
                            with m.Case(_RET_PROMPT):
                                m.next = "PROMPT"
                            with m.Case(_RET_INPUT):
                                m.next = "INPUT"
                            with m.Case(_RET_PARSE):
                                m.next = "PARSE"
                            with m.Case(_RET_RXWAIT):
                                m.next = "RX_INIT"
                            with m.Case(_RET_RXBODY):
                                m.next = "RX_BODY"
                    with m.Else():
                        m.d.comb += [tx_load_i.eq(1), tx_byte_i.eq(rom_rd.data)]
                        m.d.sync += pr_ptr.eq(pr_ptr + 1)

            with m.State("PROMPT"):
                m.d.sync += llen.eq(0)
                start_print("PROMPT", _RET_INPUT)

            # Read + echo a line into lbuf until CR/LF.
            with m.State("INPUT"):
                with m.If(rx_pending):
                    m.d.sync += rx_pending.eq(0)
                    with m.If((rx_hold == 0x0D) | (rx_hold == 0x0A)):
                        start_print("CRLF", _RET_PARSE)
                    with m.Elif((rx_hold == 0x08) | (rx_hold == 0x7F)):
                        with m.If(llen != 0):
                            m.d.sync += llen.eq(llen - 1)
                            start_print("BKSP", _RET_INPUT)
                    with m.Elif((rx_hold >= 0x20) & (rx_hold < 0x7F)):
                        with m.If(llen != N):
                            m.d.comb += lbuf_wr.en.eq(1)   # store at addr=llen
                            m.d.sync += [llen.eq(llen + 1), ech.eq(rx_hold)]
                            m.next = "ECHO"

            # Echo one stored character.
            with m.State("ECHO"):
                with m.If(~tx_busy):
                    m.d.comb += [tx_load_i.eq(1), tx_byte_i.eq(ech)]
                    m.next = "INPUT"

            # Decide which command the line holds — set up the sequential scan.
            with m.State("PARSE"):
                with m.If(llen == 0):
                    m.next = "PROMPT"
                with m.Else():
                    m.d.sync += [si.eq(0), mh.eq(1), muni.eq(1),
                                 mbc.eq(1), mbca.eq(1), dh.eq(0),
                                 dbc.eq(0), dbca.eq(0)]
                    m.next = "SCAN_RD"

            # Scan the first up-to-14 bytes, comparing each against all four
            # command templates in parallel (one byte/cycle from block RAM).
            with m.State("SCAN_RD"):
                m.d.comb += la.eq(si)          # address the current byte
                m.next = "SCAN_USE"
            with m.State("SCAN_USE"):
                m.d.comb += la.eq(si)
                c  = lbuf_rd.data
                lc = low(c)
                inb = si < llen                # byte position is within the line
                with m.If(si < len(_TH)):
                    m.d.sync += mh.eq(mh & inb & (lc == THarr[si]))
                with m.If(si == len(_TH)):
                    m.d.sync += dh.eq((si >= llen) | (c == ord(' ')))
                with m.If(si < len(_TU)):
                    m.d.sync += muni.eq(muni & inb & (lc == TUarr[si]))
                with m.If(si < len(_TB)):
                    m.d.sync += mbc.eq(mbc & inb & (lc == TBarr[si]))
                with m.If(si == len(_TB)):
                    m.d.sync += dbc.eq((si >= llen) | (c == ord(' ')))
                with m.If(si < len(_TA)):
                    m.d.sync += mbca.eq(mbca & inb & (lc == TAarr[si]))
                with m.If(si == len(_TA)):
                    m.d.sync += dbca.eq((si >= llen) | (c == ord(' ')))
                with m.If(si == len(_TB)):     # scanned enough to decide
                    m.next = "EVAL"
                with m.Else():
                    m.d.sync += si.eq(si + 1)
                    m.next = "SCAN_RD"

            # Classify from the accumulated match/delimiter bits (same priority
            # and semantics as the old parallel parser).
            with m.State("EVAL"):
                with m.If(mh & dh):
                    start_print("HELP", _RET_PROMPT)
                with m.Elif(muni):             # "udp unicast " prefix present
                    m.d.sync += [ip_pos.eq(len(_TU)), octet.eq(0), noct.eq(0),
                                 ip_b[0].eq(0), ip_b[1].eq(0),
                                 ip_b[2].eq(0), ip_b[3].eq(0)]
                    m.next = "IP_RD"
                with m.Elif(mbc & dbc):
                    m.d.sync += [kwend.eq(len(_TB)),
                                 ip_b[0].eq(0xFF), ip_b[1].eq(0xFF),
                                 ip_b[2].eq(0xFF), ip_b[3].eq(0xFF)]
                    m.next = "BC_RD"
                with m.Elif(mbca & dbca):
                    m.d.sync += [kwend.eq(len(_TA)),
                                 ip_b[0].eq(0xFF), ip_b[1].eq(0xFF),
                                 ip_b[2].eq(0xFF), ip_b[3].eq(0xFF)]
                    m.next = "BC_RD"
                with m.Else():
                    start_print("ERR", _RET_PROMPT)

            # Sequential dotted-quad IPv4 parse from lbuf[ip_pos ...].
            with m.State("IP_RD"):
                m.d.comb += la.eq(ip_pos)
                m.next = "IP_USE"
            with m.State("IP_USE"):
                m.d.comb += la.eq(ip_pos)
                c = lbuf_rd.data
                with m.If((ip_pos == llen) | (c == ord(' '))):
                    # end of address token → commit final octet + payload bounds
                    with m.If((noct == 3) & (octet <= 255)):
                        m.d.sync += ip_b[3].eq(octet)
                        with m.If(ip_pos == llen):
                            m.d.sync += use_def.eq(1)       # no payload
                        with m.Else():                      # c == ' '
                            m.d.sync += [pl_base.eq(ip_pos + 1),
                                         pl_len.eq(llen - (ip_pos + 1)),
                                         use_def.eq((ip_pos + 1) >= llen)]
                        m.next = "SEND_SETUP"
                    with m.Else():
                        start_print("BADIP", _RET_PROMPT)
                with m.Elif(c == ord('.')):
                    with m.If((noct < 3) & (octet <= 255)):
                        m.d.sync += [ip_b[noct].eq(octet), noct.eq(noct + 1),
                                     octet.eq(0), ip_pos.eq(ip_pos + 1)]
                        m.next = "IP_RD"
                    with m.Else():
                        start_print("BADIP", _RET_PROMPT)
                with m.Elif((c >= ord('0')) & (c <= ord('9'))):
                    m.d.sync += [octet.eq(octet * 10 + (c - ord('0'))),
                                 ip_pos.eq(ip_pos + 1)]
                    m.next = "IP_RD"
                with m.Else():
                    start_print("BADIP", _RET_PROMPT)

            # Broadcast payload: everything after the keyword's trailing space.
            with m.State("BC_RD"):
                m.d.comb += la.eq(kwend)
                m.next = "BC_USE"
            with m.State("BC_USE"):
                m.d.comb += la.eq(kwend)
                with m.If((llen > kwend) & (lbuf_rd.data == ord(' '))):
                    m.d.sync += [pl_base.eq(kwend + 1),
                                 pl_len.eq(llen - (kwend + 1)),
                                 use_def.eq((kwend + 1) >= llen)]
                with m.Else():
                    m.d.sync += use_def.eq(1)
                m.next = "SEND_SETUP"

            with m.State("SEND_SETUP"):
                with m.If(self.udp_test_busy):
                    start_print("BUSY", _RET_PROMPT)
                with m.Else():
                    with m.If(use_def):
                        m.d.sync += [pl_base.eq(0), pl_len.eq(DEF_LEN)]
                    m.d.sync += j.eq(0)
                    m.next = "PL_PRIME"

            # Prime the prefetch register with payload byte 0 (block RAM read
            # for typed payload; the default payload is combinational).
            with m.State("PL_PRIME"):
                m.d.comb += la.eq(pl_base)     # fetch typed byte 0
                m.next = "PL_PRIME2"
            with m.State("PL_PRIME2"):
                m.d.comb += la.eq(pl_base)
                m.d.sync += pl_byte.eq(Mux(use_def, defrom[0], lbuf_rd.data))
                m.next = "SEND_REQ"

            # Hold the request until the W5100 acknowledges by going busy.
            with m.State("SEND_REQ"):
                m.d.comb += la.eq(pl_base + j + 1)     # prefetch next byte
                m.d.sync += req_r.eq(1)
                with m.If(self.udp_test_busy):
                    m.d.sync += req_r.eq(0)
                    m.next = "SEND_STREAM"

            # Feed payload bytes as the W5100 consumes them.  `pl_byte` holds
            # the current byte; `la` prefetches the next so it is ready by the
            # time the master pulls it (block RAM 1-cycle latency hidden).
            with m.State("SEND_STREAM"):
                m.d.comb += la.eq(pl_base + j + 1)
                m.d.comb += [self.udp_pl_valid.eq(1),
                             self.udp_pl_last.eq(j + 1 == pl_len)]
                with m.If(self.udp_pl_ready):
                    with m.If(j + 1 == pl_len):
                        m.next = "SEND_WAIT"
                    with m.Else():
                        m.d.sync += [pl_byte.eq(Mux(use_def, defrom[j + 1],
                                                    lbuf_rd.data)),
                                     j.eq(j + 1)]

            with m.State("SEND_WAIT"):
                with m.If(~self.udp_test_busy):
                    # datagram is out; now wait (bounded) for a reply.
                    start_print("SENT", _RET_RXWAIT)

            # ── Wait for a UDP reply on the socket, with a timeout ────────────
            with m.State("RX_INIT"):
                m.d.sync += to_ctr.eq(self._to_cyc)
                m.next = "RX_POLL"

            # Ask the W5100 to check the socket RX buffer.
            with m.State("RX_POLL"):
                m.d.comb += self.udp_rx_req.eq(1)
                with m.If(self.udp_rx_busy):
                    m.next = "RX_WAIT"

            # One of: a datagram streams (udp_rx_valid) or none (udp_rx_none).
            with m.State("RX_WAIT"):
                with m.If(self.udp_rx_valid):
                    start_print("RXPFX", _RET_RXBODY)
                with m.Elif(self.udp_rx_none):
                    m.d.sync += gap.eq(self._gap_cyc)
                    m.next = "RX_GAP"

            # Idle a bit between polls; count down the overall timeout.
            with m.State("RX_GAP"):
                with m.If(to_ctr == 0):
                    start_print("TIMEOUT", _RET_PROMPT)
                with m.Elif(gap == 0):
                    m.next = "RX_POLL"
                with m.Else():
                    m.d.sync += [gap.eq(gap - 1), to_ctr.eq(to_ctr - 1)]

            # Stream the reply payload to the UART (throttled by tx_busy).
            with m.State("RX_BODY"):
                with m.If(self.udp_rx_valid & ~tx_busy):
                    m.d.comb += [tx_load_i.eq(1), tx_byte_i.eq(self.udp_rx_data),
                                 self.udp_rx_ready.eq(1)]
                    with m.If(self.udp_rx_eof):
                        m.next = "RX_BODY_END"

            with m.State("RX_BODY_END"):
                start_print("CRLF", _RET_PROMPT)

        return m


# ── Testbench ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import sys
    from amaranth.sim import Simulator, Period

    # Bit period is irrelevant to the shell logic, so use a tiny divisor to
    # keep the simulation fast; the real build uses 24 MHz / 115200 (div≈208).
    # Small reply timeout/gap so the timeout test finishes quickly.
    CLK, BAUD = 24_000_000, 3_000_000
    DIV = round(CLK / BAUD)               # = 8 sim cycles per bit
    dut = UARTShell(clk_freq=CLK, baud_rate=BAUD,
                    reply_timeout_cycles=4000, poll_gap_cycles=300)
    errors = []

    # ── UART line helpers (drive uart_rx, sample uart_tx) ──────────────────
    async def send_byte(ctx, val):
        ctx.set(dut.uart_rx, 0)
        await ctx.tick().repeat(DIV)
        for i in range(8):
            ctx.set(dut.uart_rx, (val >> i) & 1)
            await ctx.tick().repeat(DIV)
        ctx.set(dut.uart_rx, 1)
        await ctx.tick().repeat(DIV)

    async def send_line(ctx, s):
        for ch in s:
            await send_byte(ctx, ord(ch))
        await send_byte(ctx, 0x0D)          # Enter

    # Background collector: continuously samples uart_tx into a byte list.
    tx_chars = bytearray()

    async def tx_collector(ctx):
        while True:
            # wait for start bit
            if ctx.get(dut.uart_tx) == 0:
                await ctx.tick().repeat(DIV // 2)
                if ctx.get(dut.uart_tx) != 0:
                    continue
                b = 0
                for i in range(8):
                    await ctx.tick().repeat(DIV)
                    b |= ctx.get(dut.uart_tx) << i
                await ctx.tick().repeat(DIV)   # stop
                tx_chars.append(b)
            else:
                await ctx.tick()

    # Background responder: models the W5100 UDP interface for BOTH directions.
    # On udp_send_req it drains + records the payload (send handshake); on
    # udp_rx_req it either streams a staged reply datagram or pulses udp_rx_none
    # (nothing waiting) — mirroring the real master's poll-driven RX.
    captured = {"payload": None, "dst": None, "count": 0}
    staged   = {"reply": None}      # bytes to deliver on the next RX poll

    async def w5100_udp_model(ctx):
        while True:
            await ctx.tick()
            if ctx.get(dut.udp_send_req):
                ctx.set(dut.udp_test_busy, 1)
                dst = ctx.get(dut.udp_dst_ip)
                await ctx.tick().repeat(6)
                pl = bytearray(); last = False; guard = 0
                while not last:
                    ctx.set(dut.udp_pl_ready, 1)
                    got = ctx.get(dut.udp_pl_valid)
                    if got:
                        pl.append(ctx.get(dut.udp_pl_data))
                        last = bool(ctx.get(dut.udp_pl_last))
                    await ctx.tick()
                    if got:
                        # Space consumes ~5 cycles apart, like the W5100's bus
                        # cycle — the shell's block-RAM prefetch needs ≥1 cycle
                        # between pulls to present the next byte.
                        ctx.set(dut.udp_pl_ready, 0)
                        await ctx.tick().repeat(4)
                    guard += 1
                    if guard > 5000: break
                ctx.set(dut.udp_pl_ready, 0)
                await ctx.tick().repeat(6)
                captured["payload"] = bytes(pl)
                captured["dst"] = [(dst >> 24) & 0xFF, (dst >> 16) & 0xFF,
                                   (dst >> 8) & 0xFF, dst & 0xFF]
                captured["count"] += 1
                ctx.set(dut.udp_test_busy, 0)
            elif ctx.get(dut.udp_rx_req):
                ctx.set(dut.udp_rx_busy, 1)
                await ctx.tick().repeat(4)      # emulate RSR/RD/header latency
                rep = staged["reply"]
                if rep is None:
                    ctx.set(dut.udp_rx_none, 1)
                    await ctx.tick()
                    ctx.set(dut.udp_rx_none, 0)
                else:
                    staged["reply"] = None
                    for i, b in enumerate(rep):
                        ctx.set(dut.udp_rx_data, b)
                        ctx.set(dut.udp_rx_valid, 1)
                        ctx.set(dut.udp_rx_eof, 1 if i == len(rep) - 1 else 0)
                        g = 0
                        while not ctx.get(dut.udp_rx_ready):
                            await ctx.tick(); g += 1
                            if g > 20000: break
                        await ctx.tick()        # consume cycle
                    ctx.set(dut.udp_rx_valid, 0)
                    ctx.set(dut.udp_rx_eof, 0)
                ctx.set(dut.udp_rx_busy, 0)

    async def wait_prompt(ctx, timeout=400_000):
        """Wait until the tail of tx_chars ends with 'rebbarb> '."""
        for _ in range(timeout):
            if tx_chars.endswith(b"rebbarb> "):
                return True
            await ctx.tick()
        return False

    async def testbench(ctx):
        ctx.set(dut.uart_rx, 1)
        # udp_pl_ready / udp_test_busy are owned by w5100_udp_model — do not
        # drive them here (two testbenches on one signal deadlocks the send).

        # Boot banner + first prompt.
        if not await wait_prompt(ctx):
            errors.append("no initial prompt"); return
        print(f"boot tx: {bytes(tx_chars)!r}")

        # T1: 'help' lists commands.
        tx_chars.clear()
        await send_line(ctx, "help")
        await wait_prompt(ctx)
        if b"udp unicast" not in tx_chars or b"udp broadcast" not in tx_chars:
            errors.append(f"T1 help missing commands: {bytes(tx_chars)!r}")
        print(f"T1 help ok ({len(tx_chars)} bytes)")

        # T2: unicast with payload, and a staged reply → 'rx <reply>' printed.
        tx_chars.clear(); captured["count"] = 0
        staged["reply"] = b"pong-A"
        await send_line(ctx, "udp unicast 192.168.1.55 hello world")
        await wait_prompt(ctx)
        if captured["dst"] != [192, 168, 1, 55]:
            errors.append(f"T2 dst {captured['dst']} != [192,168,1,55]")
        if captured["payload"] != b"hello world":
            errors.append(f"T2 payload {captured['payload']!r} != b'hello world'")
        if b"sent" not in tx_chars:
            errors.append(f"T2 no 'sent' ack: {bytes(tx_chars)!r}")
        if b"rx pong-A" not in tx_chars:
            errors.append(f"T2 reply not printed: {bytes(tx_chars)!r}")
        print(f"T2 unicast+reply: dst={captured['dst']} "
              f"payload={captured['payload']!r} tx={bytes(tx_chars)!r}")

        # T3: broadcast default payload, NO reply → 'timeout' printed (no hang).
        tx_chars.clear(); staged["reply"] = None
        await send_line(ctx, "udp broadcast")
        await wait_prompt(ctx)
        if captured["dst"] != [255, 255, 255, 255]:
            errors.append(f"T3 dst {captured['dst']} != broadcast")
        if captured["payload"] != bytes(_DEFAULT_PAYLOAD):
            errors.append(f"T3 payload {captured['payload']!r} != default")
        if b"timeout" not in tx_chars:
            errors.append(f"T3 no timeout on no-reply: {bytes(tx_chars)!r}")
        print(f"T3 broadcast default→timeout: dst={captured['dst']} "
              f"payload={captured['payload']!r}")

        # T4: broadcast with payload + reply.
        tx_chars.clear(); staged["reply"] = b"Hello World"
        await send_line(ctx, "udp broadcast ping123")
        await wait_prompt(ctx)
        if captured["payload"] != b"ping123":
            errors.append(f"T4 payload {captured['payload']!r} != b'ping123'")
        if b"rx Hello World" not in tx_chars:
            errors.append(f"T4 reply not printed: {bytes(tx_chars)!r}")
        print(f"T4 broadcast+reply: payload={captured['payload']!r} "
              f"tx={bytes(tx_chars)!r}")

        # T5: case-insensitive command + uppercase preserved in payload.
        tx_chars.clear(); staged["reply"] = b"ok"
        await send_line(ctx, "UDP UNICAST 10.0.0.9 MixedCase")
        await wait_prompt(ctx)
        if captured["dst"] != [10, 0, 0, 9]:
            errors.append(f"T5 dst {captured['dst']} != [10,0,0,9]")
        if captured["payload"] != b"MixedCase":
            errors.append(f"T5 payload {captured['payload']!r} != b'MixedCase'")
        print(f"T5 case: dst={captured['dst']} payload={captured['payload']!r}")

        # T6: bad IP → 'bad ip', no send.
        tx_chars.clear(); before = captured["count"]
        await send_line(ctx, "udp unicast 1.2.3 x")
        await wait_prompt(ctx)
        if b"bad ip" not in tx_chars:
            errors.append(f"T6 no 'bad ip': {bytes(tx_chars)!r}")
        if captured["count"] != before:
            errors.append("T6 sent despite bad ip")
        print(f"T6 bad ip ok (no send)")

        # T7: unknown command → error.
        tx_chars.clear()
        await send_line(ctx, "frobnicate")
        await wait_prompt(ctx)
        if b"?" not in tx_chars:
            errors.append(f"T7 no error marker: {bytes(tx_chars)!r}")
        print(f"T7 unknown ok")

        # T8: backspace editing (type 'helX' <bs> 'p' → 'help').  The backspace
        # echo is 3 bytes ("\b \b"), longer than one incoming byte time, so a
        # realistic source pauses after it (a human always does).  Settle
        # between edits — this is the one spot the 1-deep RX register can't
        # absorb, and it never occurs with paste (no backspaces) or typing.
        tx_chars.clear()
        for ch in "helX":
            await send_byte(ctx, ord(ch))
            await ctx.tick().repeat(DIV * 15)   # > 1-byte echo (10 bits)
        await send_byte(ctx, 0x08)              # backspace removes 'X'
        await ctx.tick().repeat(DIV * 40)       # > 3-byte "\b \b" echo (~30 bits)
        await send_byte(ctx, ord('p'))
        await ctx.tick().repeat(DIV * 15)
        await send_byte(ctx, 0x0D)
        await wait_prompt(ctx)
        if b"udp unicast" not in tx_chars:
            errors.append(f"T8 backspace edit failed: {bytes(tx_chars)!r}")
        print(f"T8 backspace edit ok")

    sim = Simulator(dut)
    sim.add_clock(Period(MHz=24))
    sim.add_testbench(testbench)
    sim.add_testbench(tx_collector, background=True)
    sim.add_testbench(w5100_udp_model, background=True)
    sim.run()

    if errors:
        print("\nFAILURES:")
        for e in errors:
            print(" ", e)
        sys.exit(1)
    print("\nAll UARTShell tests passed.")