import pytest

from pathlib import Path

from kernbench.policy.address.phyaddr import PhysAddr, UnitType
from kernbench.policy.routing.router import AddressResolver, PathRouter, RoutingError
from kernbench.topology.builder import load_topology

TOPOLOGY_PATH = Path(__file__).parent.parent / "topology.yaml"


def _graph():
    return load_topology(TOPOLOGY_PATH)


# ── AddressResolver ──────────────────────────────────────────────────


def test_resolve_hbm_addr():
    """HBM address -> sip{S}.cube{C}.hbm_ctrl.slice{P}"""
    g = _graph()
    resolver = AddressResolver(g)
    # hbm_offset=0x1000, slice_size=6GB -> slice 0
    pa = PhysAddr.hbm_addr(rack_id=0, sip_id=0, cube_id=3, hbm_offset=0x1000)
    assert resolver.resolve(pa) == "sip0.cube3.hbm_ctrl.slice0"


def test_resolve_hbm_addr_slice4():
    """HBM address in PE4's slice range -> slice4."""
    g = _graph()
    resolver = AddressResolver(g)
    # slice_size = 6GB; PE4 offset starts at 4*6GB = 24GB = 0x600000000
    pa = PhysAddr.hbm_addr(rack_id=0, sip_id=0, cube_id=0, hbm_offset=0x600000000)
    assert resolver.resolve(pa) == "sip0.cube0.hbm_ctrl.slice4"


def test_resolve_pe_tcm_addr():
    """PE TCM address → sip{S}.cube{C}.pe{P}.pe_tcm"""
    g = _graph()
    resolver = AddressResolver(g)
    pa = PhysAddr.pe_tcm_addr(rack_id=0, sip_id=1, cube_id=5, pe_id=7, tcm_offset=0x400)
    assert resolver.resolve(pa) == "sip1.cube5.pe7.pe_tcm"


def test_resolve_sram_addr():
    """SRAM address → sip{S}.cube{C}.sram"""
    g = _graph()
    resolver = AddressResolver(g)
    pa = PhysAddr.cube_sram_addr(rack_id=0, sip_id=0, cube_id=10, sram_offset=0x800)
    assert resolver.resolve(pa) == "sip0.cube10.sram"


def test_resolve_mcpu_addr():
    """MCPU pe_resource address → sip{S}.cube{C}.m_cpu"""
    g = _graph()
    resolver = AddressResolver(g)
    pa = PhysAddr(
        rack_id=0, sip_id=0, sip_seg=2, local_offset=(UnitType.MCPU << 34),
        kind="pe_resource", cube_id=2, unit_type=UnitType.MCPU,
    )
    assert resolver.resolve(pa) == "sip0.cube2.m_cpu"


def test_resolve_nonexistent_node():
    """Address pointing to a node outside the compiled topology raises RoutingError."""
    g = _graph()
    resolver = AddressResolver(g)
    # sip_id=15 doesn't exist in the 2-SIP topology
    pa = PhysAddr.hbm_addr(rack_id=0, sip_id=15, cube_id=0, hbm_offset=0)
    with pytest.raises(RoutingError):
        resolver.resolve(pa)


# ── PathRouter: local HBM (same xbar half) ──────────────────────────


def test_path_local_hbm_same_half():
    """PE0 -> slice0 (local): pe_dma -> noc -> xbar_top -> hbm_ctrl.slice0."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice0")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert "sip0.cube0.noc" in path
    assert "sip0.cube0.xbar_top" in path
    assert path[-1] == "sip0.cube0.hbm_ctrl.slice0"
    assert not any("bridge" in n for n in path)
    assert len(path) == 4  # pe_dma → noc → xbar_top → slice0


# ── PathRouter: same-half remote HBM ────────────────────────────────


def test_path_same_half_remote_hbm():
    """PE0 -> slice1: same-half via noc → xbar_top, no bridge."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice1")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert "sip0.cube0.noc" in path
    assert "sip0.cube0.xbar_top" in path
    assert path[-1] == "sip0.cube0.hbm_ctrl.slice1"
    assert not any("bridge" in n for n in path)
    assert len(path) == 4  # pe_dma → noc → xbar_top → slice1


