"""RX frame assembler — sync domain (24 MHz).

Receives raw ethernet frames from W5500SPIMaster and writes them into the SPRAM
ring buffer in MX98730EC format.

Ring buffer layout (SPRAM byte addresses)
------------------------------------------
0x0100–0x0FFF  15 pages × 256 bytes = 3840 bytes
Pages 0x01–0x0F; page 0x00 is reserved.
Page wrap: after 0x0F → 0x01 (skip 0x00).

Frame descriptor (4 bytes at page start)
-----------------------------------------------
Byte 0: LRPS (last received packet status) — 0x00
Byte 1: 0x00
Byte 2: total_length[15:8]   (big-endian; includes 4 descriptor bytes)
Byte 3: total_length[7:0]
Bytes 4+: raw ethernet frame

Write sequence
--------------
1. Issue 4 SPRAM writes of 0x00 (placeholder descriptor).
2. For each byte received from W5500, issue one SPRAM write.
3. After EOF: rewrite descriptor bytes 2 and 3 with actual length.
4. Advance RWP, push to rx_wptr FIFO, pulse rx_irq.
"""

from amaranth import *

__all__ = ["RXFrameAssembler"]

_RX_PAGE_FIRST = 0x01
_RX_PAGE_LAST  = 0x0F
_PAGES_TOTAL   = _RX_PAGE_LAST - _RX_PAGE_FIRST + 1   # 15


