import pytest

from kernbench.policy.address.allocator import AddressConfig, AllocationError, PEMemAllocator
from kernbench.policy.address.phyaddr import PhysAddr, PhysAddrError, UnitType

_MB = 1 << 20
_GB = 1 << 30

# Topology-matching config: 48GB HBM / 8 slices / 16MB TCM / 4MB reserved / 32MB SRAM
_CFG = AddressConfig(
    sip_count=2,
    cubes_per_sip=16,
    pes_per_cube=8,
    hbm_bytes_per_cube=48 * _GB,
    hbm_slices_per_cube=8,
    tcm_bytes_per_pe=16 * _MB,
    tcm_scheduler_reserved_bytes=4 * _MB,
    sram_bytes_per_cube=32 * _MB,
)


# ── Immutability & value semantics ──────────────────────────────────


def test_physaddr_immutable():
    pa = PhysAddr.hbm_addr(rack_id=0, sip_id=0, cube_id=0, hbm_offset=0)
    with pytest.raises(AttributeError):
        pa.rack_id = 1  # type: ignore[misc]
    # hashable
    {pa}
    # comparable
    pa2 = PhysAddr.hbm_addr(rack_id=0, sip_id=0, cube_id=0, hbm_offset=0)
    assert pa == pa2


# ── HBM encode/decode roundtrip ────────────────────────────────────


def test_hbm_encode_decode_roundtrip():
    pa = PhysAddr.hbm_addr(rack_id=2, sip_id=3, cube_id=5, hbm_offset=0x1000)
    raw = pa.encode()
    dec = PhysAddr.decode(raw)
    assert dec.rack_id == 2
    assert dec.sip_id == 3
    assert dec.cube_id == 5
    assert dec.kind == "hbm"
    assert dec.hbm_offset == 0x1000


# ── PE resource encode/decode roundtrip ─────────────────────────────


def test_pe_resource_encode_decode_roundtrip():
    pa = PhysAddr(
        rack_id=1, sip_id=2, sip_seg=7, local_offset=0,
        kind="pe_resource", cube_id=7,
        unit_type=UnitType.PE, pe_id=3, ext=1, sub_offset=0xFF,
    )
    # manually build local_offset matching bit layout
    local_offset = (UnitType.PE << 34) | (3 << 30) | (1 << 29) | 0xFF
    pa2 = PhysAddr(
        rack_id=1, sip_id=2, sip_seg=7, local_offset=local_offset,
        kind="pe_resource", cube_id=7,
        unit_type=UnitType.PE, pe_id=3, ext=1, sub_offset=0xFF,
    )
    raw = pa2.encode()
    dec = PhysAddr.decode(raw)
    assert dec.kind == "pe_resource"
    assert dec.unit_type == UnitType.PE
    assert dec.pe_id == 3
    assert dec.ext == 1
    assert dec.sub_offset == 0xFF


# ── pe_hbm_addr factory ────────────────────────────────────────────


def test_pe_hbm_addr_factory():
    SLICE = 6 * (1 << 30)  # 6 GB per PE slice
    pa = PhysAddr.pe_hbm_addr(
        rack_id=0, sip_id=0, cube_id=0,
        pe_id=2, pe_local_hbm_offset=1024, slice_size_bytes=SLICE,
    )
    assert pa.kind == "hbm"
    assert pa.cube_id == 0
    assert pa.hbm_offset == 2 * SLICE + 1024


def test_pe_hbm_addr_overflow():
    SLICE = 6 * (1 << 30)
    with pytest.raises(PhysAddrError, match="pe_local_hbm_offset"):
        PhysAddr.pe_hbm_addr(
            rack_id=0, sip_id=0, cube_id=0,
            pe_id=0, pe_local_hbm_offset=SLICE, slice_size_bytes=SLICE,
        )


# ── Invalid unit_type decode (fix #1) ──────────────────────────────


def test_invalid_unit_type_raises():
    # Craft a PE-resource address with unit_type=7 (invalid)
    local_offset = (7 << 34) | (0 << 30) | 0
    pa_raw = PhysAddr(
        rack_id=0, sip_id=0, sip_seg=0, local_offset=local_offset,
    )
    raw = pa_raw.encode()
    with pytest.raises(PhysAddrError, match="unit_type"):
        PhysAddr.decode(raw)


# ── hbm_pe_id utility (fix #3) ─────────────────────────────────────


def test_hbm_pe_id_utility():
    SLICE = 6 * (1 << 30)  # 6 GB
    pa = PhysAddr.pe_hbm_addr(
        rack_id=0, sip_id=0, cube_id=0,
        pe_id=5, pe_local_hbm_offset=256, slice_size_bytes=SLICE,
    )
    assert PhysAddr.hbm_pe_id(pa.hbm_offset, SLICE) == 5


# ── UnitType.SRAM exists (fix #5) ──────────────────────────────────


