from __future__ import annotations

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

# ── Parameter defaults (matching cam_params.svh with NUM_ROWS=4096, HASH_BITS=512) ──
K = 4
ROW_BITS = 12
SCORE_BITS = 10  # $clog2(512 + 1)
ROW_SENTINEL = (1 << ROW_BITS) - 1  # all-ones


def golden(candidates, k=K):
    """Return top-k candidates sorted by score desc, then row asc (higher score wins; lower row breaks ties)."""
    return sorted(candidates, key=lambda item: (-item[1], item[0]))[:k]


def unpack(dut, k=K):
    """Read packed topk_rows_o and topk_scores_o; return list of (row, score) tuples ordered by rank."""
    rows_packed = int(dut.topk_rows_o.value)
    scores_packed = int(dut.topk_scores_o.value)
    result = []
    for i in range(k):
        row = (rows_packed >> (i * ROW_BITS)) & ((1 << ROW_BITS) - 1)
        score = (scores_packed >> (i * SCORE_BITS)) & ((1 << SCORE_BITS) - 1)
        result.append((row, score))
    return result


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


async def reset_tracker(dut):
    """Assert reset for 5 cycles, release, then wait 2 cycles."""
    dut.rst_n.value = 0
    dut.clear_i.value = 0
    dut.candidate_valid_i.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 push_candidate(dut, row: int, score: int):
    """Drive one candidate handshake (candidate_ready_o is always 1)."""
    dut.candidate_valid_i.value = 1
    dut.candidate_row_i.value = row
    dut.candidate_score_i.value = score
    await RisingEdge(dut.clk)
    dut.candidate_valid_i.value = 0


# ── Test 1: Basic ordering ───────────────────────────────────────────────────────


@cocotb.test()
async def tracker_orders_by_score_then_lower_row(dut):
    """
    Push candidates [(5,10), (3,11), (2,11), (9,7), (1,11), (0,3)] and verify
    the tracker's Top-4 matches golden: [(1,11), (2,11), (3,11), (5,10)].
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_tracker(dut)

    candidates = [(5, 10), (3, 11), (2, 11), (9, 7), (1, 11), (0, 3)]

    for row, score in candidates:
        await push_candidate(dut, row, score)

    result = unpack(dut)
    expected = golden(candidates)
    assert result == expected, f"Expected {expected}, got {result}"
    assert dut.update_pending_o.value == 0, "update_pending_o should be 0"


# ── Test 2: Clear restarts query ─────────────────────────────────────────────────


@cocotb.test()
async def tracker_clear_starts_new_query(dut):
    """
    After candidate (7,60), pulse clear_i, send candidate (4,12).
    Expect rank 0 = (4,12), remaining slots = sentinel rows, score 0.
    """
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    await reset_tracker(dut)

    # Push (7, 60)
    await push_candidate(dut, 7, 60)

    # Pulse clear
    dut.clear_i.value = 1
    await RisingEdge(dut.clk)
    dut.clear_i.value = 0

    # Push (4, 12)
    await push_candidate(dut, 4, 12)

    result = unpack(dut)
    expected = [(4, 12)] + [(ROW_SENTINEL, 0)] * (K - 1)
    assert result == expected, f"Expected {expected}, got {result}"


# ── Test 3: Random validation against golden ─────────────────────────────────────


@cocotb.test()
async def tracker_matches_random_python_golden(dut):
    """
    Push 100 candidates (row in [0,63], score in [0,127]) seeded at 42.
    Verify final Top-4 matches Python golden model.
    """
    import random
    random.seed(42)

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

    max_row = 64
    max_score = 128
    candidates = []

    for _ in range(100):
        row = random.randint(0, max_row - 1)
        score = random.randint(0, max_score - 1)
        candidates.append((row, score))
        await push_candidate(dut, row, score)

    result = unpack(dut)
    expected = golden(candidates)
    assert result == expected, f"Expected {expected}, got {result}"