class RXFrameAssembler(Elaboratable):
    """Writes incoming ethernet frames into the SPRAM ring buffer.

    W5500 streaming interface (sync domain)
    ----------------------------------------
    rx_data / rx_valid / rx_ready : byte stream
    rx_sof / rx_eof               : frame delimiters (same cycle as rx_valid)

    SPRAM write interface (to SPRAMArbiter, sync domain)
    -----------------------------------------------------
    eth_wr_addr / eth_wr_data / eth_wr_valid / eth_wr_ready

    CDC outputs (wired by BBATop)
    -----------------------------
    rx_wptr_w_data / rx_wptr_w_en / rx_wptr_w_rdy
    rx_irq : 1-cycle pulse → PulseSynchronizer input
    rx_enabled : controlled by NCRA SR bit (from BBARegisterFile)
    """

    def __init__(self):
        # W5500 stream in
        self.rx_data   = Signal(8)
        self.rx_valid  = Signal()
        self.rx_ready  = Signal()
        self.rx_sof    = Signal()
        self.rx_eof    = Signal()

        # SPRAM write out
        self.eth_wr_addr  = Signal(16)
        self.eth_wr_data  = Signal(8)
        self.eth_wr_valid = Signal()
        self.eth_wr_ready = Signal()

        # RWP FIFO write-side (sync→exi)
        self.rx_wptr_w_data = Signal(8)
        self.rx_wptr_w_en   = Signal()
        self.rx_wptr_w_rdy  = Signal()

        # rx_irq pulse (→ PulseSynchronizer)
        self.rx_irq = Signal()

        # RX gate from NCRA SR bit
        self.rx_enabled = Signal()

    def elaborate(self, platform):
        m = Module()

        # ── Ring-buffer state ─────────────────────────────────────────────
        rwp = Signal(8, init=_RX_PAGE_FIRST)   # current RX write page (1–15)

        # Write address within current frame
        wr_addr  = Signal(16)
        # Number of frame data bytes received
        data_ctr = Signal(12)
        # Total length = data_ctr + 4
        total_len = Signal(12)

        # Descriptor base (rwp*256) — saved when frame starts
        desc_base = Signal(16)

        # Placeholder descriptor byte counter (0..3)
        desc_ctr = Signal(2)

        # Number of pages consumed by this frame (rounded up)
        pages_used = Signal(5)

        # Default: no pulses
        m.d.sync += self.rx_irq.eq(0)
        m.d.sync += self.rx_wptr_w_en.eq(0)

        # Combinatorial outputs
        m.d.comb += total_len.eq(data_ctr + 4)

        with m.FSM(domain="sync", name="rx_fsm"):

            with m.State("IDLE"):
                m.d.comb += self.rx_ready.eq(0)
                m.d.sync += self.eth_wr_valid.eq(0)
                with m.If(self.rx_valid & self.rx_sof & self.rx_enabled):
                    frame_base = Signal(16)
                    m.d.comb += frame_base.eq(Cat(Const(0, 8), rwp))
                    m.d.sync += desc_base.eq(frame_base)
                    m.d.sync += wr_addr.eq(frame_base)
                    m.d.sync += data_ctr.eq(0)
                    m.d.sync += desc_ctr.eq(0)
                    m.next = "WRITE_PLACEHOLDER"

            with m.State("WRITE_PLACEHOLDER"):
                # Write 4 bytes of 0x00 as placeholder descriptor
                m.d.sync += self.eth_wr_addr.eq(wr_addr)
                m.d.sync += self.eth_wr_data.eq(0x00)
                m.d.sync += self.eth_wr_valid.eq(1)
                with m.If(self.eth_wr_ready):
                    m.d.sync += wr_addr.eq(wr_addr + 1)
                    with m.If(desc_ctr == 3):
                        m.d.sync += self.eth_wr_valid.eq(0)
                        m.next = "RECV_AND_WRITE"
                    with m.Else():
                        m.d.sync += desc_ctr.eq(desc_ctr + 1)

            with m.State("RECV_AND_WRITE"):
                # Accept bytes from W5500 and write each to SPRAM immediately
                m.d.comb += self.rx_ready.eq(~self.eth_wr_valid | self.eth_wr_ready)
                with m.If(self.rx_valid & (~self.eth_wr_valid | self.eth_wr_ready)):
                    m.d.sync += self.eth_wr_addr.eq(wr_addr)
                    m.d.sync += self.eth_wr_data.eq(self.rx_data)
                    m.d.sync += self.eth_wr_valid.eq(1)
                    m.d.sync += wr_addr.eq(wr_addr + 1)
                    m.d.sync += data_ctr.eq(data_ctr + 1)
                    with m.If(self.rx_eof):
                        m.next = "WAIT_LAST_WRITE"
                with m.Elif(self.eth_wr_valid & self.eth_wr_ready):
                    m.d.sync += self.eth_wr_valid.eq(0)

            with m.State("WAIT_LAST_WRITE"):
                # Wait for the last data byte write to be accepted
                with m.If(~self.eth_wr_valid | self.eth_wr_ready):
                    m.d.sync += self.eth_wr_valid.eq(0)
                    # Compute pages used: ceil((data_ctr + 4) / 256)
                    # = (total_len + 255) >> 8 = total_len[11:8] + (total_len[7:0] != 0)
                    m.d.sync += pages_used.eq(total_len[8:12] + (total_len[:8] != 0))
                    m.next = "WRITE_LEN_HI"

            with m.State("WRITE_LEN_HI"):
                # Overwrite descriptor byte 2 with total_len[15:8]
                m.d.sync += self.eth_wr_addr.eq(desc_base + 2)
                m.d.sync += self.eth_wr_data.eq(total_len[8:12])
                m.d.sync += self.eth_wr_valid.eq(1)
                with m.If(self.eth_wr_ready):
                    m.d.sync += self.eth_wr_valid.eq(0)
                    m.next = "WRITE_LEN_LO"

            with m.State("WRITE_LEN_LO"):
                # Overwrite descriptor byte 3 with total_len[7:0]
                m.d.sync += self.eth_wr_addr.eq(desc_base + 3)
                m.d.sync += self.eth_wr_data.eq(total_len[:8])
                m.d.sync += self.eth_wr_valid.eq(1)
                with m.If(self.eth_wr_ready):
                    m.d.sync += self.eth_wr_valid.eq(0)
                    m.next = "ADVANCE_RWP"

            with m.State("ADVANCE_RWP"):
                # next_rwp = ((rwp - 1 + pages_used) % 15) + 1
                next_rwp_raw = Signal(8)
                m.d.comb += next_rwp_raw.eq(rwp + pages_used)
                with m.If(next_rwp_raw > _RX_PAGE_LAST):
                    m.d.sync += rwp.eq(next_rwp_raw - _PAGES_TOTAL)
                with m.Else():
                    m.d.sync += rwp.eq(next_rwp_raw)
                m.next = "PUSH_WPT"

            with m.State("PUSH_WPT"):
                with m.If(self.rx_wptr_w_rdy):
                    m.d.sync += self.rx_wptr_w_data.eq(rwp)
                    m.d.sync += self.rx_wptr_w_en.eq(1)
                    m.d.sync += self.rx_irq.eq(1)
                    m.next = "IDLE"

        return m


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

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

    dut = RXFrameAssembler()
    errors = []

    # Track all SPRAM writes issued by the DUT
    spram_writes = []

    async def testbench(ctx):
        # Setup: acknowledge all SPRAM writes immediately
        ctx.set(dut.eth_wr_ready,    1)
        ctx.set(dut.rx_wptr_w_rdy,  1)
        ctx.set(dut.rx_enabled,     1)
        await ctx.tick("sync").repeat(2)

        # ── T1: 10-byte frame → pages_used=1, rwp advances 1→2 ──────────────
        # Send SOF + first byte
        frame = [0xAA, 0xBB, 0xCC, 0xDD, 0x08, 0x00, 0x45, 0x00, 0x00, 0x01]

        ctx.set(dut.rx_data,  frame[0])
        ctx.set(dut.rx_valid, 1)
        ctx.set(dut.rx_sof,   1)
        await ctx.tick("sync").repeat(1)
        ctx.set(dut.rx_sof, 0)

        for i, b in enumerate(frame[1:], start=1):
            ctx.set(dut.rx_data, b)
            ctx.set(dut.rx_eof,  1 if i == len(frame) - 1 else 0)
            await ctx.tick("sync").repeat(1)

        ctx.set(dut.rx_valid, 0)
        ctx.set(dut.rx_eof,   0)

        # Poll for up to 30 ticks until rx_irq pulses (1-cycle pulse)
        t1_irq_seen = False
        t1_wptr_d   = 0
        for _ in range(30):
            await ctx.tick("sync").repeat(1)
            if ctx.get(dut.rx_irq):
                t1_irq_seen = True
                t1_wptr_d   = ctx.get(dut.rx_wptr_w_data)
                break

        print(f"T1 rx_irq_seen={t1_irq_seen}  wptr_data=0x{t1_wptr_d:02X}")
        if not t1_irq_seen:
            errors.append("T1: rx_irq never pulsed")
        if t1_wptr_d != 2:
            errors.append(f"T1: rwp should be 2 (page 1→2), got {t1_wptr_d}")

        await ctx.tick("sync").repeat(4)

        # ── T2: Send a second frame; verify rwp advances further ────────────
        frame2 = [0x11, 0x22, 0x33, 0x44, 0x55, 0x66]
        ctx.set(dut.rx_data,  frame2[0])
        ctx.set(dut.rx_valid, 1)
        ctx.set(dut.rx_sof,   1)
        await ctx.tick("sync").repeat(1)
        ctx.set(dut.rx_sof, 0)

        for i, b in enumerate(frame2[1:], start=1):
            ctx.set(dut.rx_data, b)
            ctx.set(dut.rx_eof,  1 if i == len(frame2) - 1 else 0)
            await ctx.tick("sync").repeat(1)

        ctx.set(dut.rx_valid, 0)
        ctx.set(dut.rx_eof,   0)

        t2_irq_seen = False
        t2_wptr_d   = 0
        for _ in range(30):
            await ctx.tick("sync").repeat(1)
            if ctx.get(dut.rx_irq):
                t2_irq_seen = True
                t2_wptr_d   = ctx.get(dut.rx_wptr_w_data)
                break

        print(f"T2 rx_irq_seen={t2_irq_seen}  wptr_data=0x{t2_wptr_d:02X}")
        if not t2_irq_seen:
            errors.append("T2: rx_irq never pulsed after second frame")
        if t2_wptr_d != 3:
            errors.append(f"T2: rwp should be 3 (page 2→3), got {t2_wptr_d}")

        # ── T3: RX disabled — SOF must be ignored ──────────────────────────
        ctx.set(dut.rx_enabled, 0)
        ctx.set(dut.rx_data,  0xDE)
        ctx.set(dut.rx_valid, 1)
        ctx.set(dut.rx_sof,   1)
        await ctx.tick("sync").repeat(4)
        ctx.set(dut.rx_valid, 0)
        ctx.set(dut.rx_sof, 0)

        # No SPRAM write should have been issued
        wr_valid = ctx.get(dut.eth_wr_valid)
        if wr_valid:
            errors.append("T3: SPRAM write issued while rx_enabled=0")
        print(f"T3 rx disabled: eth_wr_valid={wr_valid} (expected 0)")

    sim = Simulator(dut)
    sim.add_clock(Period(MHz=24), domain="sync")
    sim.add_testbench(testbench)

    with sim.write_vcd("RXFrameAssembler.vcd"):
        sim.run()

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