def test_sram_unit_type_exists():
    assert UnitType.SRAM == 2


# ── cube_sram_addr factory + roundtrip ──────────────────────────────


def test_cube_sram_addr_roundtrip():
    pa = PhysAddr.cube_sram_addr(
        rack_id=0, sip_id=1, cube_id=3, sram_offset=0x800,
    )
    assert pa.kind == "pe_resource"
    assert pa.unit_type == UnitType.SRAM
    assert pa.cube_id == 3
    assert pa.sub_offset == 0x800
    # encode → decode roundtrip
    dec = PhysAddr.decode(pa.encode())
    assert dec.unit_type == UnitType.SRAM
    assert dec.cube_id == 3
    assert dec.sub_offset == 0x800


def test_cube_sram_addr_range_check():
    with pytest.raises(PhysAddrError):
        PhysAddr.cube_sram_addr(
            rack_id=0, sip_id=0, cube_id=0,
            sram_offset=(1 << 29),  # exceeds 29-bit sub_offset
        )


# ── pe_tcm_addr factory + roundtrip ────────────────────────────────


def test_pe_tcm_addr_roundtrip():
    pa = PhysAddr.pe_tcm_addr(
        rack_id=0, sip_id=0, cube_id=2, pe_id=7, tcm_offset=0x400,
    )
    assert pa.kind == "pe_resource"
    assert pa.unit_type == UnitType.PE
    assert pa.pe_id == 7
    assert pa.cube_id == 2
    assert pa.sub_offset == 0x400
    # encode → decode roundtrip
    dec = PhysAddr.decode(pa.encode())
    assert dec.unit_type == UnitType.PE
    assert dec.pe_id == 7
    assert dec.sub_offset == 0x400


def test_pe_tcm_addr_range_check():
    with pytest.raises(PhysAddrError):
        PhysAddr.pe_tcm_addr(
            rack_id=0, sip_id=0, cube_id=0, pe_id=0,
            tcm_offset=(1 << 29),  # exceeds 29-bit sub_offset
        )


# ── AddressConfig ───────────────────────────────────────────────────


def test_address_config_derived_sizes():
    assert _CFG.hbm_slice_bytes == 6 * _GB
    assert _CFG.tcm_allocatable_bytes == 12 * _MB


# ── PEMemAllocator: HBM ────────────────────────────────────────────


def _make_alloc(pe_id: int = 0) -> PEMemAllocator:
    return PEMemAllocator(rack_id=0, sip_id=0, cube_id=0, pe_id=pe_id, cfg=_CFG)


def test_allocator_hbm_basic():
    a = _make_alloc(pe_id=3)
    pa = a.alloc_hbm(4096)
    assert pa.kind == "hbm"
    assert pa.sip_id == 0
    assert pa.cube_id == 0
    # hbm_offset should be pe3's slice start
    assert pa.hbm_offset == 3 * 6 * _GB


def test_allocator_hbm_sequential():
    a = _make_alloc()
    pa1 = a.alloc_hbm(1024)
    pa2 = a.alloc_hbm(2048)
    assert pa1.hbm_offset == 0  # pe0 slice start + 0
    assert pa2.hbm_offset == 1024  # pe0 slice start + 1024


def test_allocator_hbm_overflow():
    a = _make_alloc()
    a.alloc_hbm(6 * _GB - 256)
    with pytest.raises(AllocationError, match="HBM"):
        a.alloc_hbm(512)


# ── PEMemAllocator: TCM ────────────────────────────────────────────


def test_allocator_tcm_basic():
    a = _make_alloc(pe_id=5)
    pa = a.alloc_tcm(256)
    assert pa.kind == "pe_resource"
    assert pa.unit_type == UnitType.PE
    assert pa.pe_id == 5
    assert pa.sub_offset == 0


def test_allocator_tcm_respects_reserved():
    a = _make_alloc()
    # allocatable = 12 MB, should succeed
    a.alloc_tcm(12 * _MB)
    assert a.tcm_used == 12 * _MB
    assert a.tcm_total == 12 * _MB


def test_allocator_tcm_overflow():
    a = _make_alloc()
    a.alloc_tcm(12 * _MB)
    with pytest.raises(AllocationError, match="TCM"):
        a.alloc_tcm(1)


# ── PEMemAllocator: stats & determinism ─────────────────────────────


def test_allocator_stats():
    a = _make_alloc()
    a.alloc_hbm(1000)
    a.alloc_tcm(500)
    assert a.hbm_used == 1000
    assert a.hbm_total == 6 * _GB
    assert a.tcm_used == 500
    assert a.tcm_total == 12 * _MB


def test_allocator_deterministic():
    a1 = _make_alloc(pe_id=2)
    a2 = _make_alloc(pe_id=2)
    assert a1.alloc_hbm(4096) == a2.alloc_hbm(4096)
    assert a1.alloc_tcm(256) == a2.alloc_tcm(256)