from __future__ import annotations

import cocotb
import numpy as np
from cocotb.clock import Clock
from cocotb.triggers import RisingEdge
from model.ref_model import (  # noqa: E402
    match_top1,
    pack_lanes_flat,
    random_hashes,
    unpack_score_debug_flat,
)

NUM_ROWS = 512
HASH_BITS = 512
LANES = 16
SCORE_BITS = 10


async def reset_dut(dut):
    dut.rst_n.value = 0
    dut.wr_en.value = 0
    dut.wr_row.value = 0
    dut.wr_hash.value = 0
    dut.query_valid.value = 0
    dut.query_hash.value = 0
    dut.result_ready.value = 1

    if hasattr(dut, "noise_mask_lanes_flat"):
        dut.noise_mask_lanes_flat.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 write_rows(dut, rows):
    for idx, value in enumerate(rows):
        dut.wr_row.value = idx
        dut.wr_hash.value = int(value)
        dut.wr_en.value = 1
        await RisingEdge(dut.clk)

    dut.wr_en.value = 0
    await RisingEdge(dut.clk)


async def query_once(dut, query, noise_masks=None):
    dut.query_hash.value = int(query)
    dut.query_valid.value = 1

    await RisingEdge(dut.clk)
    dut.query_valid.value = 0

    # Feed lane noise masks batch by batch while DUT is scanning.
    # For no-noise builds this signal is absent and ignored.
    base = 0
    while int(dut.result_valid.value) == 0:
        if hasattr(dut, "noise_mask_lanes_flat") and noise_masks is not None:
            lane_masks = []
            for lane in range(LANES):
                row = base + lane
                lane_masks.append(noise_masks[row] if row < NUM_ROWS else 0)
            dut.noise_mask_lanes_flat.value = pack_lanes_flat(
                lane_masks, width=HASH_BITS
            )
            base += LANES

        await RisingEdge(dut.clk)

    top1_index = int(dut.top1_index.value)
    top1_score = int(dut.top1_score.value)

    score_debug = None
    if hasattr(dut, "score_debug_flat"):
        score_debug = unpack_score_debug_flat(
            int(dut.score_debug_flat.value),
            NUM_ROWS,
            SCORE_BITS,
        )

    await RisingEdge(dut.clk)
    return top1_index, top1_score, score_debug


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

    rng = np.random.default_rng(1)
    rows = random_hashes(rng, NUM_ROWS, width=HASH_BITS)
    query_index = 123
    query = rows[query_index]

    await write_rows(dut, rows)
    top1_index, top1_score, score_debug = await query_once(dut, query)

    expected = match_top1(query, rows, width=HASH_BITS)

    assert top1_index == expected.top1_index
    assert top1_score == expected.top1_score
    assert top1_index == query_index
    assert top1_score == HASH_BITS

    if score_debug is not None:
        assert np.array_equal(score_debug, expected.scores)


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

    rows = [0] * NUM_ROWS
    rows[0] = 0
    rows[1] = (1 << HASH_BITS) - 1

    query = 0
    await write_rows(dut, rows)
    top1_index, top1_score, score_debug = await query_once(dut, query)

    assert top1_score == HASH_BITS
    assert top1_index == 0

    if score_debug is not None:
        assert int(score_debug[0]) == HASH_BITS
        assert int(score_debug[1]) == 0


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

    query = 0
    rows = [0] * NUM_ROWS
    rows[10] = (1 << 7) - 1  # Hamming distance = 7
    rows[11] = (1 << 31) - 1  # Hamming distance = 31
    rows[12] = (1 << 128) - 1  # Hamming distance = 128

    await write_rows(dut, rows)
    top1_index, top1_score, score_debug = await query_once(dut, query)

    # Many rows are identical to query; tie-break must select row 0.
    assert top1_index == 0
    assert top1_score == HASH_BITS

    if score_debug is not None:
        assert int(score_debug[10]) == HASH_BITS - 7
        assert int(score_debug[11]) == HASH_BITS - 31
        assert int(score_debug[12]) == HASH_BITS - 128


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

    rng = np.random.default_rng(2)
    rows = random_hashes(rng, NUM_ROWS, width=HASH_BITS)
    query = rows[200]
    rows[10] = query
    rows[20] = query
    rows[200] = query

    await write_rows(dut, rows)
    top1_index, top1_score, _ = await query_once(dut, query)

    assert top1_index == 10
    assert top1_score == HASH_BITS


@cocotb.test()
async def external_noise_mask(dut):
    # This test is meaningful only when compiled with SIM_NOISE and SIM_DEBUG.
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    if not hasattr(dut, "noise_mask_lanes_flat"):
        dut._log.warning("SIM_NOISE not enabled; skipping exact noise-mask behavior.")
        return

    rng = np.random.default_rng(3)
    rows = random_hashes(rng, NUM_ROWS, width=HASH_BITS)
    query_index = 42
    query = rows[query_index]

    noise_masks = [0] * NUM_ROWS
    noise_masks[query_index] = (1 << 13) - 1  # flip exactly 13 bits

    await write_rows(dut, rows)
    top1_index, top1_score, score_debug = await query_once(
        dut,
        query,
        noise_masks=noise_masks,
    )

    expected = match_top1(query, rows, width=HASH_BITS, noise_masks=noise_masks)

    assert top1_index == expected.top1_index
    assert top1_score == expected.top1_score

    if score_debug is not None:
        assert int(score_debug[query_index]) == HASH_BITS - 13
        assert np.array_equal(score_debug, expected.scores)