from __future__ import annotations

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


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


async def reset_popcount(dut):
    """Assert reset for 5 cycles, then release and wait 2 cycles."""
    dut.rst_n.value = 0
    dut.valid_i.value = 0
    dut.row_i.value = 0
    dut.bits_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 drive_and_expect(dut, row: int, bits: int, expected_count: int):
    """Drive valid_i for one cycle, wait for valid_o, assert row_o / count_o."""
    dut.row_i.value = row
    dut.bits_i.value = bits
    dut.valid_i.value = 1
    await RisingEdge(dut.clk)
    dut.valid_i.value = 0

    # Pipeline depth = 3 (valid_s1_q -> valid_s2_q -> valid_o); wait with timeout
    for _ in range(20):
        await RisingEdge(dut.clk)
        if int(dut.valid_o.value):
            break

    assert int(dut.valid_o.value) == 1, "valid_o never asserted"
    assert int(dut.row_o.value) == row, (
        f"row_o mismatch: got {int(dut.row_o.value)}, expected {row}"
    )
    assert int(dut.count_o.value) == expected_count, (
        f"count_o mismatch: got {int(dut.count_o.value)}, expected {expected_count}"
    )


def _width(dut) -> int:
    """Return the WIDTH of bits_i in the DUT."""
    return len(dut.bits_i)


# ── Test ─────────────────────────────────────────────────────────────────────


@cocotb.test()
async def popcount_pipeline_counts_bits_and_preserves_row_metadata(dut):
    """
    Popcount pipeline: for each deterministic vector, assert row_o matches row_i
    and count_o equals the bit-count of bits_i.
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_popcount(dut)

    w = _width(dut)
    all_ones = (1 << w) - 1

    # Fixed literal: "0123456789ABCDEF" repeated enough for width, masked.
    hex_repeat = (w + 63) // 64
    fixed_literal = int("0123456789ABCDEF" * hex_repeat, 16) & all_ones

    # Alternating 0xAA pattern
    alternating_aa = int.from_bytes(b"\xAA" * (w // 8), "little")

    vectors = [
        (0, 0, 0),
        (1, all_ones, w),
        (2, 1 << 0, 1),
        (3, 1 << min(127, w - 1), 1),
        (4, 1 << min(511, w - 1), 1),
        (5, alternating_aa, alternating_aa.bit_count()),
        (6, fixed_literal, fixed_literal.bit_count()),
    ]

    for row, bits, expected in vectors:
        await drive_and_expect(dut, row, bits, expected)