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,
    match_top1_with_read_noise,
    random_hashes,
    unpack_score_debug_flat,
)

DEFAULT_NUM_ROWS = 4096
DEFAULT_HASH_BITS = 512
DEFAULT_LANES = 8
DEFAULT_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


def dut_num_rows(dut):
    val = _get_param(dut, "NUM_ROWS", None)
    if val is not None:
        return val
    # Derive from wr_addr width (ROW_BITS): NUM_ROWS = 2^ROW_BITS
    return 1 << len(dut.wr_addr)


def dut_hash_bits(dut):
    val = _get_param(dut, "HASH_BITS", None)
    if val is not None:
        return val
    # Derive from write_hash signal width
    return len(dut.write_hash)


def dut_lanes(dut):
    val = _get_param(dut, "LANES", None)
    if val is not None:
        return val
    # Derive from rd_resp_row_ids width / ROW_BITS
    return len(dut.rd_resp_row_ids) // len(dut.wr_addr)


def dut_score_bits(dut):
    val = _get_param(dut, "SCORE_BITS", None)
    if val is not None:
        return val
    # Derive from top1_score signal width
    return len(dut.top1_score)


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


async def collect_topk(dut, timeout_cycles: int = 2000):
    """Collect all serial Top-K result beats with result_ready held high.

    Returns list of (rank, row, score, last) tuples.
    Raises AssertionError if the stream does not complete within timeout.
    """
    dut.result_ready.value = 1
    beats = []
    for _ in range(timeout_cycles):
        if int(dut.result_valid.value):
            rank = int(dut.result_rank.value)
            row = int(dut.result_row.value)
            score = int(dut.result_score.value)
            last = int(dut.result_last.value)
            beats.append((rank, row, score, last))
            if last:
                return beats
        await RisingEdge(dut.clk)
    raise AssertionError("Top-K result stream did not finish")


async def query_topk_once(dut, query, timeout_cycles=2000):
    """Issue a query, collect full serial Top-K stream, and return beats + Top-1 metadata.

    Returns (beats, top1_index, top1_score, score_debug).
    After this call the full result stream has been consumed and the DUT is idle.
    """
    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

    # Consume full serial result stream
    beats = await collect_topk(dut, timeout_cycles=timeout_cycles)

    # score_debug is available after query completes
    num_rows = dut_num_rows(dut)
    score_bits = dut_score_bits(dut)
    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,
        )

    return beats, beats[0][1], beats[0][2], score_debug


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."""
    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 write pipeline to drain
    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).

    Consumes the full serial Top-K stream and returns rank-0 data.
    """
    _, top1_index, top1_score, score_debug = await query_topk_once(dut, query)
    return top1_index, top1_score, score_debug


# ── Compile smoke test ────────────────────────────────────────────────────────


@cocotb.test()
async def compile_includes_grouped_noise_helper(dut):
    """Compilation test: new grouped noise helper must elaborate with cam_top."""
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)
    assert int(dut.wr_ready.value) in (0, 1)


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


