from __future__ import annotations

import cocotb
from cocotb.clock import Clock
from cocotb.triggers import RisingEdge


# ── Helpers ──────────────────────────────────────────────────────────────────


async def reset_fifo(dut):
    """Assert reset for 5 cycles, release, then wait 2 cycles."""
    dut.rst_n.value = 0
    dut.wr_valid_i.value = 0
    dut.rd_ready_i.value = 0
    dut.wr_row_i.value = 0
    dut.wr_score_i.value = 0
    for _ in range(5):
        await RisingEdge(dut.clk)
    dut.rst_n.value = 1
    for _ in range(2):
        await RisingEdge(dut.clk)


async def push(dut, row: int, score: int):
    """Push a (row, score) entry, waiting until wr_ready_o is asserted."""
    while not dut.wr_ready_o.value:
        await RisingEdge(dut.clk)
    dut.wr_row_i.value = row
    dut.wr_score_i.value = score
    dut.wr_valid_i.value = 1
    await RisingEdge(dut.clk)
    dut.wr_valid_i.value = 0


async def pop(dut) -> tuple[int, int]:
    """Pop an entry, waiting until rd_valid_o is asserted. Return (row, score)."""
    while not dut.rd_valid_o.value:
        await RisingEdge(dut.clk)
    # Sample data in the cycle rd_valid_o is high, before accepting
    row = int(dut.rd_row_o.value)
    score = int(dut.rd_score_o.value)
    dut.rd_ready_i.value = 1
    await RisingEdge(dut.clk)
    dut.rd_ready_i.value = 0
    return (row, score)


async def push_pop_simultaneous(dut, row: int, score: int) -> tuple[int, int]:
    """Push and pop in the same cycle. Return the popped (row, score)."""
    # Capture read data before accepting (combinational read)
    row_out = int(dut.rd_row_o.value)
    score_out = int(dut.rd_score_o.value)
    # Assert both handshakes
    dut.wr_row_i.value = row
    dut.wr_score_i.value = score
    dut.wr_valid_i.value = 1
    dut.rd_ready_i.value = 1
    await RisingEdge(dut.clk)
    dut.wr_valid_i.value = 0
    dut.rd_ready_i.value = 0
    return (row_out, score_out)


# ── Test 1: Reset  ───────────────────────────────────────────────────────────


@cocotb.test()
async def reset_asserts_empty_and_defaults(dut):
    """
    After reset:
      empty_o    == 1
      full_o     == 0
      rd_valid_o == 0
      wr_ready_o == 1
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_fifo(dut)

    assert dut.empty_o.value == 1, "empty_o should be 1 after reset"
    assert dut.full_o.value == 0, "full_o should be 0 after reset"
    assert dut.rd_valid_o.value == 0, "rd_valid_o should be 0 after reset"
    assert dut.wr_ready_o.value == 1, "wr_ready_o should be 1 after reset"


# ── Test 2: Order preservation  ──────────────────────────────────────────────


@cocotb.test()
async def fifo_preserves_order(dut):
    """
    Push [(3,44),(1,55),(7,22),(2,55)] then pop and verify each entry appears
    in the same order.  After draining, empty_o == 1.
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_fifo(dut)

    entries = [(3, 44), (1, 55), (7, 22), (2, 55)]

    for row, score in entries:
        await push(dut, row, score)

    for expected_row, expected_score in entries:
        row, score = await pop(dut)
        assert row == expected_row, (
            f"row mismatch: got {row}, expected {expected_row}"
        )
        assert score == expected_score, (
            f"score mismatch: got {score}, expected {expected_score}"
        )

    assert dut.empty_o.value == 1, "fifo should be empty after draining all entries"


# ── Test 3: Full / back-pressure / drain  ────────────────────────────────────


