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
    generate_write_flip_mask,
    match_top1,
    random_hashes,
    unpack_score_debug_flat,
    xorshift64,
)

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


def _get_param(dut, name, default=None):
    """Read a Verilator-exposed parameter from the DUT."""
    try:
        val = getattr(dut, name, None)
        if val is not None:
            return int(val.value)
    except Exception:
        pass
    return default


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


async def reset_dut(dut):
    """Reset the DUT with new handshake interface."""
    dut.rst_n.value = 0
    dut.wr_valid.value = 0
    dut.wr_addr.value = 0
    dut.write_hash.value = 0
    dut.query_valid.value = 0
    dut.query_hash.value = 0
    dut.result_ready.value = 1

    for _ in range(5):
        await RisingEdge(dut.clk)

    dut.rst_n.value = 1

    for _ in range(2):
        await RisingEdge(dut.clk)


async def wait_idle(dut):
    """Wait until both wr_ready=1 and query_ready=1 (system fully idle)."""
    while not (int(dut.wr_ready.value) and int(dut.query_ready.value)):
        await RisingEdge(dut.clk)


async def write_row(dut, addr, value):
    """Write a single row using wr_valid/wr_ready handshake.

    1. Wait for system idle
    2. Assert wr_valid + set addr/hash
    3. Wait for handshake (wr_ready=1 on clock edge)
    4. Deassert wr_valid
    5. Wait for wr_ready to return 1 (commit complete)
    """
    await wait_idle(dut)

    dut.wr_addr.value = addr
    dut.write_hash.value = int(value)
    dut.wr_valid.value = 1

    # Wait for handshake
    while True:
        await RisingEdge(dut.clk)
        if int(dut.wr_ready.value):
            break

    dut.wr_valid.value = 0

    # Wait for cam_noisy to finish GEN_MASK/COMMIT
    await wait_idle(dut)


async def write_rows(dut, rows):
    """Write all rows sequentially."""
    for idx, value in enumerate(rows):
        await write_row(dut, idx, value)


async def query_once(dut, query):
    """Issue a query and return (top1_index, top1_score, score_debug).

    1. Wait for system idle
    2. Assert query_valid + set query_hash
    3. Wait for query_ready handshake
    4. Deassert query_valid
    5. Wait for result_valid
    6. Read result, pulse result_ready to consume
    """
    await wait_idle(dut)

    dut.query_hash.value = int(query)
    dut.query_valid.value = 1

    # Wait for handshake
    while True:
        await RisingEdge(dut.clk)
        if int(dut.query_ready.value):
            break

    dut.query_valid.value = 0

    # Wait for result
    while int(dut.result_valid.value) == 0:
        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,
        )

    dut.result_ready.value = 1
    await RisingEdge(dut.clk)
    dut.result_ready.value = 0

    return top1_index, top1_score, score_debug


# ── Test A: Baseline (NOISE_EN=0) ────────────────────────────────────────────


@cocotb.test()
async def baseline_no_noise(dut):
    """Verify write+query works exactly like the old CAM when NOISE_EN=0."""
    noise_en = _get_param(dut, "NOISE_EN", 0)
    if noise_en:
        dut._log.info("Skipping baseline_no_noise: requires NOISE_EN=0.")
        return
    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)


# ── Test B: Zero noise rate (NOISE_EN=1, RATE_NUM=0) ────────────────────────


@cocotb.test()
async def zero_rate_noise(dut):
    """Noise module connected but THRESHOLD=0 → no flips, equivalent to NOISE_EN=0."""
    noise_en = _get_param(dut, "NOISE_EN", 1)
    rate_num = _get_param(dut, "NOISE_RATE_NUM", 1)
    if not noise_en or rate_num != 0:
        dut._log.info("Skipping zero_rate_noise: requires NOISE_EN=1, RATE_NUM=0.")
        return
    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)


# ── Test C: 100% noise rate (RATE_NUM=1, RATE_DEN=1) ───────────────────────


