from __future__ import annotations

import os

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

# ── Parameter defaults (matching cam_params.svh: NUM_ROWS=4096, HASH_BITS=512) ──
# K is read from env so Makefile override (e.g. make TOPK_K=1) propagates to tests.
K = int(os.environ.get('TOPK_K', '4'))
ROW_BITS = 12
SCORE_BITS = 10  # $clog2(512 + 1)
# Mirror the RTL's RANK_BITS formula: (K <= 1) ? 1 : $clog2(K)
RANK_BITS = 1 if K <= 1 else (K - 1).bit_length()


def pack(values: list[int], width: int) -> int:
    """Pack a list of values into a flattened integer (LSB = index 0)."""
    result = 0
    for idx, value in enumerate(values):
        result |= value << (idx * width)
    return result


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


async def reset_serializer(dut):
    """Drive reset active-low for 5 cycles, then release; zero all inputs."""
    dut.rst_n.value = 0
    dut.start_i.value = 0
    dut.result_ready_i.value = 0
    dut.topk_rows_i.value = 0
    dut.topk_scores_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 assert_output_stable(dut, rank: int, row: int, score: int,
                               last: bool, valid: bool):
    """Assert output signals match expected values."""
    assert int(dut.result_rank_o.value) == rank, \
        f"rank: expected {rank}, got {int(dut.result_rank_o.value)}"
    assert int(dut.result_row_o.value) == row, \
        f"row: expected {row}, got {int(dut.result_row_o.value)}"
    assert int(dut.result_score_o.value) == score, \
        f"score: expected {score}, got {int(dut.result_score_o.value)}"
    assert bool(dut.result_last_o.value) == last, \
        f"last: expected {last}, got {bool(dut.result_last_o.value)}"
    assert bool(dut.result_valid_o.value) == valid, \
        f"valid: expected {valid}, got {bool(dut.result_valid_o.value)}"


# ── Test 1: Serializer outputs rank order and last ─────────────────────────────


@cocotb.test()
async def serializer_outputs_rank_order_and_last(dut):
    """
    Start serialization with K=4 results. Clock through all 4 ranks, verifying
    row/score/last at each step. Assert done_o pulses after last rank.

    Gated: requires TOPK_K=4 (the default). Use ``make test-k1`` for K=1.
    """
    if K != 4:
        cocotb.log.info("Skipping (designed for TOPK_K=4, got %d)", K)
        return

    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_serializer(dut)

    # Pre-defined top-4 results
    rows = [42, 17, 99, 5]
    scores = [10, 9, 8, 7]

    # Drive packed inputs and assert start
    dut.topk_rows_i.value = pack(rows, ROW_BITS)
    dut.topk_scores_i.value = pack(scores, SCORE_BITS)
    dut.start_i.value = 1
    await RisingEdge(dut.clk)
    dut.start_i.value = 0

    # DUT should be busy, valid, showing rank 0
    assert bool(dut.busy_o.value), "busy_o should be high after start"
    await assert_output_stable(dut, rank=0, row=rows[0], score=scores[0],
                               last=False, valid=True)

    # Cycle through remaining ranks with ready held high
    for r in range(1, K):
        dut.result_ready_i.value = 1
        await RisingEdge(dut.clk)
        is_last = (r == K - 1)
        await assert_output_stable(dut, rank=r, row=rows[r], score=scores[r],
                                   last=is_last, valid=True)

    # After last rank fire, next cycle should show done_o and deassert busy
    await RisingEdge(dut.clk)
    assert bool(dut.done_o.value), "done_o should be high after last rank fire"
    assert not bool(dut.busy_o.value), "busy_o should be low after serialization done"
    assert not bool(dut.result_valid_o.value), \
        "result_valid_o should be low after serialization done"

    # done_o should be single-cycle pulse
    await RisingEdge(dut.clk)
    assert not bool(dut.done_o.value), "done_o should be single-cycle pulse"
    assert int(dut.result_rank_o.value) == 0, \
        "rank should reset to 0 after done"


