from __future__ import annotations

import cocotb
from cocotb.clock import Clock
from cocotb.triggers import RisingEdge, ReadOnly, Timer


# ── Helpers (local reimplementation — no imports from test_cam_basic) ─────────


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
    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
    return len(dut.write_hash)


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


def _pipeline_depth(num_rows, lanes):
    """Number of match-engine pipeline stages (rows processed per query)."""
    return num_rows // lanes


# ---------------------------------------------------------------------------
# Bounded wait helper  (RisingEdge only — keep the caller writeable)
# ---------------------------------------------------------------------------


async def _wait_bounded(dut, condition_fn, max_cycles, label):
    """Wait for *condition_fn* to become ``True``, sampling at each RisingEdge.

    Raises ``AssertionError`` with *label* if *max_cycles* is exceeded.
    The caller is left in a writeable phase after return.
    """
    for _ in range(max_cycles):
        await RisingEdge(dut.clk)
        if condition_fn():
            return
    raise AssertionError(
        f"{label}: timed out after {max_cycles} clock cycles"
    )


def _condition_signal_high(signal):
    """Return a nullary callable that checks *signal* is high (integer 1)."""
    def _check():
        return int(signal.value) == 1
    return _check


# ---------------------------------------------------------------------------
# DUT helpers
# ---------------------------------------------------------------------------


async def reset_dut(dut):
    """Reset the DUT with new handshake interface.

    ``result_ready`` is held at 0 during measurement; the performance test
    pulses it high for exactly one cycle after capturing ``result_valid``.
    """
    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 = 0

    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, max_cycles=10):
    """Wait until both wr_ready=1 and query_ready=1 (system fully idle)."""
    await _wait_bounded(
        dut,
        lambda: int(dut.wr_ready.value) and int(dut.query_ready.value),
        max_cycles,
        "wait_idle",
    )