# ── PathRouter: cross-half HBM ──────────────────────────────────────


def test_path_cross_half_hbm():
    """PE0 -> slice4 (cross-half): pe_dma → noc → xbar_top → bridge → xbar_bot → slice4."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice4")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert "sip0.cube0.xbar_top" in path
    assert any("bridge" in n for n in path), "cross-half HBM must traverse bridge"
    assert "sip0.cube0.xbar_bot" in path
    assert path[-1] == "sip0.cube0.hbm_ctrl.slice4"
    assert len(path) == 6  # pe_dma → noc → xbar_top → bridge → xbar_bot → slice4


def test_path_cross_half_via_xbar_top():
    """PE4 (bottom) -> slice2 (top) goes through xbar_top via NOC.

    NOC connects directly to xbar_top (low routing weight), so
    bottom PEs access top-half HBM through noc → xbar_top.
    """
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe4", "sip0.cube0.hbm_ctrl.slice2")
    assert "sip0.cube0.xbar_top" in path
    assert path[-1] == "sip0.cube0.hbm_ctrl.slice2"


def test_cross_half_distance_greater():
    """Cross-half HBM access must have greater distance than local-half."""
    g = _graph()
    router = PathRouter(g)
    _, dist_local = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice0")
    _, dist_cross = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice4")
    assert dist_cross > dist_local


def test_path_same_half_same_distance():
    """Same-half HBM slices (PE0->slice0 vs PE0->slice3) have same distance.

    With xbar_top/bot, all top-half slices are equidistant via noc → xbar_top.
    """
    g = _graph()
    router = PathRouter(g)
    _, dist_local = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice0")
    _, dist_remote = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice3")
    assert dist_remote == dist_local, (
        f"same-half slices should have equal distance: "
        f"slice0={dist_local:.2f}mm, slice3={dist_remote:.2f}mm"
    )


def test_path_remote_cube_hbm():
    """PE0 in cube0 can reach HBM in cube1 via UCIe (ADR-0004 D4)."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube1.hbm_ctrl.slice0")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert path[-1] == "sip0.cube1.hbm_ctrl.slice0"
    # inter-cube path must cross a UCIe link
    assert any("ucie" in n for n in path), "remote cube path must traverse UCIe"
    # must not be trivially short (needs noc + ucie + remote noc + xbar)
    assert len(path) >= 5


# ── PathRouter: SRAM via NOC ────────────────────────────────────────


def test_path_sram_via_noc():
    """PE → SRAM must go through NOC (non-HBM data path)."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube0.sram")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert "sip0.cube0.noc" in path
    assert path[-1] == "sip0.cube0.sram"
    # should NOT go through xbar (SRAM is non-HBM path)
    assert not any("xbar" in n for n in path)


# ── PathRouter: PE TCM (local) ──────────────────────────────────────


def test_path_local_tcm():
    """PE0 → own TCM is PE-internal, not via xbar or noc."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube0.pe0.pe_tcm")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert path[-1] == "sip0.cube0.pe0.pe_tcm"
    # PE-internal path, no fabric
    assert not any("xbar" in n or "noc" in n for n in path)


# ── PathRouter: distance monotonic ──────────────────────────────────


def test_path_distance_positive():
    """All routed paths must have accumulated distance > 0 (ADR-0002 D4)."""
    g = _graph()
    router = PathRouter(g)
    _, dist = router.find_path_with_distance("sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.slice0")
    assert dist > 0


def test_path_deterministic():
    """Same (src, dst) must always produce the same path."""
    g = _graph()
    r1 = PathRouter(g)
    r2 = PathRouter(g)
    p1 = r1.find_path("sip0.cube0.pe3", "sip0.cube0.hbm_ctrl.slice3")
    p2 = r2.find_path("sip0.cube0.pe3", "sip0.cube0.hbm_ctrl.slice3")
    assert p1 == p2


def test_remote_cube_path_no_routing_error():
    """Routing to remote cube HBM must not raise RoutingError (ADR-0004 D4)."""
    g = _graph()
    router = PathRouter(g)
    # cube0.PE0 -> cube1.slice0 (adjacent cube, E direction)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube1.hbm_ctrl.slice0")
    assert len(path) >= 1  # succeeds without exception