# ── Test 2: Holds output when not ready ────────────────────────────────────────


@cocotb.test()
async def serializer_holds_output_when_not_ready(dut):
    """
    After start, with result_ready_i=0, output remains stable.
    After ready is asserted, advances to next rank.

    Gated: requires TOPK_K=4 (the default).  Use ``make test-k1`` for K=1.
    """
    if K != 4:
        cocotb.log.info("Skipping (designed for TOPK_K=4, got %d)", K)
        return

    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_serializer(dut)

    rows = [10, 20, 30, 40]
    scores = [5, 4, 3, 2]

    dut.topk_rows_i.value = pack(rows, ROW_BITS)
    dut.topk_scores_i.value = pack(scores, SCORE_BITS)
    dut.result_ready_i.value = 0  # not ready
    dut.start_i.value = 1
    await RisingEdge(dut.clk)
    dut.start_i.value = 0

    # Output should present rank 0 but hold because ready is low
    await assert_output_stable(dut, rank=0, row=rows[0], score=scores[0],
                               last=False, valid=True)

    # Stall for 3 cycles — mutating inputs to verify they don't leak through
    # (snapshot must hold original values captured at start)
    for cycle in range(3):
        # Corrupt the live inputs — a buggy combinational slice would change
        dut.topk_rows_i.value = pack([cycle + 100, 0, 0, 0], ROW_BITS)
        dut.topk_scores_i.value = pack([cycle + 50, 0, 0, 0], SCORE_BITS)
        await RisingEdge(dut.clk)
        await assert_output_stable(dut, rank=0, row=rows[0], score=scores[0],
                                   last=False, valid=True)

    # Restore original inputs and assert ready for 1 cycle; rank advances to 1
    dut.topk_rows_i.value = pack(rows, ROW_BITS)
    dut.topk_scores_i.value = pack(scores, SCORE_BITS)
    dut.result_ready_i.value = 1
    await RisingEdge(dut.clk)
    await assert_output_stable(dut, rank=1, row=rows[1], score=scores[1],
                               last=False, valid=True)

    # Stall again at rank 1 for 2 cycles — mutate inputs again
    dut.result_ready_i.value = 0
    for cycle in range(2):
        dut.topk_rows_i.value = pack([cycle + 200, 0, 0, 0], ROW_BITS)
        dut.topk_scores_i.value = pack([cycle + 100, 0, 0, 0], SCORE_BITS)
        await RisingEdge(dut.clk)
        await assert_output_stable(dut, rank=1, row=rows[1], score=scores[1],
                                   last=False, valid=True)


# ── Test 3: K=1 edge case ──────────────────────────────────────────────────────


@cocotb.test()
async def serializer_handles_k1(dut):
    """
    When K=1 (overridden via generics), start and fire produce done_o and
    deassert busy on the next cycle.  Run via ``make test-k1``.
    """
    if K != 1:
        cocotb.log.info("Skipping (designed for TOPK_K=1, got %d)", K)
        return

    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_serializer(dut)

    rows = [100]
    scores = [50]

    dut.topk_rows_i.value = pack(rows, ROW_BITS)
    dut.topk_scores_i.value = pack(scores, SCORE_BITS)
    dut.start_i.value = 1
    await RisingEdge(dut.clk)
    dut.start_i.value = 0

    # K=1: rank 0 is also the last; valid asserted
    assert bool(dut.busy_o.value), "busy_o should be high after start"
    await assert_output_stable(dut, rank=0, row=rows[0], score=scores[0],
                               last=True, valid=True)

    # Fire the one and only result
    dut.result_ready_i.value = 1
    await RisingEdge(dut.clk)

    # done_o should pulse, busy deasserted
    assert bool(dut.done_o.value), "done_o should pulse after K=1 fire"
    assert not bool(dut.busy_o.value), "busy_o should be low after done"
    assert not bool(dut.result_valid_o.value), "valid should be low after done"
    assert int(dut.result_rank_o.value) == 0, "rank should reset to 0"