@cocotb.test()
async def fifo_full_backpressure_and_drain(dut):
    """
    Fill the FIFO until full (discover depth), confirm backpressure,
    pop one, push a new item, drain all, verify final inserted item
    (99, 999) is preserved.
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_fifo(dut)

    # ── Fill the FIFO until full ──
    depth = 0
    while not dut.full_o.value:
        await push(dut, depth, depth * 10)
        depth += 1

    assert depth > 0, "should have pushed at least one item"
    assert dut.full_o.value == 1, "fifo should report full after DEPTH pushes"
    assert dut.wr_ready_o.value == 0, "wr_ready should be 0 when full"

    # ── Pop one ──
    row, score = await pop(dut)
    assert row == 0 and score == 0, (
        f"first popped entry should be (0,0), got ({row},{score})"
    )

    assert dut.full_o.value == 0, "fifo should no longer be full after one pop"
    assert dut.wr_ready_o.value == 1, "wr_ready should be 1 after one pop"

    # ── Push a new item ──
    await push(dut, 99, 999)

    # ── Drain all entries ──
    # Expected order: (1,10), (2,20), … (depth-1, (depth-1)*10), (99,999)
    for i in range(1, depth):
        row, score = await pop(dut)
        assert row == i, (
            f"row mismatch during drain: got {row}, expected {i}"
        )
        assert score == i * 10, (
            f"score mismatch during drain: got {score}, expected {i * 10}"
        )

    row, score = await pop(dut)
    assert row == 99 and score == 999, (
        f"final inserted item wrong: got ({row},{score}), expected (99,999)"
    )

    assert dut.empty_o.value == 1, "fifo should be empty after draining everything"


# ── Test 4: Simultaneous push+pop at full  ───────────────────────────────────


@cocotb.test()
async def test_simultaneous_push_pop_when_full(dut):
    """
    When the FIFO is full, a simultaneous read should allow a write to pass
    through in the same cycle (wr_ready_o goes high combinationally because
    of the read handshake).  Verify the pop returns the oldest item, the
    push inserts a new item, and the fill level stays at depth.
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_fifo(dut)

    # ── Fill to full ──
    depth = 0
    while not dut.full_o.value:
        await push(dut, depth, depth * 10)
        depth += 1

    assert dut.full_o.value == 1
    assert dut.wr_ready_o.value == 0  # backpressure active

    # ── Simultaneous push+pop ──
    # wr_ready_o should be 1 this cycle because rd_valid_o && rd_ready_i
    old_row, old_score = await push_pop_simultaneous(dut, 99, 999)
    assert old_row == 0 and old_score == 0, (
        f"simultaneous pop should return the oldest item, got ({old_row},{old_score})"
    )

    # Fill level unchanged (pop frees a slot, push immediately refills it)
    assert dut.full_o.value == 1, "should remain full after simultaneous push+pop"

    # ── Drain — expected: (1,10) … (depth-1,…), (99,999) ──
    for i in range(1, depth):
        row, score = await pop(dut)
        assert row == i, (
            f"row mismatch: got {row}, expected {i}"
        )

    row, score = await pop(dut)
    assert row == 99 and score == 999, (
        f"final item wrong: got ({row},{score})"
    )

    assert dut.empty_o.value == 1


# ── Test 5: Non-power-of-two depth  ──────────────────────────────────────────


@cocotb.test()
async def test_non_power_of_two_depth(dut):
    """
    Validate pointer wrap at a non-power-of-two DEPTH boundary (e.g. 10).
    Fill the FIFO, then perform multiple pop+push cycles to force the write
    pointer past a binary-wrap boundary, then drain and verify data integrity.
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_fifo(dut)

    # ── Fill to full ──
    depth = 0
    while not dut.full_o.value:
        await push(dut, depth, depth)
        depth += 1

    # ── Force several pointer wraps by pop+push ──
    # For a broken RTL that wraps at 2**clog2(DEPTH) instead of DEPTH,
    # enough iterations will hit out-of-bounds storage or corrupt data.
    wraps = depth * 3  # well past the binary wrap point
    for i in range(wraps):
        row, score = await pop(dut)
        await push(dut, 500 + i, 600 + i)

    # ── Drain and verify count ──
    items = []
    while dut.rd_valid_o.value:
        row, score = await pop(dut)
        items.append((row, score))

    assert len(items) == depth, (
        f"expected {depth} items after {wraps} push/pop cycles, got {len(items)}"
    )

    # All items should be from the latest wrap cycle (order preserved)
    # Final batch started at iteration wraps - depth:
    for idx, (row, _score) in enumerate(items):
        expected_row = 500 + (wraps - depth + idx)
        assert row == expected_row, (
            f"item {idx}: expected row {expected_row}, got {row}"
        )

    assert dut.empty_o.value == 1


# ── Test 6: Depth = 1  ───────────────────────────────────────────────────────


@cocotb.test()
async def test_depth_one(dut):
    """
    Validate DEPTH=1 FIFO: push one item → full, pop → empty,
    push another → full, pop → empty.

    Gated: pushes one item; if the FIFO does *not* become full, we are
    compiled with DEPTH > 1 and the test silently passes (skips).
    (Use ``make test-depth-1`` to run with FIFO_DEPTH=1.)
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_fifo(dut)

    # ── Probe: is DEPTH == 1? ────────────────────────────────────────────
    await push(dut, 42, 100)
    if not dut.full_o.value:
        # Not compiled at DEPTH=1; pop the item back and return.
        await pop(dut)
        return

    # ── Full DEPTH=1 logic follows ───────────────────────────────────────
    assert dut.wr_ready_o.value == 0, "wr_ready should be 0 when full"

    row, score = await pop(dut)
    assert row == 42 and score == 100, (
        f"popped ({row},{score}), expected (42,100)"
    )
    assert dut.empty_o.value == 1, "depth-1 FIFO should be empty after one pop"
    assert dut.full_o.value == 0

    await push(dut, 7, 55)
    assert dut.full_o.value == 1
    assert dut.empty_o.value == 0

    row, score = await pop(dut)
    assert row == 7 and score == 55

    assert dut.empty_o.value == 1
    assert dut.full_o.value == 0