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.pe{X} (per-PE controller, ADR-0019 D1)."""
    g = _graph()
    resolver = AddressResolver(g)
    # offset 0x1000 falls inside PE0's slice (slice_size = 6 GB)
    pa = PhysAddr.hbm_addr(sip_id=0, die_id=3, hbm_offset=0x1000)
    assert resolver.resolve(pa) == "sip0.cube3.hbm_ctrl.pe0"


def test_resolve_hbm_addr_high_offset():
    """HBM offset that lands in PE4's slice must resolve to hbm_ctrl.pe4."""
    g = _graph()
    resolver = AddressResolver(g)
    # 0x600000000 / (6 GB) = 4
    pa = PhysAddr.hbm_addr(sip_id=0, die_id=0, hbm_offset=0x600000000)
    assert resolver.resolve(pa) == "sip0.cube0.hbm_ctrl.pe4"


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(sip_id=1, die_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(sip_id=0, die_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.mcpu_resource_addr(
        sip_id=0, die_id=2,
        mcpu_sub_unit=0, sub_offset=0,
    )
    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(sip_id=15, die_id=0, hbm_offset=0)
    with pytest.raises(RoutingError):
        resolver.resolve(pa)


# ── PathRouter: local HBM via router mesh ────────────────────────────


def test_path_local_hbm():
    """PE0 -> own slice: pe_dma -> r0c0 -> hbm_ctrl.pe0 (1 mesh hop)."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.pe0")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert path[-1] == "sip0.cube0.hbm_ctrl.pe0"
    # Path must go through at least one router node
    assert any(n.startswith("sip0.cube0.r") for n in path), \
        "HBM path must traverse router mesh"
    # No xbar or bridge nodes in the new topology
    assert not any("xbar" in n or "bridge" in n for n in path)


# ── PathRouter: remote PE HBM (different corner, same cube) ──────────


def test_path_remote_pe_hbm():
    """PE4 (bottom half) -> its own slice: routes through router mesh."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe4", "sip0.cube0.hbm_ctrl.pe4")
    assert path[0] == "sip0.cube0.pe4.pe_dma"
    assert path[-1] == "sip0.cube0.hbm_ctrl.pe4"
    assert any(n.startswith("sip0.cube0.r") for n in path)
    assert not any("xbar" in n or "bridge" in n for n in path)


# ── PathRouter: cross-PE HBM distance reflects mesh hops (ADR-0019 D4) ─


def test_cross_pe_hbm_distance_increases_with_mesh_hops():
    """Restored ADR-0019 D4 behavior: accessing another PE's HBM slice
    must take more routing distance than accessing one's own slice,
    because each per-PE hbm_ctrl is reachable only via its PE's router.

    Replaces a previous ``test_all_pe_hbm_equidistant`` that asserted the
    over-consolidated (spec-violating) behavior introduced in 5917b34.
    """
    g = _graph()
    router = PathRouter(g)
    _, dist_local = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.pe0")
    _, dist_to_pe7 = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.pe7")
    assert dist_to_pe7 > dist_local, (
        f"pe0→pe7_slice should require more mesh distance than pe0→pe0_slice; "
        f"got local={dist_local}, to_pe7={dist_to_pe7}"
    )


def test_remote_pe_distance_not_less_than_local():
    """PE4 -> pe0_slice distance >= PE0 -> pe0_slice distance.

    Both access pe0's slice (hbm_ctrl.pe0). PE0's path is shortest; PE4
    must mesh-route up to r0c0 before entering the slice.
    """
    g = _graph()
    router = PathRouter(g)
    _, dist_pe0 = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.pe0")
    _, dist_pe4 = router.find_path_with_distance(
        "sip0.cube0.pe4", "sip0.cube0.hbm_ctrl.pe0")
    assert dist_pe4 >= dist_pe0


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


# ── PathRouter: SRAM via router mesh ─────────────────────────────────


def test_path_sram_via_router_mesh():
    """PE -> SRAM must go through router mesh nodes."""
    g = _graph()
    router = PathRouter(g)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube0.sram")
    assert path[0] == "sip0.cube0.pe0.pe_dma"
    assert path[-1] == "sip0.cube0.sram"
    # Must traverse at least one router node
    assert any(n.startswith("sip0.cube0.r") for n in path), \
        "SRAM path must traverse router mesh"
    # No xbar nodes
    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 router mesh."""
    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 n.startswith("sip0.cube0.r") for n in path)


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


def test_path_distance_positive():
    """Routed paths that traverse the mesh must have positive accumulated
    distance (ADR-0002 D4). Use a cross-PE target so the path includes
    inter-router mesh edges (which have non-zero distance_mm). The
    single-hop pe0→pe0_slice path stays at 0 because PE_DMA↔router and
    router↔hbm_ctrl are zero-length placements within the same corner."""
    g = _graph()
    router = PathRouter(g)
    _, dist = router.find_path_with_distance(
        "sip0.cube0.pe0", "sip0.cube0.hbm_ctrl.pe7")
    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.pe0")
    p2 = r2.find_path("sip0.cube0.pe3", "sip0.cube0.hbm_ctrl.pe0")
    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.hbm_ctrl (adjacent cube, E direction)
    path = router.find_path("sip0.cube0.pe0", "sip0.cube1.hbm_ctrl.pe0")
    assert len(path) >= 1  # succeeds without exception