async def write_row(dut, addr, value, max_cycles=10):
    """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

    await _wait_bounded(
        dut,
        _condition_signal_high(dut.wr_ready),
        max_cycles,
        f"write_row(addr={addr}) handshake",
    )

    dut.wr_valid.value = 0
    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_with_latency(dut, query, max_result_cycles):
    """Issue a query and return (top1_index, top1_score, latency_cycles, total_cycles).

    Parameters
    ----------
    max_result_cycles : int
        Hard bound on cycles from query acceptance to *result_valid*.
        Derived from ``NUM_ROWS / LANES + pipeline_slack`` by the caller.

    Cycle counting
    --------------
    ``query_ready`` is a combinational signal that goes low *immediately* after
    the RisingEdge at which the query is accepted (the state machine
    transitions from S_IDLE to S_SCAN).  We therefore sample ``query_ready``
    **before** ``ReadOnly`` (i.e. at the RisingEdge time-point itself) to
    capture the handshake.

    ``result_valid`` is a registered output that stays high until consumed, so
    we sample it in the **settled phase** after ``ReadOnly``.

    ``latency_cycles`` is the number of RisingEdge events between the cycle
    where the query is accepted and the cycle where ``result_valid`` is
    observed.
    """
    await wait_idle(dut)

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

    # ── Phase 1: Accept handshake ───────────────────────────────────────
    #    query_ready is combinational — sample before ReadOnly.
    accept_edge = None
    for _ in range(10):
        await RisingEdge(dut.clk)
        edge_count += 1
        q_ready = int(dut.query_ready.value)  # sample before state transition
        await ReadOnly()
        await Timer(1, units="step")          # exit ReadOnly for driving
        if q_ready:
            accept_edge = edge_count
            break
    assert accept_edge is not None, (
        "Query accept handshake timed out after 10 cycles"
    )
    dut.query_valid.value = 0

    # ── Phase 2: Wait for result_valid (measurement window) ─────────────
    #    result_valid is registered — sample after ReadOnly is fine.
    result_edge = None
    for _ in range(max_result_cycles):
        await RisingEdge(dut.clk)
        edge_count += 1
        await ReadOnly()
        if int(dut.result_valid.value):
            result_edge = edge_count
            break
    assert result_edge is not None, (
        f"Query result_valid timed out after {max_result_cycles} cycles "
        f"(accepted at edge {accept_edge})"
    )
    await Timer(1, units="step")

    latency_cycles = result_edge - accept_edge

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

    # ── Phase 3: Consume result (pulse result_ready for one cycle) ──────
    dut.result_ready.value = 1
    await RisingEdge(dut.clk)
    edge_count += 1
    await ReadOnly()
    await Timer(1, units="step")
    dut.result_ready.value = 0

    return top1_index, top1_score, latency_cycles, edge_count


def deterministic_rows(num_rows, hash_bits, query_hash):
    """Create deterministic rows where only row 0 stores *query_hash*."""
    mask = (1 << hash_bits) - 1
    rows = [0] * num_rows
    rows[0] = query_hash
    for i in range(1, num_rows):
        # Deterministic non-matching value; golden-ratio-like spread
        val = ((i + 1) * 0x9E3779B97F4A7C15) & mask
        if val == query_hash or val == 0:
            val = ((val ^ query_hash) ^ 0xA5A5A5A5A5A5A5A5) & mask
            if val == query_hash or val == 0:
                val = (val ^ 1) & mask
        rows[i] = val
    return rows


# ── Performance Test ──────────────────────────────────────────────────────────


@cocotb.test()
async def cam_perf_benchmark(dut):
    """Performance benchmark: measure query latency in cycles."""
    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)
    write_noise_en = _get_param(dut, "WRITE_NOISE_EN", 1)
    read_noise_en = _get_param(dut, "READ_NOISE_EN", 0)

    # ── Deterministic query ─────────────────────────────────────────────
    query_hash = 0xAA55_AA55_AA55_AA55_AA55_AA55_AA55_AA55
    query_hash &= (1 << hash_bits) - 1
    if query_hash == 0:
        query_hash = 1

    rows = deterministic_rows(num_rows, hash_bits, query_hash)
    await write_rows(dut, rows)

    # Bound: pipeline depth plus generous slack for read + drain stages
    pipeline = _pipeline_depth(num_rows, lanes)
    max_result_cycles = pipeline + 30

    top1_index, top1_score, latency_cycles, total_cycles = (
        await query_once_with_latency(dut, query_hash, max_result_cycles)
    )

    # ── Performance assertions ──────────────────────────────────────────
    assert latency_cycles > 0, (
        f"latency_cycles must be positive, got {latency_cycles}"
    )
    assert total_cycles > 0, (
        f"total_cycles must be positive, got {total_cycles}"
    )

    # ── Correctness assertions (conditional on noise state) ─────────────
    if not write_noise_en and not read_noise_en:
        # Without noise: stored hash at row 0 == query_hash → exact match.
        assert top1_index == 0, (
            f"Noise disabled: expected top1_index=0 (exact match), got "
            f"{top1_index}"
        )
        assert top1_score == hash_bits, (
            f"Noise disabled: expected top1_score={hash_bits}, got "
            f"{top1_score}"
        )
    else:
        # With noise: write/read flip masks may corrupt stored values, so
        # we cannot reliably assert the exact match.  Instead, confirm a
        # valid non-zero score was produced (the match engine ran).
        assert top1_score > 0, (
            f"Noise enabled: expected top1_score > 0, got {top1_score}. "
            "Match engine returned invalid result."
        )
        dut._log.info(
            "Noise enabled (WRITE_NOISE_EN=%s, READ_NOISE_EN=%s) — "
            "skipping exact top1_index/top1_score assertion. "
            "top1_index=%d top1_score=%d",
            write_noise_en, read_noise_en, top1_index, top1_score,
        )

    # ── Machine-readable performance marker ─────────────────────────────
    queries_per_cycle = 1.0 / total_cycles
    dut._log.info(
        "PERF_RESULT latency_cycles=%d total_cycles=%d "
        "accepted_queries=1 completed_queries=1 "
        "queries_per_cycle=%.6f status=pass",
        latency_cycles, total_cycles, queries_per_cycle,
    )