@cocotb.test()
async def full_rate_noise(dut):
    """THRESHOLD=256 → all bits flip. stored == ~written."""
    noise_en = _get_param(dut, "NOISE_EN", 1)
    rate_num = _get_param(dut, "NOISE_RATE_NUM", 1)
    rate_den = _get_param(dut, "NOISE_RATE_DEN", 100)
    if not noise_en or rate_num != 1 or rate_den != 1:
        dut._log.info("Skipping full_rate_noise: requires NOISE_EN=1, RATE_NUM=1, RATE_DEN=1.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    all_zero = 0
    all_one = (1 << HASH_BITS) - 1

    # Write all-zero to row 0, all-one to row 1, rest zero
    rows = [0] * NUM_ROWS
    rows[0] = all_zero
    rows[1] = all_one

    await write_rows(dut, rows)

    # After 100% flip: row 0 stored = ~0 = all_one, row 1 stored = ~all_one = all_zero
    # Query all-zero → matches row 1 (stored all_zero) with score = HASH_BITS
    top1_index, top1_score, _ = await query_once(dut, all_zero)
    assert top1_index == 1
    assert top1_score == HASH_BITS

    # Query all-one → matches row 0 (stored all_one) with score = HASH_BITS
    top1_index, top1_score, _ = await query_once(dut, all_one)
    assert top1_index == 0
    assert top1_score == HASH_BITS

    # Query all-zero → score against row 0 (stored all_one) = 0
    # Re-query to verify: row 0 stored is all_one, query all_zero → score 0
    # But top1 picks row 1 with score HASH_BITS, so top1_index=1


# ── Test D: Default ~1% noise, reproducible ────────────────────────────────


@cocotb.test()
async def default_noise_reproducible(dut):
    """Fixed seed → deterministic write noise. Two identical runs produce same results."""
    noise_en = _get_param(dut, "NOISE_EN", 1)
    if not noise_en:
        dut._log.info("Skipping default_noise_reproducible: requires NOISE_EN=1.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    rng = np.random.default_rng(42)
    rows = random_hashes(rng, NUM_ROWS, width=HASH_BITS)

    # ── First run ──
    await write_rows(dut, rows)
    query = rows[50]
    top1_index_1, top1_score_1, _ = await query_once(dut, query)

    # Reset for second run
    await reset_dut(dut)

    # ── Second run with same data ──
    await write_rows(dut, rows)
    top1_index_2, top1_score_2, _ = await query_once(dut, query)

    # Deterministic: same seed → same PRNG sequence → same stored hashes → same result
    assert top1_index_1 == top1_index_2
    assert top1_score_1 == top1_score_2


# ── Preserved legacy tests (only meaningful for NOISE_EN=0) ──────────────────


@cocotb.test()
async def known_hamming_distance(dut):
    """Hamming distance verification — exact scores only valid without noise."""
    if _get_param(dut, "NOISE_EN", 1):
        dut._log.info("Skipping known_hamming_distance: NOISE_EN=1, stored hashes may differ.")
        return

    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
    rows[11] = (1 << 31) - 1
    rows[12] = (1 << 128) - 1

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

    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):
    """Tie-break: lowest row index wins — only verified without noise."""
    if _get_param(dut, "NOISE_EN", 1):
        dut._log.info("Skipping tie_break_policy: NOISE_EN=1, stored hashes may differ.")
        return

    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 all_zero_all_one_boundary(dut):
    """All-zero / all-one boundary — only verified without noise."""
    if _get_param(dut, "NOISE_EN", 1):
        dut._log.info("Skipping all_zero_all_one_boundary: NOISE_EN=1, stored hashes may differ.")
        return

    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


# ── Test E: Exact RTL-vs-model PRNG mask match ──────────────────────────────


@cocotb.test()
async def exact_noise_model_match(dut):
    """Verify RTL stored hashes match ref_model.py for a known seed and rate.

    Writes rows with noise enabled, then queries back via score_debug to
    reconstruct stored hashes, and compares against Python model predictions.
    """
    noise_en = _get_param(dut, "NOISE_EN", 1)
    rate_num = _get_param(dut, "NOISE_RATE_NUM", 1)
    rate_den = _get_param(dut, "NOISE_RATE_DEN", 100)
    if not noise_en or rate_num == 0:
        dut._log.info("Skipping exact_noise_model_match: requires NOISE_EN=1, RATE_NUM>0.")
        return
    if not hasattr(dut, "score_debug_flat"):
        dut._log.info("Skipping exact_noise_model_match: requires SIM_DEBUG.")
        return

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

    noise_gen_bits = _get_param(dut, "NOISE_GEN_BITS", 8)
    noise_sample_bits = _get_param(dut, "NOISE_SAMPLE_BITS", 8)

    # Use a small subset to keep test fast
    n_test_rows = 4
    rng = np.random.default_rng(99)
    rows = random_hashes(rng, n_test_rows, width=HASH_BITS)

    # Predict stored hashes with Python model using the same seed
    # RTL default seed: 64'hB504_F32D_B504_F32D
    RTL_SEED = 0xB504_F32D_B504_F32D
    prng_state = RTL_SEED
    expected_stored = []
    for row in rows:
        flip, prng_state = generate_write_flip_mask(
            prng_state, HASH_BITS, noise_gen_bits, noise_sample_bits, rate_num, rate_den,
        )
        expected_stored.append(row ^ flip)

    # Write only test rows (rest stay at 0 from reset)
    for idx, val in enumerate(rows):
        await write_row(dut, idx, val)

    # Query all-zero to get Hamming distances (= HASH_BITS - popcount(stored ^ 0) = HASH_BITS - popcount(stored))
    # So popcount(stored) = HASH_BITS - score
    # This gives us the number of set bits but not the exact value.
    # Instead, query each expected_stored value — it should score HASH_BITS if match is exact.
    for idx, expected in enumerate(expected_stored):
        top1_index, top1_score, score_debug = await query_once(dut, expected)
        # The stored hash at idx should exactly match expected, so score == HASH_BITS
        assert score_debug is not None, "score_debug required for mask match verification"
        assert int(score_debug[idx]) == HASH_BITS, (
            f"Row {idx}: expected stored hash to match model prediction, "
            f"score={score_debug[idx]} != {HASH_BITS}"
        )


# ── Test F: Half-duplex write-priority arbitration ───────────────────────────


@cocotb.test()
async def half_duplex_write_priority(dut):
    """When wr_valid and query_valid are both high, write wins and query is held off.
    Only runs with NOISE_EN=0 so stored hashes are predictable.
    """
    if _get_param(dut, "NOISE_EN", 1):
        dut._log.info("Skipping half_duplex_write_priority: requires NOISE_EN=0 for exact scores.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    # Write a known value to row 0
    test_val = (1 << HASH_BITS) - 1  # all-ones
    await write_row(dut, 0, test_val)

    # Now system is idle: wr_ready=1, query_ready=1
    await wait_idle(dut)
    assert int(dut.wr_ready.value) == 1
    assert int(dut.query_ready.value) == 1

    # Drive both wr_valid and query_valid simultaneously
    dut.wr_valid.value = 1
    dut.wr_addr.value = 1
    dut.write_hash.value = 0  # write all-zeros to row 1
    dut.query_valid.value = 1
    dut.query_hash.value = test_val  # query for all-ones (in row 0)

    await RisingEdge(dut.clk)

    # Write should have been accepted (wr_ready was 1), query should NOT have been accepted
    # because write-priority gates query_ready when wr_valid=1
    wr_accepted = int(dut.wr_ready.value) == 0  # after handshake, wr_ready drops
    # query_ready should have been 0 during the simultaneous assertion
    # (it's !wr_valid gated), so query was blocked

    # Deassert both
    dut.wr_valid.value = 0
    dut.query_valid.value = 0

    # Wait for write to complete (noise generation + commit)
    await wait_idle(dut)

    # Now query should work — row 0 has all-ones (written first)
    top1_index, top1_score, _ = await query_once(dut, test_val)
    assert top1_index == 0
    assert top1_score == HASH_BITS

    # Verify row 1 was written (all-zeros) — query all-zeros
    top1_index, top1_score, _ = await query_once(dut, 0)
    assert top1_index == 1
    assert top1_score == HASH_BITS