@cocotb.test()
async def baseline_no_noise(dut):
    """Verify write+query works exactly like the old CAM when noise disabled."""
    noise_en = _get_param(dut, "WRITE_NOISE_EN", 0)
    read_noise_en = _get_param(dut, "READ_NOISE_EN", 0)
    if noise_en or read_noise_en:
        dut._log.info("Skipping baseline_no_noise: requires noise disabled.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)
    rng = np.random.default_rng(1)
    rows = random_hashes(rng, num_rows, width=hash_bits)
    query_index = min(123, num_rows - 1)
    query = rows[query_index]

    await write_rows(dut, rows)
    beats, top1_index, top1_score, score_debug = await query_topk_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

    # Serial Top-K stream verification: beats from the single query above
    assert beats[0][0] == 0, "First beat must have rank 0"
    assert beats[-1][3] == 1, "Last beat must assert result_last"
    # Verify top1 aliases match first beat after stream fully consumed
    assert int(dut.top1_index.value) == beats[0][1]
    assert int(dut.top1_score.value) == beats[0][2]
    # Verify returned top1 matches first beat rank0
    assert top1_index == beats[0][1]
    assert top1_score == beats[0][2]

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


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


@cocotb.test()
async def zero_rate_noise(dut):
    """Noise module connected but THRESHOLD=0 → no flips."""
    noise_en = _get_param(dut, "WRITE_NOISE_EN", 1)
    rate_num = _get_param(dut, "WRITE_NOISE_RATE_NUM", 1)
    read_noise_en = _get_param(dut, "READ_NOISE_EN", 0)
    if not noise_en or rate_num != 0 or read_noise_en:
        dut._log.info("Skipping zero_rate_noise: requires WRITE_NOISE_EN=1, RATE_NUM=0, READ_NOISE_EN=0.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)
    rng = np.random.default_rng(1)
    rows = random_hashes(rng, num_rows, width=hash_bits)
    query_index = min(123, num_rows - 1)
    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):
    """WRITE_NOISE_RATE_NUM=1, WRITE_NOISE_RATE_DEN=1 → every group flips."""
    noise_en = _get_param(dut, "WRITE_NOISE_EN", 1)
    rate_num = _get_param(dut, "WRITE_NOISE_RATE_NUM", 1)
    rate_den = _get_param(dut, "WRITE_NOISE_RATE_DEN", 100)
    if not noise_en or rate_num != 1 or rate_den != 1:
        dut._log.info("Skipping full_rate_noise: requires WRITE_NOISE_EN=1, RATE_NUM=1, RATE_DEN=1.")
        return
    if not hasattr(dut, "score_debug_flat"):
        dut._log.info("Skipping full_rate_noise: requires SIM_DEBUG (score_debug_flat).")
        return

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

    hash_bits = dut_hash_bits(dut)
    num_rows = dut_num_rows(dut)
    noise_bits = _get_param(dut, "WRITE_NOISE_BITS", 8)
    all_zero = 0
    all_one = (1 << hash_bits) - 1

    RTL_SEED = 0xB504_F32D_B504_F32D
    prng_state = (RTL_SEED << 64) | RTL_SEED

    flip0, prng_state = generate_write_flip_mask(
        prng_state, hash_bits, noise_bits, rate_num, rate_den,
    )
    expected_row0 = all_zero ^ flip0

    flip1, prng_state = generate_write_flip_mask(
        prng_state, hash_bits, noise_bits, rate_num, rate_den,
    )
    expected_row1 = all_one ^ flip1

    rows = [0] * num_rows
    rows[0] = all_zero
    rows[1] = all_one

    await write_rows(dut, rows)

    top1_index, top1_score, score_debug = await query_once(dut, expected_row0)
    assert score_debug is not None, "score_debug required for full_rate_noise"
    assert int(score_debug[0]) == hash_bits, (
        f"Row 0: expected exact match, score={score_debug[0]} != {hash_bits}"
    )

    top1_index, top1_score, score_debug = await query_once(dut, expected_row1)
    assert score_debug is not None
    assert int(score_debug[1]) == hash_bits, (
        f"Row 1: expected exact match, score={score_debug[1]} != {hash_bits}"
    )


# ── 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, "WRITE_NOISE_EN", 1)
    if not noise_en:
        dut._log.info("Skipping default_noise_reproducible: requires WRITE_NOISE_EN=1.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)
    rng = np.random.default_rng(42)
    rows = random_hashes(rng, num_rows, width=hash_bits)

    await write_rows(dut, rows)
    query = rows[min(50, num_rows - 1)]
    top1_index_1, top1_score_1, _ = await query_once(dut, query)

    await reset_dut(dut)

    await write_rows(dut, rows)
    top1_index_2, top1_score_2, _ = await query_once(dut, query)

    assert top1_index_1 == top1_index_2
    assert top1_score_1 == top1_score_2


# ── Preserved legacy tests (only meaningful for noise disabled) ──────────────


@cocotb.test()
async def known_hamming_distance(dut):
    """Hamming distance verification — exact scores only valid without noise."""
    if _get_param(dut, "WRITE_NOISE_EN", 1) or _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping known_hamming_distance: requires noise disabled.")
        return

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

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)

    query = 0
    rows = [0] * num_rows
    rows[min(10, num_rows - 1)] = (1 << 7) - 1
    rows[min(11, num_rows - 1)] = (1 << 31) - 1
    rows[min(12, num_rows - 1)] = (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[min(10, num_rows - 1)]) == hash_bits - 7
        assert int(score_debug[min(11, num_rows - 1)]) == hash_bits - 31
        assert int(score_debug[min(12, num_rows - 1)]) == 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, "WRITE_NOISE_EN", 1) or _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping tie_break_policy: requires noise disabled.")
        return

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

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)
    rng = np.random.default_rng(2)
    rows = random_hashes(rng, num_rows, width=hash_bits)
    query = rows[min(200, num_rows - 1)]
    rows[10] = query
    rows[20] = query
    rows[min(200, num_rows - 1)] = 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, "WRITE_NOISE_EN", 1) or _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping all_zero_all_one_boundary: requires noise disabled.")
        return

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

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(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."""
    noise_en = _get_param(dut, "WRITE_NOISE_EN", 1)
    rate_num = _get_param(dut, "WRITE_NOISE_RATE_NUM", 1)
    rate_den = _get_param(dut, "WRITE_NOISE_RATE_DEN", 100)
    if not noise_en or rate_num == 0:
        dut._log.info("Skipping exact_noise_model_match: requires WRITE_NOISE_EN=1, RATE_NUM>0.")
        return
    if _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping exact_noise_model_match: requires READ_NOISE_EN=0 (read noise corrupts score comparison).")
        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)

    hash_bits = dut_hash_bits(dut)
    noise_bits = _get_param(dut, "WRITE_NOISE_BITS", 8)

    n_test_rows = 4
    rng = np.random.default_rng(99)
    rows = random_hashes(rng, n_test_rows, width=hash_bits)

    RTL_SEED = 0xB504_F32D_B504_F32D
    prng_state = (RTL_SEED << 64) | RTL_SEED
    expected_stored = []
    for row in rows:
        flip, prng_state = generate_write_flip_mask(
            prng_state, hash_bits, noise_bits, rate_num, rate_den,
        )
        expected_stored.append(row ^ flip)

    for idx, val in enumerate(rows):
        await write_row(dut, idx, val)

    for idx, expected in enumerate(expected_stored):
        top1_index, top1_score, score_debug = await query_once(dut, expected)
        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."""
    if _get_param(dut, "WRITE_NOISE_EN", 1) or _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping half_duplex_write_priority: requires noise disabled.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    hash_bits = dut_hash_bits(dut)
    test_val = (1 << hash_bits) - 1
    await write_row(dut, 0, test_val)

    await wait_idle(dut)
    assert int(dut.wr_ready.value) == 1
    assert int(dut.query_ready.value) == 1

    dut.wr_valid.value = 1
    dut.wr_addr.value = 1
    dut.write_hash.value = 0
    dut.query_valid.value = 1
    dut.query_hash.value = test_val

    await RisingEdge(dut.clk)

    dut.wr_valid.value = 0
    dut.query_valid.value = 0

    await wait_idle(dut)

    top1_index, top1_score, _ = await query_once(dut, test_val)
    assert top1_index == 0
    assert top1_score == hash_bits

    top1_index, top1_score, _ = await query_once(dut, 0)
    assert top1_index == 1
    assert top1_score == hash_bits


# ── Test G: Banked pipeline no-noise Top-1 ───────────────────────────────────


@cocotb.test()
async def banked_pipeline_no_noise_top1(dut):
    """No-noise banked pipeline returns the same Top-1 as the pure model."""
    if _get_param(dut, "WRITE_NOISE_EN", 0) or _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping banked_pipeline_no_noise_top1: requires noise disabled.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)
    rng = np.random.default_rng(7)
    rows = random_hashes(rng, num_rows, width=hash_bits)
    query_index = min(17, num_rows - 1)
    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


