"""Tests for VirtualAllocator: device-wide VA space management.

Validates:
  T7. Basic VA allocation (contiguous, non-overlapping)
  T8. VA free + reallocation (free-list reuse)
  T9. VA space exhaustion raises AllocationError
"""
import pytest

from kernbench.policy.address.va_allocator import VirtualAllocator

_KB = 1024
_MB = 1024 * 1024
_GB = 1024 * 1024 * 1024


# ── T7. Basic VA allocation ──────────────────────────────────────────


def test_alloc_returns_aligned_va():
    """First allocation returns va_base."""
    va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
    addr = va.alloc(4096)
    assert addr == 0x1_0000_0000


def test_alloc_sequential_non_overlapping():
    """Two allocations return contiguous, non-overlapping VA ranges."""
    va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
    a1 = va.alloc(4096)
    a2 = va.alloc(8192)
    assert a1 == 0x1_0000_0000
    assert a2 == 0x1_0000_1000  # a1 + 4096
    # No overlap
    assert a2 >= a1 + 4096


def test_alloc_page_aligned():
    """Allocations are page-aligned even if requested size is not page-multiple."""
    va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
    a1 = va.alloc(100)   # < 1 page, but occupies 1 page
    a2 = va.alloc(100)
    assert a2 == 0x1_0000_1000  # aligned to next page


def test_alloc_large_contiguous():
    """Large allocation (multiple pages) is contiguous."""
    va = VirtualAllocator(va_base=0x0, va_size=1 * _GB, page_size=2 * _MB)
    addr = va.alloc(8 * _MB)  # 4 pages
    assert addr == 0x0
    # Next alloc starts after 8 MB
    addr2 = va.alloc(2 * _MB)
    assert addr2 == 8 * _MB


# ── T8. VA free + reallocation ───────────────────────────────────────


def test_free_and_realloc():
    """Freed VA range can be reused by subsequent allocation."""
    va = VirtualAllocator(va_base=0x1_0000_0000, va_size=1 * _GB, page_size=4096)
    a1 = va.alloc(4096)
    a2 = va.alloc(4096)
    va.free(a1, 4096)

    # New alloc should reuse a1's range
    a3 = va.alloc(4096)
    assert a3 == a1


def test_free_coalesce():
    """Freeing adjacent blocks allows larger reallocation."""
    va = VirtualAllocator(va_base=0x0, va_size=1 * _GB, page_size=4096)
    a1 = va.alloc(4096)
    a2 = va.alloc(4096)
    a3 = va.alloc(4096)

    # Free first two (adjacent)
    va.free(a1, 4096)
    va.free(a2, 4096)

    # Should be able to allocate 8192 contiguous from the freed region
    a4 = va.alloc(8192)
    assert a4 == a1  # reuses coalesced region


def test_free_out_of_order():
    """Free in non-sequential order still works."""
    va = VirtualAllocator(va_base=0x0, va_size=1 * _GB, page_size=4096)
    a1 = va.alloc(4096)
    a2 = va.alloc(4096)
    a3 = va.alloc(4096)

    va.free(a2, 4096)  # free middle
    a4 = va.alloc(4096)
    assert a4 == a2  # reuses middle slot


# ── T9. VA space exhaustion ──────────────────────────────────────────


def test_alloc_exhaustion():
    """Allocation beyond VA space raises AllocationError."""
    va = VirtualAllocator(va_base=0x0, va_size=8192, page_size=4096)
    va.alloc(4096)
    va.alloc(4096)
    with pytest.raises(Exception, match="[Aa]lloc|[Ee]xhaust|[Oo]ut of"):
        va.alloc(4096)


def test_alloc_after_partial_free():
    """After freeing some, can allocate again within freed space."""
    va = VirtualAllocator(va_base=0x0, va_size=8192, page_size=4096)
    a1 = va.alloc(4096)
    a2 = va.alloc(4096)

    # Space is full
    with pytest.raises(Exception):
        va.alloc(4096)

    # Free one, now can allocate again
    va.free(a1, 4096)
    a3 = va.alloc(4096)
    assert a3 == a1


# ── Stats / inspection ───────────────────────────────────────────────


def test_used_and_total():
    """used and total properties reflect allocation state."""
    va = VirtualAllocator(va_base=0x0, va_size=1 * _MB, page_size=4096)
    assert va.used == 0
    assert va.total == 1 * _MB

    va.alloc(4096)
    assert va.used == 4096

    va.alloc(8192)
    assert va.used == 4096 + 8192  # page-aligned: 4096 + 8192 = 12288