from __future__ import annotations

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


async def reset_match(dut):
    dut.rst_n.value = 0
    dut.query_valid.value = 0
    dut.query_hash.value = 0
    dut.result_ready.value = 1
    dut.rd_valid_i.value = 0
    dut.rd_row_ids_i.value = 0
    dut.rd_hashes_i.value = 0
    dut.rd_lane_valid_i.value = 0
    for _ in range(3):
        await RisingEdge(dut.clk)
    dut.rst_n.value = 1
    for _ in range(2):
        await RisingEdge(dut.clk)


@cocotb.test()
async def match_engine_returns_top1_after_pipeline_drain(dut):
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_match(dut)

    LANES = 8
    ROW_BITS = len(dut.rd_row_ids_i) // LANES
    HASH_BITS = len(dut.rd_hashes_i) // LANES
    NUM_ROWS = 1 << len(dut.rd_base_row_o)
    TARGET_ROW = 9

    query = (1 << HASH_BITS) - 1
    dut.query_hash.value = query
    dut.query_valid.value = 1
    await RisingEdge(dut.clk)
    dut.query_valid.value = 0

    for base in range(0, NUM_ROWS, LANES):
        while int(dut.rd_valid_o.value) == 0:
            await RisingEdge(dut.clk)
        assert int(dut.rd_base_row_o.value) == base
        rows = 0
        hashes = 0
        for lane in range(LANES):
            row_id = base + lane
            rows |= row_id << (lane * ROW_BITS)
            row_hash = query if row_id == TARGET_ROW else 0
            hashes |= row_hash << (lane * HASH_BITS)
        dut.rd_row_ids_i.value = rows
        dut.rd_hashes_i.value = hashes
        dut.rd_lane_valid_i.value = 0xFF
        dut.rd_valid_i.value = 1
        await RisingEdge(dut.clk)
        dut.rd_valid_i.value = 0

    for _ in range(20):
        if int(dut.result_valid.value):
            break
        await RisingEdge(dut.clk)

    assert int(dut.result_valid.value) == 1
    assert int(dut.result_row.value) == TARGET_ROW
    assert int(dut.result_score.value) == HASH_BITS