# ── Test H: Query scan blocks writes until result consumed ───────────────────


@cocotb.test()
async def query_scan_blocks_writes_until_result_consumed(dut):
    """Half-duplex: active query scan deasserts wr_ready."""
    if _get_param(dut, "WRITE_NOISE_EN", 0) or _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping query_scan_blocks_writes: requires noise disabled.")
        return
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_dut(dut)

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)
    rows = [0] * num_rows
    rows[0] = (1 << hash_bits) - 1
    await write_rows(dut, rows)

    await wait_idle(dut)
    dut.query_hash.value = rows[0]
    dut.query_valid.value = 1
    await RisingEdge(dut.clk)
    dut.query_valid.value = 0

    dut.wr_valid.value = 1
    dut.wr_addr.value = 1
    dut.write_hash.value = 0
    await RisingEdge(dut.clk)
    assert int(dut.wr_ready.value) == 0
    dut.wr_valid.value = 0

    # Consume full serial stream so the DUT returns idle
    beats = await collect_topk(dut, timeout_cycles=2000)
    assert len(beats) > 0
    assert beats[-1][3] == 1  # last asserted


# ── Test I: Read noise model match ──────────────────────────────────────────


@cocotb.test()
async def read_noise_model_match(dut):
    """Read noise uses grouped masks and matches the Python model for one query."""
    if not _get_param(dut, "READ_NOISE_EN", 0):
        dut._log.info("Skipping read_noise_model_match: requires READ_NOISE_EN=1.")
        return

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

    num_rows = dut_num_rows(dut)
    hash_bits = dut_hash_bits(dut)
    lanes = dut_lanes(dut)
    rng = np.random.default_rng(123)
    rows = random_hashes(rng, num_rows, width=hash_bits)

    # If write noise is enabled, apply write flip masks to predict stored rows
    stored_rows = list(rows)
    if _get_param(dut, "WRITE_NOISE_EN", 0):
        seed = 0xB504_F32D_B504_F32D
        prng_state = (seed << 64) | seed
        stored_rows = []
        for row in rows:
            flip, prng_state = generate_write_flip_mask(
                prng_state,
                hash_bits,
                _get_param(dut, "WRITE_NOISE_BITS", 8),
                _get_param(dut, "WRITE_NOISE_RATE_NUM", 1),
                _get_param(dut, "WRITE_NOISE_RATE_DEN", 100),
            )
            stored_rows.append(row ^ flip)

    query = rows[min(5, num_rows - 1)]

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

    expected = match_top1_with_read_noise(
        query,
        stored_rows,
        width=hash_bits,
        lanes=lanes,
        noise_bits=_get_param(dut, "READ_NOISE_BITS", 8),
        rate_num=_get_param(dut, "READ_NOISE_RATE_NUM", 1),
        rate_den=_get_param(dut, "READ_NOISE_RATE_DEN", 100),
        seed=0x6A09_E667_F3BC_C909,
    )

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