import numpy as np

from contiguous import ContiguousPool, requests_that_fit_contiguous, requests_that_fit_paged


def test_free_memory_in_two_gaps_cannot_hold_one_big_request():
    """External fragmentation: 12 slots are free, but never 12 in a row."""
    pool = ContiguousPool(32)
    for rid, size in (("R1", 12), ("R2", 8), ("R3", 8)):
        pool.allocate(rid, size)
    pool.used.update(R1=5, R2=3, R3=6)
    pool.release("R2")
    assert pool.free_runs() == [(12, 8), (28, 4)]
    assert pool.free_total() == 12 and pool.largest_free_run() == 8
    assert pool.allocate("D", 12) is None
    assert pool.internal_waste() == (12 - 5) + (8 - 6)     # reserved slots that hold no token


def test_equal_reservations_turn_the_third_request_away():
    """Reserving 12 slots each: the lesson's 32-slot pool fits two requests, the paged pool fits all three."""
    pool = ContiguousPool(32)
    assert [pool.allocate(r, 12) for r in "ABC"] == [0, 12, None]


def test_paged_waste_is_under_one_block_per_request():
    lengths = np.random.default_rng(1).integers(1, 500, size=300)
    result = requests_that_fit_paged(lengths, pool_tokens=20000, block_size=16)
    assert result["admitted"] > requests_that_fit_contiguous(20000, 2048)
    assert 0 <= result["waste"] < result["admitted"] * 16     # rounding down would make it negative


def test_lesson_capacity_example_matches_the_text():
    """16 GiB of Llama 3 8B cache is 131,072 slots; seed 0 draws 500 lengths from 100 to 2,000."""
    pool_tokens = 131072
    lengths = np.random.default_rng(0).integers(100, 2001, size=500)
    fit_contiguous = requests_that_fit_contiguous(pool_tokens, 8192)
    paged = requests_that_fit_paged(lengths, pool_tokens, block_size=16)
    assert fit_contiguous == 16
    assert round(lengths[:fit_contiguous].sum() / (fit_contiguous * 8192), 3) == 0.129
    assert (paged["admitted"], paged["allocated"], paged["waste"]) == (122, 130800, 913)